""" Evaluation script for comparing baseline and gradient ascent pipelines using multiple metrics. This script evaluates both pipelines on COCO or Pick-a-Pic validation sets and computes various preference and quality metrics. """ import warnings warnings.filterwarnings("ignore") import torch import torch.nn as nn import json import os import sys import logging from glob import glob from pathlib import Path from PIL import Image from diffusers import StableDiffusionPipeline, DDIMScheduler, UNet2DConditionModel from models import LRMRewardModel from pipelines.sd15_gradient_ascent_pipeline import StableDiffusionGradientAscentPipeline from torchmetrics.image.fid import FrechetInceptionDistance from torchmetrics.multimodal import CLIPScore from transformers import CLIPModel, CLIPProcessor from tqdm import tqdm import numpy as np import argparse from datasets import load_dataset from grad_ascent_configs import get_config, list_configs import matplotlib.pyplot as plt import matplotlib matplotlib.use('Agg') # Use non-interactive backend # Import evaluation metrics sys.path.append('../evaluation') from huggingface_hub import hf_hub_download import random def configure_hf_runtime(hf_cache_dir=None, force_offline=False): """Set Hugging Face cache/offline environment for cluster-safe execution.""" cache_dir = hf_cache_dir or os.getenv("HF_HUB_CACHE") or os.getenv("HUGGINGFACE_HUB_CACHE") if cache_dir: os.environ["HF_HUB_CACHE"] = cache_dir os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir os.environ["HF_HOME"] = os.path.dirname(cache_dir) env_offline = os.getenv("HF_HUB_OFFLINE", "0").strip().lower() in {"1", "true", "yes", "on"} offline_enabled = bool(force_offline or env_offline) if offline_enabled: os.environ["HF_DATASETS_OFFLINE"] = "1" os.environ["HF_METRICS_OFFLINE"] = "1" os.environ["HF_MODULES_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" os.environ["DIFFUSERS_OFFLINE"] = "1" os.environ["HF_HUB_OFFLINE"] = "1" return cache_dir, offline_enabled def resolve_default_lrm_model(): """Prefer local LRM checkout when available; otherwise fall back to HF repo id.""" project_root = Path(__file__).resolve().parents[1] local_lrm = project_root / "lrm" / "lrm_15" / "LRM" if local_lrm.exists(): return str(local_lrm) return "casiatao/LRM" def load_pickapic_prompts(max_samples=None, cache_dir=None, offline=False): """Load Pick-a-Pic prompts with robust offline fallback to cached parquet shards.""" split = "validation_unique" if not offline: try: ds = load_dataset("pickapic-anonymous/pickapic_v1", split=split, streaming=True) prompts = [] for i, sample in enumerate(ds): prompts.append(sample["caption"]) if max_samples and i + 1 >= max_samples: break return prompts except Exception as e: print(f"Warning: online streaming load failed ({e}). Trying cached offline parquet shards.") cache_candidates = [] for p in [ cache_dir, os.getenv("HF_HUB_CACHE"), os.getenv("HUGGINGFACE_HUB_CACHE"), (os.path.join(os.getenv("HF_HOME"), "hub") if os.getenv("HF_HOME") else None), os.path.expanduser("~/.cache/huggingface/hub"), "/scratch/rr81/ma5430/.cache/huggingface/hub", ]: if p and p not in cache_candidates: cache_candidates.append(p) for cache_root in cache_candidates: repo_cache = os.path.join(cache_root, "datasets--pickapic-anonymous--pickapic_v1") if not os.path.isdir(repo_cache): continue snapshot_dir = None ref_main = os.path.join(repo_cache, "refs", "main") if os.path.isfile(ref_main): revision = open(ref_main, "r", encoding="utf-8").read().strip() candidate = os.path.join(repo_cache, "snapshots", revision) if os.path.isdir(candidate): snapshot_dir = candidate if snapshot_dir is None: snapshots = sorted(glob(os.path.join(repo_cache, "snapshots", "*"))) if snapshots: snapshot_dir = snapshots[-1] if snapshot_dir is None: continue data_dir = os.path.join(snapshot_dir, "data") if not os.path.isdir(data_dir): continue selected_split = split parquet_files = sorted(glob(os.path.join(data_dir, f"{selected_split}-*.parquet"))) if not parquet_files: for alt_split in ("test_unique", "test"): alt_files = sorted(glob(os.path.join(data_dir, f"{alt_split}-*.parquet"))) if alt_files: selected_split = alt_split parquet_files = alt_files print(f"Offline cache missing split '{split}', falling back to '{selected_split}'.") break if not parquet_files: continue print( f"Loading cached Pick-a-Pic split '{selected_split}' from {len(parquet_files)} parquet shards\n" f"cache={repo_cache}" ) ds = load_dataset("parquet", data_files=parquet_files, split="train") prompts = ds["caption"] if max_samples: prompts = prompts[:max_samples] return list(prompts) raise RuntimeError( "Could not load pickapic prompts in offline mode. " "Set --hf_cache_dir to a cache that contains datasets--pickapic-anonymous--pickapic_v1." ) def resolve_scorer_device(requested_device, generation_device, min_free_gb_for_gpu=14.0): """Choose where metric scorers should run to avoid GPU OOM/cudnn init failures.""" if requested_device == "cpu": return "cpu" if not torch.cuda.is_available() or not str(generation_device).startswith("cuda"): return "cpu" if requested_device == "cuda": return generation_device # Auto mode: only keep scorers on GPU if enough headroom remains after loading generation models. try: free_bytes, total_bytes = torch.cuda.mem_get_info(torch.device(generation_device)) free_gb = free_bytes / (1024 ** 3) total_gb = total_bytes / (1024 ** 3) print(f"GPU memory before scorer load: {free_gb:.2f} GB free / {total_gb:.2f} GB total") if free_gb >= min_free_gb_for_gpu: return generation_device print( f"⚠ Low free VRAM ({free_gb:.2f} GB). Running scorers on CPU to keep diffusion stable. " f"Use --scorer_device cuda to force GPU scorers." ) return "cpu" except Exception as e: print(f"Warning: could not inspect CUDA free memory ({e}). Falling back to CPU scorers.") return "cpu" def configure_cudnn_safely(device): """Disable cuDNN when the current GPU or runtime cannot initialize it safely.""" if not torch.cuda.is_available() or not str(device).startswith("cuda"): return try: major, minor = torch.cuda.get_device_capability(torch.device(device)) if (major, minor) < (7, 5): print( f"⚠ Detected compute capability sm_{major}{minor} (< 75). " "Disabling cuDNN to prevent runtime initialization failures." ) torch.backends.cudnn.enabled = False return # Force a cuDNN init probe early so failures are handled once at startup. _ = torch.backends.cudnn.version() except Exception as e: print(f"⚠ cuDNN init probe failed ({e}). Disabling cuDNN for this run.") torch.backends.cudnn.enabled = False def seed_everything(seed: int): """Locks down all random number generators for absolute reproducibility.""" # 1. Python & Numpy random.seed(seed) np.random.seed(seed) # 2. PyTorch Base torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) # For multi-GPU # 3. cuDNN Determinism (Crucial for consistent gradients) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # 4. Optional: Force deterministic algorithms for PyTorch 2.0+ # Uncomment if variance persists, but it may slow down generation slightly # torch.use_deterministic_algorithms(True) class MLP(nn.Module): """MLP for aesthetic scoring.""" def __init__(self): super().__init__() self.layers = nn.Sequential( nn.Linear(768, 1024), nn.Dropout(0.2), nn.Linear(1024, 128), nn.Dropout(0.2), nn.Linear(128, 64), nn.Dropout(0.1), nn.Linear(64, 16), nn.Linear(16, 1), ) @torch.no_grad() def forward(self, embed): return self.layers(embed) class AestheticScorer(torch.nn.Module): """Aesthetic scorer using CLIP and MLP.""" def __init__(self, dtype, device, clip_name_or_path="openai/clip-vit-large-patch14", aesthetic_path="./sac+logos+ava1-l14-linearMSE.pth"): super().__init__() self.clip = CLIPModel.from_pretrained(clip_name_or_path) self.processor = CLIPProcessor.from_pretrained(clip_name_or_path) self.mlp = MLP() # Load aesthetic weights if os.path.exists(aesthetic_path): state_dict = torch.load(aesthetic_path, map_location='cpu') self.mlp.load_state_dict(state_dict) else: print(f"Warning: Aesthetic weights not found at {aesthetic_path}") self.dtype = dtype self.to(device) self.eval() @torch.no_grad() def __call__(self, images): device = next(self.parameters()).device inputs = self.processor(images=images, return_tensors="pt") inputs = {k: v.to(self.dtype).to(device) for k, v in inputs.items()} embed = self.clip.get_image_features(**inputs) # normalize embedding embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True) return self.mlp(embed).squeeze(1) class TeeLogger: """Logger that writes to both console and file.""" def __init__(self, log_file): self.terminal = sys.stdout self.log = open(log_file, 'w') def write(self, message): self.terminal.write(message) self.log.write(message) self.log.flush() def flush(self): self.terminal.flush() self.log.flush() def close(self): self.log.close() def setup_logging(output_dir): """Setup logging to both console and file.""" output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) log_file = output_path / "log.log" # Redirect stdout to both console and file tee = TeeLogger(log_file) sys.stdout = tee return tee, log_file def load_validation_data(data_dir, max_samples=None, dataset_type="coco", hf_cache_dir=None, offline=False): """Load validation prompts and image paths. Args: data_dir: Path to data directory max_samples: Maximum number of samples to load dataset_type: Type of dataset ("coco" or "pickapic") Returns: prompts: List of text prompts image_paths: List of image paths (None for pickapic streaming dataset) """ if dataset_type == "coco": data_dir = Path(data_dir) val_json = data_dir / "coco" / "caption_val.json" if not val_json.exists(): raise FileNotFoundError(f"Validation JSON not found: {val_json}") with open(val_json, 'r') as f: data = json.load(f) # Validate that image folder exists val_img_dir = data_dir / "coco" / "images" / "val" if not val_img_dir.exists(): raise FileNotFoundError(f"Validation image directory not found: {val_img_dir}") # Parse data prompts = [] image_paths = [] for img_path, caption in data.items(): full_path = data_dir / "coco" / img_path if full_path.exists(): prompts.append(caption) image_paths.append(str(full_path)) else: print(f"Warning: Image not found: {full_path}") if max_samples: prompts = prompts[:max_samples] image_paths = image_paths[:max_samples] print(f"Loaded {len(prompts)} COCO validation samples") return prompts, image_paths elif dataset_type == "pickapic": print("Loading Pick-a-Pic validation prompts...") prompts = load_pickapic_prompts(max_samples=max_samples, cache_dir=hf_cache_dir, offline=offline) print(f"Loaded {len(prompts)} Pick-a-Pic validation samples") return prompts, None # No reference images for Pick-a-Pic else: raise ValueError(f"Unknown dataset type: {dataset_type}. Choose 'coco' or 'pickapic'.") def generate_and_evaluate( pipeline, prompts, image_paths, device, dtype, num_inference_steps=20, guidance_scale=7.5, seed=42, batch_size=1, apply_gradient_ascent=False, mode_name="baseline", log_interval=10, output_dir=None, save_images=False, clip_scorer=None, aesthetic_scorer=None, pick_scorer=None, hpsv2_scorer=None, hpsv21_scorer=None, imagereward_scorer=None, compute_fid=True, capture_trajectory=False ): """Generate images and update FID metric.""" pipeline.to(device) print(f"\nGenerating images with {mode_name} mode...") all_rewards = [] all_clip_scores = [] all_aesthetic_scores = [] all_pick_scores = [] all_hpsv2_scores = [] all_hpsv21_scores = [] all_imagereward_scores = [] lr_history_first_image = None # Store LR history for first image trajectory_first_image = [] num_batches = (len(prompts) + batch_size - 1) // batch_size # Create output directory if saving images if save_images and output_dir: mode_output_dir = Path(output_dir) / mode_name mode_output_dir.mkdir(parents=True, exist_ok=True) # Disable internal progress bars pipeline.set_progress_bar_config(disable=True) for idx, i in enumerate(tqdm(range(0, len(prompts), batch_size), desc=f"Generating {mode_name}")): batch_prompts = prompts[i:i+batch_size] batch_real_paths = image_paths[i:i+batch_size] if image_paths is not None else None batch_num = idx + 1 # Initialize FID metric if needed fid_metric = None real_images_tensor = None if compute_fid and batch_real_paths is not None: fid_metric = FrechetInceptionDistance().to(device) # Load and update FID with real images for this batch real_images = [] for path in batch_real_paths: img = Image.open(path).convert("RGB") img = img.resize((512, 512)) # Inception v3 input size img_array = np.array(img) real_images.append(img_array) # Convert to tensor [B, H, W, C] -> [B, C, H, W] real_images_tensor = torch.from_numpy(np.stack(real_images)).permute(0, 3, 1, 2).float() real_images_tensor = real_images_tensor.to(device) # Generate images generator = torch.Generator(device=device).manual_seed(seed + i) # Only capture trajectory for the very first batch to save RAM def trajectory_callback(step, timestep, latents): if idx == 0 and capture_trajectory: # Detach and move to CPU immediately to prevent VRAM OOM trajectory_first_image.append(latents.detach().cpu().clone()) with torch.no_grad(): result = pipeline( prompt=batch_prompts, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator, track_rewards=True, print_rewards=False, apply_gradient_ascent=apply_gradient_ascent, verbose_grad=False, callback=trajectory_callback if capture_trajectory else None, callback_steps=1 ) # Process generated images images = result.images # Update FID metric if computing it if compute_fid and fid_metric is not None: image_tensors = [] for img in images: img_resized = img.resize((512, 512)) # Inception v3 input size img_array = np.array(img_resized) image_tensors.append(img_array) # Convert to tensor and update FID images_tensor = torch.from_numpy(np.stack(image_tensors)).permute(0, 3, 1, 2).float() images_tensor = images_tensor.to(device) if batch_size == 1: real_images_tensor = torch.cat([real_images_tensor, real_images_tensor], dim=0).to(dtype=torch.uint8) images_tensor = torch.cat([images_tensor, images_tensor], dim=0).to(dtype=torch.uint8) fid_metric.update(real_images_tensor, real=True) fid_metric.update(images_tensor, real=False) # Track rewards - get the final timestep reward (t=0) current_batch_final_reward = None current_batch_final_timestep = None if hasattr(pipeline, 'reward_history') and pipeline.reward_history: # For each image, get the reward from the last denoising step (t=0 or closest to 0) num_steps_per_image = num_inference_steps # Get the last entry which corresponds to the final timestep of the last image in batch final_entry = pipeline.reward_history[-1] current_batch_final_reward = final_entry['reward_score'] current_batch_final_timestep = final_entry['timestep'] all_rewards.append(current_batch_final_reward) # Capture LR history from first image if gradient ascent is enabled if apply_gradient_ascent and idx == 0 and lr_history_first_image is None: if hasattr(pipeline, 'grad_guidance') and pipeline.grad_guidance: grad_stats = pipeline.grad_guidance.get_statistics() if grad_stats and 'detailed_stats' in grad_stats: # Extract LR history from the gradient ascent statistics lr_history_first_image = { 'prompt': batch_prompts[0], 'timesteps': [], 'learning_rates': [], # All LR values from all gradient steps 'rewards': [] } for stat in grad_stats['detailed_stats']: lr_history_first_image['timesteps'].append(stat['timestep']) if 'lr_history' in stat: # Extend with all LR values from this timestep's gradient steps lr_history_first_image['learning_rates'].extend(stat['lr_history']) # Collect all rewards from reward_history for each gradient step if 'reward_history' in stat: lr_history_first_image['rewards'].extend(stat['reward_history']) # Compute CLIP score if clip_scorer is not None: clip_device = next(clip_scorer.parameters()).device # Convert PIL images to tensor format for CLIP score [C, H, W] in range [0, 1] for img, prompt in zip(images, batch_prompts): img_array = np.array(img).astype(np.float32) img_tensor = torch.from_numpy(img_array).permute(2, 0, 1).unsqueeze(0).to(clip_device) clip_score = clip_scorer(img_tensor, [prompt]).item() all_clip_scores.append(clip_score) # Compute aesthetic score if aesthetic_scorer is not None: aesthetic_scores = aesthetic_scorer(images) if isinstance(aesthetic_scores, torch.Tensor): aesthetic_scores = aesthetic_scores.cpu().numpy() if aesthetic_scores.ndim == 0: aesthetic_scores = [aesthetic_scores.item()] all_aesthetic_scores.extend(aesthetic_scores.tolist() if hasattr(aesthetic_scores, 'tolist') else [aesthetic_scores]) # Compute PickScore if pick_scorer is not None: for img, prompt in zip(images, batch_prompts): pick_score = pick_scorer(prompt, [img])[0] all_pick_scores.append(pick_score) # Compute HPSv2 score if hpsv2_scorer is not None: for img, prompt in zip(images, batch_prompts): hpsv2_score = hpsv2_scorer.score(img, prompt)[0] all_hpsv2_scores.append(hpsv2_score) # Compute HPSv2.1 score if hpsv21_scorer is not None: for img, prompt in zip(images, batch_prompts): hpsv21_score = hpsv21_scorer.score(img, prompt)[0] all_hpsv21_scores.append(hpsv21_score) # Compute ImageReward score if imagereward_scorer is not None: for img, prompt in zip(images, batch_prompts): imagereward_score = imagereward_scorer.score(prompt, img) all_imagereward_scores.append(imagereward_score) # Save generated images if requested if save_images and output_dir: for img_idx, img in enumerate(images): global_idx = i + img_idx img_path = mode_output_dir / f"sample_{global_idx:05d}.png" img.save(img_path) # Log intermediate FID and metrics every log_interval batches if batch_num % log_interval == 0 or batch_num == num_batches: num_samples_processed = min(i + batch_size, len(prompts)) log_msg = f"\n[{mode_name}] Batch {batch_num}/{num_batches} | Samples: {num_samples_processed}/{len(prompts)}" # Add FID if computing if compute_fid and fid_metric is not None: try: current_fid = fid_metric.compute().item() log_msg += f" | FID: {current_fid:.4f}" except Exception as e: log_msg += f" | FID: Computing..." # Add reward - show both final timestep reward and average if all_rewards: avg_reward = np.mean(all_rewards) if current_batch_final_reward is not None: log_msg += f" | Reward (t={current_batch_final_timestep}): {current_batch_final_reward:.4f}" log_msg += f" | Reward (Avg): {avg_reward:.4f}" else: log_msg += f" | Reward (Avg): {avg_reward:.4f}" # Add CLIP if computing if clip_scorer is not None and all_clip_scores: log_msg += f" | CLIP: {np.mean(all_clip_scores):.4f}" # Add aesthetic if computing if aesthetic_scorer is not None and all_aesthetic_scores: log_msg += f" | Aesthetic: {np.mean(all_aesthetic_scores):.4f}" # Add PickScore if pick_scorer is not None and all_pick_scores: log_msg += f" | PickScore: {np.mean(all_pick_scores):.4f}" # Add HPSv2 if hpsv2_scorer is not None and all_hpsv2_scores: log_msg += f" | HPSv2: {np.mean(all_hpsv2_scores):.4f}" # Add HPSv2.1 if hpsv21_scorer is not None and all_hpsv21_scores: log_msg += f" | HPSv2.1: {np.mean(all_hpsv21_scores):.4f}" # Add ImageReward if imagereward_scorer is not None and all_imagereward_scores: log_msg += f" | ImageReward: {np.mean(all_imagereward_scores):.4f}" print(log_msg) # Re-enable progress bars pipeline.set_progress_bar_config(disable=False) avg_reward = np.mean(all_rewards) if all_rewards else 0.0 avg_clip_score = np.mean(all_clip_scores) if all_clip_scores else 0.0 avg_aesthetic_score = np.mean(all_aesthetic_scores) if all_aesthetic_scores else 0.0 avg_pick_score = np.mean(all_pick_scores) if all_pick_scores else 0.0 avg_hpsv2_score = np.mean(all_hpsv2_scores) if all_hpsv2_scores else 0.0 avg_hpsv21_score = np.mean(all_hpsv21_scores) if all_hpsv21_scores else 0.0 avg_imagereward_score = np.mean(all_imagereward_scores) if all_imagereward_scores else 0.0 return avg_reward, fid_metric, avg_clip_score, avg_aesthetic_score, avg_pick_score, avg_hpsv2_score, avg_hpsv21_score, avg_imagereward_score, lr_history_first_image, trajectory_first_image def auto_increment_path(base_path): """ Create an auto-incrementing run folder inside base_path. Returns: base_path/run_1, base_path/run_2, etc. """ base_path = Path(base_path) base_path.mkdir(parents=True, exist_ok=True) # Ensure base directory exists i = 1 while True: new_path = base_path / f"run_{i}" if not new_path.exists(): return new_path i += 1 def main(): parser = argparse.ArgumentParser(description="Evaluate baseline and gradient ascent pipelines") parser.add_argument("--data_dir", type=str, default="./data", help="Path to data directory") parser.add_argument("--dataset_type", type=str, default="coco", choices=["coco", "pickapic"], help="Dataset to use for evaluation: coco or pickapic (default: coco)") parser.add_argument("--base_model", type=str, default="stable-diffusion-v1-5/stable-diffusion-v1-5", help="Base model path") parser.add_argument("--model_variant", type=str, default="origin", choices=["origin", "spo", "diffusion_dpo", "lpo"], help="SD1.5 model variant to use (default: origin)") parser.add_argument("--lrm_model", type=str, default=None, help="LRM model path. Defaults to local lrm/lrm_15/LRM when present.") parser.add_argument("--hf_cache_dir", type=str, default="/scratch/rr81/ma5430/.cache/huggingface/hub", help="Shared HF cache directory") parser.add_argument("--offline", action="store_true", help="Force fully offline mode (recommended on GPU nodes)") parser.add_argument("--num_steps", type=int, default=50, help="Number of inference steps") parser.add_argument("--cfg_scale", type=float, default=7.5, help="Classifier-free guidance scale") parser.add_argument("--seed", type=int, default=42, help="Random seed") parser.add_argument("--max_samples", type=int, default=None, help="Max samples to evaluate (None for all)") parser.add_argument("--batch_size", type=int, default=1, help="Batch size for generation (use 1 for reward model compatibility)") parser.add_argument("--fid_batch_size", type=int, default=32, help="Batch size for FID computation") parser.add_argument("--log_interval", type=int, default=10, help="Log FID and metrics every N batches") parser.add_argument("--output_dir", type=str, default="eval_outputs", help="Directory to save generated images and results") parser.add_argument("--save_images", action="store_true", help="Save all generated images to output directory") parser.add_argument("--mode", type=str, default="both", choices=["baseline", "gradient_ascent", "both"], help="Which evaluation to run: baseline, gradient_ascent, or both (default: both)") # Metrics selection parser.add_argument("--metrics", type=str, nargs="+", default=["clip", "aesthetic"], choices=["fid", "clip", "aesthetic", "pickscore", "hpsv2", "hpsv21", "imagereward"], help="Which metrics to evaluate (default: clip aesthetic)") parser.add_argument("--scorer_device", type=str, default="auto", choices=["auto", "cpu", "cuda"], help="Device for metric scorers. auto keeps scorers on GPU only when enough VRAM is free.") # Gradient ascent config parser.add_argument("--grad_config", type=str, default=None, help=f"Gradient ascent config preset (available: {', '.join(list_configs())}). " "If provided, overrides individual grad_* arguments.") parser.add_argument("--grad_range_start", type=int, default=0, help="Gradient timestep range start") parser.add_argument("--grad_range_end", type=int, default=700, help="Gradient timestep range end") parser.add_argument("--grad_steps", type=int, default=5, help="Number of gradient steps per timestep (use 5 for better reward improvement)") parser.add_argument("--grad_step_size", type=float, default=0.1, help="Gradient step size (initial LR)") # Config overrides (these override values from grad_config if specified) parser.add_argument("--override_momentum", type=float, default=None, help="Override momentum value from grad_config") parser.add_argument("--override_num_grad_steps", type=int, default=None, help="Override num_grad_steps from grad_config") parser.add_argument("--override_grad_step_size", type=float, default=None, help="Override grad_step_size from grad_config") # Cuda parser.add_argument("--cuda", type=int, default=0, help="Use CUDA device id") args = parser.parse_args() hf_cache_dir, offline_enabled = configure_hf_runtime(args.hf_cache_dir, force_offline=args.offline) if args.lrm_model is None: args.lrm_model = resolve_default_lrm_model() seed_everything(args.seed) # Configuration device = f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu" dtype = torch.float16 #if torch.cuda.is_available() else torch.float32 configure_cudnn_safely(device) # Create auto-incremented output directory args.output_dir = auto_increment_path(args.output_dir) # Setup logging to file tee_logger, log_file = setup_logging(args.output_dir) print("="*70) print("FID EVALUATION: BASELINE vs GRADIENT ASCENT") print("="*70) print(f"\nLogging to: {log_file}") print(f"\nDevice: {device}") print(f"Dataset: {args.dataset_type.upper()}") print(f"Data directory: {args.data_dir}") print(f"Base model: {args.base_model}") print(f"Model variant: {args.model_variant}") print(f"LRM model: {args.lrm_model}") print(f"HF cache dir: {hf_cache_dir or 'default'}") print(f"HF offline mode: {offline_enabled}") print(f"Inference steps: {args.num_steps}") print(f"CFG scale: {args.cfg_scale}") print(f"Batch size: {args.batch_size}") print(f"Max samples: {args.max_samples or 'All'}") print(f"Output directory: {args.output_dir}") print(f"Save images: {args.save_images}") print(f"Evaluation mode: {args.mode}") print(f"Metrics to evaluate: {', '.join(args.metrics).upper()}") if args.grad_config: print(f"Gradient ascent config: {args.grad_config}") # Load validation data print("\n" + "="*70) print("1. LOADING VALIDATION DATA") print("="*70) prompts, image_paths = load_validation_data( args.data_dir, args.max_samples, args.dataset_type, hf_cache_dir=hf_cache_dir, offline=offline_enabled, ) # Automatically disable FID if no reference images available (e.g., Pick-a-Pic dataset) can_compute_fid = image_paths is not None if not can_compute_fid and "fid" in args.metrics: print("\n⚠ Warning: FID metric requested but no reference images available. FID will be skipped.") args.metrics = [m for m in args.metrics if m != "fid"] # Load reward model print("\n" + "="*70) print("2. LOADING REWARD MODEL") print("="*70) reward_model = LRMRewardModel( pretrained_model_name_or_path=args.base_model, lrm_model_path=args.lrm_model, guidance_scale=args.cfg_scale, device=device ) if dtype == torch.float16: reward_model = reward_model.half() reward_model.eval() print("✓ Reward model loaded") # Load pipeline print("\n" + "="*70) print("3. LOADING PIPELINE") print("="*70) pretrained_kwargs = {"local_files_only": offline_enabled} if hf_cache_dir: pretrained_kwargs["cache_dir"] = hf_cache_dir # Load model based on variant if args.model_variant == "origin": base_pipeline = StableDiffusionPipeline.from_pretrained( args.base_model, torch_dtype=dtype, safety_checker=None, **pretrained_kwargs, ) print(f"✓ Loaded origin SD1.5 model") elif args.model_variant == "spo": base_pipeline = StableDiffusionPipeline.from_pretrained( 'SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep', torch_dtype=dtype, safety_checker=None, **pretrained_kwargs, ) args.cfg_scale = 5.0 # SPO uses CFG 5.0 print(f"✓ Loaded SPO SD1.5 model (cfg_scale adjusted to 5.0)") elif args.model_variant == "diffusion_dpo": unet = UNet2DConditionModel.from_pretrained( 'mhdang/dpo-sd1.5-text2image-v1', subfolder="unet", torch_dtype=dtype, **pretrained_kwargs, ) base_pipeline = StableDiffusionPipeline.from_pretrained( args.base_model, torch_dtype=dtype, safety_checker=None, unet=unet, **pretrained_kwargs, ) print(f"✓ Loaded Diffusion-DPO SD1.5 model") elif args.model_variant == "lpo": unet = UNet2DConditionModel.from_pretrained( 'casiatao/LPO', subfolder="lpo_sd15_merge/unet", torch_dtype=dtype, **pretrained_kwargs, ) base_pipeline = StableDiffusionPipeline.from_pretrained( args.base_model, torch_dtype=dtype, safety_checker=None, unet=unet, **pretrained_kwargs, ) args.cfg_scale = 5.0 # LPO uses CFG 5.0 print(f"✓ Loaded LPO SD1.5 model (cfg_scale adjusted to 5.0)") pipeline = StableDiffusionGradientAscentPipeline(**base_pipeline.components) pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config) pipeline = pipeline.to(device) pipeline.set_reward_model(reward_model) print("✓ Pipeline loaded") scorer_device = resolve_scorer_device(args.scorer_device, device) scorer_dtype = dtype if str(scorer_device).startswith("cuda") else torch.float32 print(f"Scorer device: {scorer_device}") if torch.cuda.is_available(): torch.cuda.empty_cache() # Load CLIP scorer print("\n" + "="*70) print("3.5. LOADING CLIP AND AESTHETIC SCORERS") print("="*70) # Only load scorers for requested metrics clip_scorer = None aesthetic_scorer = None pick_scorer = None hpsv2_scorer = None hpsv21_scorer = None imagereward_scorer = None if "clip" in args.metrics: try: clip_scorer = CLIPScore(model_name_or_path="openai/clip-vit-large-patch14").to(scorer_device) print("✓ CLIP scorer loaded") except Exception as e: print(f"Warning: Could not load CLIP scorer: {e}") clip_scorer = None else: print("⊘ CLIP scorer skipped (not in selected metrics)") if "aesthetic" in args.metrics: try: aesthetic_scorer = AestheticScorer(dtype=scorer_dtype, device=scorer_device) print("✓ Aesthetic scorer loaded") except Exception as e: print(f"Warning: Could not load Aesthetic scorer: {e}") aesthetic_scorer = None else: print("⊘ Aesthetic scorer skipped (not in selected metrics)") if "pickscore" in args.metrics: try: from pick_score import PickScorer pick_scorer = PickScorer( processor_name_or_path="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", model_pretrained_name_or_path="yuvalkirstain/PickScore_v1", device=scorer_device ) print("✓ PickScore scorer loaded") except Exception as e: print(f"Warning: Could not load PickScore scorer: {e}") pick_scorer = None else: print("⊘ PickScore scorer skipped (not in selected metrics)") if "hpsv2" in args.metrics: try: from hpsv2_score import HPSv2Scorer hf_dl_kwargs = {"local_files_only": offline_enabled} if hf_cache_dir: hf_dl_kwargs["cache_dir"] = hf_cache_dir hpsv2_scorer = HPSv2Scorer( clip_pretrained_name_or_path=hf_hub_download( repo_id="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", filename="open_clip_pytorch_model.bin", **hf_dl_kwargs, ), model_pretrained_name_or_path=hf_hub_download( repo_id="xswu/HPSv2", filename="HPS_v2_compressed.pt", **hf_dl_kwargs, ), device=scorer_device ) print("✓ HPSv2 scorer loaded") except Exception as e: print(f"Warning: Could not load HPSv2 scorer: {e}") hpsv2_scorer = None else: print("⊘ HPSv2 scorer skipped (not in selected metrics)") if "hpsv21" in args.metrics: try: from hpsv2_score import HPSv2Scorer hf_dl_kwargs = {"local_files_only": offline_enabled} if hf_cache_dir: hf_dl_kwargs["cache_dir"] = hf_cache_dir hpsv21_scorer = HPSv2Scorer( clip_pretrained_name_or_path=hf_hub_download( repo_id="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", filename="open_clip_pytorch_model.bin", **hf_dl_kwargs, ), model_pretrained_name_or_path=hf_hub_download( repo_id="xswu/HPSv2", filename="HPS_v2.1_compressed.pt", **hf_dl_kwargs, ), device=scorer_device ) print("✓ HPSv2.1 scorer loaded") except Exception as e: print(f"Warning: Could not load HPSv2.1 scorer: {e}") hpsv21_scorer = None else: print("⊘ HPSv2.1 scorer skipped (not in selected metrics)") if "imagereward" in args.metrics: try: from imagereward_score import load_imagereward hf_dl_kwargs = {"local_files_only": offline_enabled} if hf_cache_dir: hf_dl_kwargs["cache_dir"] = hf_cache_dir imagereward_scorer = load_imagereward( model_path=hf_hub_download(repo_id="THUDM/ImageReward", filename="ImageReward.pt", **hf_dl_kwargs), med_config=hf_hub_download(repo_id="THUDM/ImageReward", filename="med_config.json", **hf_dl_kwargs), device=scorer_device ) print("✓ ImageReward scorer loaded") except Exception as e: print(f"Warning: Could not load ImageReward scorer: {e}") imagereward_scorer = None else: print("⊘ ImageReward scorer skipped (not in selected metrics)") # Configure gradient ascent print("\n" + "="*70) print("4. CONFIGURING GRADIENT ASCENT") print("="*70) # Use config preset if provided, otherwise use individual args if args.grad_config: print(f"Loading gradient ascent config: {args.grad_config}") grad_config = get_config(args.grad_config) print(f"Config loaded: {grad_config}") # Apply overrides if specified if args.override_momentum is not None: grad_config['momentum'] = args.override_momentum print(f" Overriding momentum: {args.override_momentum}") if args.override_num_grad_steps is not None: grad_config['num_grad_steps'] = args.override_num_grad_steps print(f" Overriding num_grad_steps: {args.override_num_grad_steps}") if args.override_grad_step_size is not None: grad_config['grad_step_size'] = args.override_grad_step_size print(f" Overriding grad_step_size: {args.override_grad_step_size}") else: grad_config = { "grad_timestep_range": (args.grad_range_start, args.grad_range_end), "num_grad_steps": args.grad_steps, "grad_step_size": args.grad_step_size, } print(f"Using manual gradient ascent configuration") print(f"Gradient timestep range: {grad_config.get('grad_timestep_range', (args.grad_range_start, args.grad_range_end))}") print(f"Gradient steps: {grad_config.get('num_grad_steps', args.grad_steps)}") print(f"Gradient step size (initial LR): {grad_config.get('grad_step_size', args.grad_step_size)}") if grad_config.get('lr_scheduler_type'): print(f"LR Scheduler: {grad_config['lr_scheduler_type']}") if grad_config.get('use_momentum'): print(f"Momentum: {grad_config.get('momentum', 0.9)} (Nesterov: {grad_config.get('use_nesterov', False)})") pipeline.enable_gradient_ascent(**grad_config) # Initialize result variables fid_score_baseline = None avg_reward_baseline = None clip_score_baseline = None aesthetic_score_baseline = None pick_score_baseline = None hpsv2_score_baseline = None hpsv21_score_baseline = None imagereward_score_baseline = None fid_score_grad = None avg_reward_grad = None clip_score_grad = None aesthetic_score_grad = None pick_score_grad = None hpsv2_score_grad = None hpsv21_score_grad = None imagereward_score_grad = None grad_stats = None # ========== BASELINE EVALUATION ========== if args.mode in ["baseline", "both"]: print("\n" + "="*70) print("5. EVALUATING BASELINE") print("="*70) # Generate and evaluate baseline avg_reward_baseline, fid_baseline, clip_score_baseline, aesthetic_score_baseline, pick_score_baseline, hpsv2_score_baseline, hpsv21_score_baseline, imagereward_score_baseline, _, baseline_trajectory = generate_and_evaluate( pipeline=pipeline, prompts=prompts, image_paths=image_paths, device=device, dtype=dtype, num_inference_steps=args.num_steps, guidance_scale=args.cfg_scale, seed=args.seed, batch_size=args.batch_size, apply_gradient_ascent=False, mode_name="baseline", log_interval=args.log_interval, output_dir=args.output_dir, save_images=args.save_images, clip_scorer=clip_scorer, aesthetic_scorer=aesthetic_scorer, pick_scorer=pick_scorer, hpsv2_scorer=hpsv2_scorer, hpsv21_scorer=hpsv21_scorer, imagereward_scorer=imagereward_scorer, compute_fid=("fid" in args.metrics and can_compute_fid), capture_trajectory=True ) # Compute FID for baseline if requested if "fid" in args.metrics and fid_baseline is not None: fid_score_baseline = fid_baseline.compute().item() print(f"\n✓ Baseline FID: {fid_score_baseline:.4f}") print(f"✓ Baseline Avg Reward: {avg_reward_baseline:.4f}") if "clip" in args.metrics: print(f"✓ Baseline Avg CLIP Score: {clip_score_baseline:.4f}") if "aesthetic" in args.metrics: print(f"✓ Baseline Avg Aesthetic Score: {aesthetic_score_baseline:.4f}") if "pickscore" in args.metrics and pick_score_baseline is not None: print(f"✓ Baseline Avg PickScore: {pick_score_baseline:.4f}") if "hpsv2" in args.metrics and hpsv2_score_baseline is not None: print(f"✓ Baseline Avg HPSv2 Score: {hpsv2_score_baseline:.4f}") if "hpsv21" in args.metrics and hpsv21_score_baseline is not None: print(f"✓ Baseline Avg HPSv2.1 Score: {hpsv21_score_baseline:.4f}") if "imagereward" in args.metrics and imagereward_score_baseline is not None: print(f"✓ Baseline Avg ImageReward: {imagereward_score_baseline:.4f}") # ========== GRADIENT ASCENT EVALUATION ========== if args.mode in ["gradient_ascent", "both"]: print("\n" + "="*70) print("6. EVALUATING GRADIENT ASCENT") print("="*70) # Generate and evaluate with gradient ascent avg_reward_grad, fid_grad, clip_score_grad, aesthetic_score_grad, pick_score_grad, hpsv2_score_grad, hpsv21_score_grad, imagereward_score_grad, lr_history, guided_trajectory = generate_and_evaluate( pipeline=pipeline, prompts=prompts, image_paths=image_paths, device=device, dtype=dtype, num_inference_steps=args.num_steps, guidance_scale=args.cfg_scale, seed=args.seed, batch_size=args.batch_size, apply_gradient_ascent=True, mode_name="gradient_ascent", log_interval=args.log_interval, output_dir=args.output_dir, save_images=args.save_images, clip_scorer=clip_scorer, aesthetic_scorer=aesthetic_scorer, pick_scorer=pick_scorer, hpsv2_scorer=hpsv2_scorer, hpsv21_scorer=hpsv21_scorer, imagereward_scorer=imagereward_scorer, compute_fid=("fid" in args.metrics and can_compute_fid), capture_trajectory=True ) # Compute FID for gradient ascent if requested if "fid" in args.metrics and fid_grad is not None: fid_score_grad = fid_grad.compute().item() print(f"\n✓ Gradient Ascent FID: {fid_score_grad:.4f}") print(f"✓ Gradient Ascent Avg Reward: {avg_reward_grad:.4f}") if "clip" in args.metrics: print(f"✓ Gradient Ascent Avg CLIP Score: {clip_score_grad:.4f}") if "aesthetic" in args.metrics: print(f"✓ Gradient Ascent Avg Aesthetic Score: {aesthetic_score_grad:.4f}") if "pickscore" in args.metrics and pick_score_grad is not None: print(f"✓ Gradient Ascent Avg PickScore: {pick_score_grad:.4f}") if "hpsv2" in args.metrics and hpsv2_score_grad is not None: print(f"✓ Gradient Ascent Avg HPSv2 Score: {hpsv2_score_grad:.4f}") if "hpsv21" in args.metrics and hpsv21_score_grad is not None: print(f"✓ Gradient Ascent Avg HPSv2.1 Score: {hpsv21_score_grad:.4f}") if "imagereward" in args.metrics and imagereward_score_grad is not None: print(f"✓ Gradient Ascent Avg ImageReward: {imagereward_score_grad:.4f}") # Get gradient stats grad_stats = pipeline.grad_guidance.get_statistics() if grad_stats: print(f"\nGradient Ascent Statistics:") print(f" Applications: {grad_stats['num_applications']}") print(f" Total reward improvement: {grad_stats['total_reward_improvement']:+.4f}") print(f" Avg reward improvement: {grad_stats['avg_reward_improvement']:+.4f}") # Plot LR curve if we captured it if lr_history is not None and lr_history['learning_rates']: plot_path = Path(args.output_dir) / "lr_curve.png" # LR values are now continuous across all gradient steps lrs = lr_history['learning_rates'] steps = list(range(len(lrs))) # Step indices (0 to total_steps-1) plt.figure(figsize=(12, 6)) plt.plot(steps, lrs, linewidth=2, color='blue', alpha=0.8) # Mark the first step with a star plt.plot(steps[0], lrs[0], marker='*', markersize=20, color='gold', markeredgecolor='darkgoldenrod', markeredgewidth=2, zorder=5) # Mark timestep boundaries num_timesteps = len(lr_history['timesteps']) num_grad_steps_per_timestep = len(lrs) // num_timesteps if num_timesteps > 0 else 0 if num_grad_steps_per_timestep > 0: for i in range(num_timesteps + 1): step_idx = i * num_grad_steps_per_timestep if step_idx <= len(lrs): plt.axvline(x=step_idx, color='red', linestyle='--', alpha=0.3, linewidth=1) if i < num_timesteps: plt.text(step_idx, plt.ylim()[1] * 0.95, f't={lr_history["timesteps"][i]}', fontsize=8, color='red', alpha=0.7, ha='left') plt.xlabel('Global Gradient Step', fontsize=12) plt.ylabel('Learning Rate', fontsize=12) plt.title(f'Learning Rate Evolution Across All Gradient Steps\\nPrompt: "{lr_history["prompt"][:60]}..."', fontsize=12, fontweight='bold') plt.grid(True, alpha=0.3) # Add info text num_timesteps = len(lr_history['timesteps']) num_grad_steps_per_timestep = len(lrs) // num_timesteps if num_timesteps > 0 else 0 plt.text(0.02, 0.98, f'Total timesteps: {num_timesteps}\\nGrad steps/timestep: {num_grad_steps_per_timestep}\\nTotal grad steps: {len(lrs)}', transform=plt.gca().transAxes, fontsize=10, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5)) plt.tight_layout() plt.savefig(plot_path, dpi=150, bbox_inches='tight') plt.close() print(f"\n✓ Saved LR curve plot to: {plot_path}") print(f" Total gradient steps: {len(lrs)}") print(f" LR range: {min(lrs):.6f} → {max(lrs):.6f}") # Plot Rewards curve if we captured it if lr_history is not None and lr_history['rewards']: plot_path = Path(args.output_dir) / "rewards_curve.png" # Reward values are now continuous across all gradient steps rewards = lr_history['rewards'] steps = list(range(len(rewards))) # Step indices (0 to total_steps-1) plt.figure(figsize=(12, 6)) plt.plot(steps, rewards, linewidth=2, color='green', alpha=0.8) # Mark the first step with a star plt.plot(steps[0], rewards[0], marker='*', markersize=20, color='gold', markeredgecolor='darkgoldenrod', markeredgewidth=2, zorder=5) # Mark timestep boundaries num_timesteps = len(lr_history['timesteps']) # rewards has one extra value at the start (initial) compared to gradient steps num_grad_steps_per_timestep = (len(rewards) - num_timesteps) // num_timesteps if num_timesteps > 0 else 0 if num_grad_steps_per_timestep > 0: for i in range(num_timesteps + 1): step_idx = i * (num_grad_steps_per_timestep + 1) # +1 because reward_history includes initial if step_idx <= len(rewards): plt.axvline(x=step_idx, color='red', linestyle='--', alpha=0.3, linewidth=1) if i < num_timesteps: plt.text(step_idx, plt.ylim()[1] * 0.95, f't={lr_history["timesteps"][i]}', fontsize=8, color='red', alpha=0.7, ha='left') plt.xlabel('Global Gradient Step', fontsize=12) plt.ylabel('Reward Score', fontsize=12) plt.title(f'Reward Evolution Across All Gradient Steps\nPrompt: "{lr_history["prompt"][:60]}..."', fontsize=12, fontweight='bold') plt.grid(True, alpha=0.3) # Add info text num_timesteps = len(lr_history['timesteps']) reward_improvement = rewards[-1] - rewards[0] if len(rewards) > 1 else 0 plt.text(0.02, 0.98, f'Total timesteps: {num_timesteps}\nTotal grad steps: {len(rewards)}\n' f'Initial reward: {rewards[0]:.4f}\nFinal reward: {rewards[-1]:.4f}\n' f'Improvement: {reward_improvement:+.4f}', transform=plt.gca().transAxes, fontsize=10, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.5)) plt.tight_layout() plt.savefig(plot_path, dpi=150, bbox_inches='tight') plt.close() print(f"\n✓ Saved Rewards curve plot to: {plot_path}") print(f" Total gradient steps: {len(rewards)}") print(f" Reward range: {min(rewards):.4f} → {max(rewards):.4f}") print(f" Total improvement: {reward_improvement:+.4f}") # ---> NEW: PLOT TRAJECTORY DIVERGENCE (MANIFOLD DRIFT) <--- if args.mode == "both" and 'baseline_trajectory' in locals() and 'guided_trajectory' in locals(): if len(baseline_trajectory) == len(guided_trajectory) and len(baseline_trajectory) > 0: print("\n" + "="*70) print("7. CALCULATING TRAJECTORY DIVERGENCE (THEOREM 1 & 2)") print("="*70) drift_path = Path(args.output_dir) / "trajectory_drift.png" l2_distances = [] # Calculate L2 norm ||z_t_guided - z_t_base||_2 for each step for b_lat, g_lat in zip(baseline_trajectory, guided_trajectory): dist = torch.norm(g_lat.float() - b_lat.float(), p=2).item() l2_distances.append(dist) steps = list(range(len(l2_distances))) plt.figure(figsize=(10, 6)) plt.plot(steps, l2_distances, linewidth=2.5, color='purple', marker='o', markersize=4) plt.xlabel('Denoising Step', fontsize=12) plt.ylabel('L2 Distance: ||z_guided - z_base||_2', fontsize=12) plt.title('Latent Trajectory Divergence (Manifold Drift)', fontsize=14, fontweight='bold') plt.grid(True, alpha=0.3) # Add interpretation text based on your theory max_drift = max(l2_distances) plt.text(0.02, 0.98, f'Max Drift: {max_drift:.4f}\n' f'Final Drift: {l2_distances[-1]:.4f}\n' f'(Matches bounded drift from Thm 1\n' f'or ODE stiffness collapse from Thm 2)', transform=plt.gca().transAxes, fontsize=10, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='thistle', alpha=0.5)) plt.tight_layout() plt.savefig(drift_path, dpi=150, bbox_inches='tight') plt.close() print(f"? Saved Manifold Drift curve to: {drift_path}") print(f" Max L2 Distance from baseline: {max_drift:.4f}") # ========== FINAL RESULTS ========== print("\n" + "="*70) print("FINAL RESULTS") print("="*70) if avg_reward_baseline is not None: print(f"\nBaseline:") if fid_score_baseline is not None: print(f" FID Score: {fid_score_baseline:.4f}") print(f" Avg Reward: {avg_reward_baseline:.4f}") if "clip" in args.metrics and clip_score_baseline is not None: print(f" Avg CLIP Score: {clip_score_baseline:.4f}") if "aesthetic" in args.metrics and aesthetic_score_baseline is not None: print(f" Avg Aesthetic: {aesthetic_score_baseline:.4f}") if "pickscore" in args.metrics and pick_score_baseline is not None: print(f" Avg PickScore: {pick_score_baseline:.4f}") if "hpsv2" in args.metrics and hpsv2_score_baseline is not None: print(f" Avg HPSv2: {hpsv2_score_baseline:.4f}") if "hpsv21" in args.metrics and hpsv21_score_baseline is not None: print(f" Avg HPSv2.1: {hpsv21_score_baseline:.4f}") if "imagereward" in args.metrics and imagereward_score_baseline is not None: print(f" Avg ImageReward: {imagereward_score_baseline:.4f}") if avg_reward_grad is not None: print(f"\nGradient Ascent:") if fid_score_grad is not None: print(f" FID Score: {fid_score_grad:.4f}") print(f" Avg Reward: {avg_reward_grad:.4f}") if "clip" in args.metrics and clip_score_grad is not None: print(f" Avg CLIP Score: {clip_score_grad:.4f}") if "aesthetic" in args.metrics and aesthetic_score_grad is not None: print(f" Avg Aesthetic: {aesthetic_score_grad:.4f}") if "pickscore" in args.metrics and pick_score_grad is not None: print(f" Avg PickScore: {pick_score_grad:.4f}") if "hpsv2" in args.metrics and hpsv2_score_grad is not None: print(f" Avg HPSv2: {hpsv2_score_grad:.4f}") if "hpsv21" in args.metrics and hpsv21_score_grad is not None: print(f" Avg HPSv2.1: {hpsv21_score_grad:.4f}") if "imagereward" in args.metrics and imagereward_score_grad is not None: print(f" Avg ImageReward: {imagereward_score_grad:.4f}") if avg_reward_baseline is not None and avg_reward_grad is not None: print(f"\nComparison:") if fid_score_baseline is not None and fid_score_grad is not None: fid_diff = fid_score_grad - fid_score_baseline print(f" FID Change: {fid_diff:+.4f} ({'worse' if fid_diff > 0 else 'better'}, lower is better)") reward_diff = avg_reward_grad - avg_reward_baseline print(f" Reward Change: {reward_diff:+.4f} ({'better' if reward_diff > 0 else 'worse'}, higher is better)") if "clip" in args.metrics and clip_score_baseline is not None and clip_score_grad is not None: clip_diff = clip_score_grad - clip_score_baseline print(f" CLIP Change: {clip_diff:+.4f} ({'better' if clip_diff > 0 else 'worse'}, higher is better)") if "aesthetic" in args.metrics and aesthetic_score_baseline is not None and aesthetic_score_grad is not None: aesthetic_diff = aesthetic_score_grad - aesthetic_score_baseline print(f" Aesthetic Change: {aesthetic_diff:+.4f} ({'better' if aesthetic_diff > 0 else 'worse'}, higher is better)") if "pickscore" in args.metrics and pick_score_baseline is not None and pick_score_grad is not None: pick_diff = pick_score_grad - pick_score_baseline print(f" PickScore Change: {pick_diff:+.4f} ({'better' if pick_diff > 0 else 'worse'}, higher is better)") if "hpsv2" in args.metrics and hpsv2_score_baseline is not None and hpsv2_score_grad is not None: hpsv2_diff = hpsv2_score_grad - hpsv2_score_baseline print(f" HPSv2 Change: {hpsv2_diff:+.4f} ({'better' if hpsv2_diff > 0 else 'worse'}, higher is better)") if "hpsv21" in args.metrics and hpsv21_score_baseline is not None and hpsv21_score_grad is not None: hpsv21_diff = hpsv21_score_grad - hpsv21_score_baseline print(f" HPSv2.1 Change: {hpsv21_diff:+.4f} ({'better' if hpsv21_diff > 0 else 'worse'}, higher is better)") if "imagereward" in args.metrics and imagereward_score_baseline is not None and imagereward_score_grad is not None: imagereward_diff = imagereward_score_grad - imagereward_score_baseline print(f" ImageReward Chg: {imagereward_diff:+.4f} ({'better' if imagereward_diff > 0 else 'worse'}, higher is better)") # Save results to file results = { "mode": args.mode, "metrics": args.metrics, "config": { "num_samples": len(prompts), "num_steps": args.num_steps, "cfg_scale": args.cfg_scale, "grad_range": [args.grad_range_start, args.grad_range_end], "grad_steps": args.grad_steps, "grad_step_size": args.grad_step_size } } if avg_reward_baseline is not None: results["baseline"] = {"avg_reward": avg_reward_baseline} if fid_score_baseline is not None: results["baseline"]["fid"] = fid_score_baseline if "clip" in args.metrics and clip_score_baseline is not None: results["baseline"]["clip_score"] = clip_score_baseline if "aesthetic" in args.metrics and aesthetic_score_baseline is not None: results["baseline"]["aesthetic_score"] = aesthetic_score_baseline if "pickscore" in args.metrics and pick_score_baseline is not None: results["baseline"]["pickscore"] = pick_score_baseline if "hpsv2" in args.metrics and hpsv2_score_baseline is not None: results["baseline"]["hpsv2_score"] = hpsv2_score_baseline if "hpsv21" in args.metrics and hpsv21_score_baseline is not None: results["baseline"]["hpsv21_score"] = hpsv21_score_baseline if "imagereward" in args.metrics and imagereward_score_baseline is not None: results["baseline"]["imagereward_score"] = imagereward_score_baseline if avg_reward_grad is not None: results["gradient_ascent"] = {"avg_reward": avg_reward_grad} if fid_score_grad is not None: results["gradient_ascent"]["fid"] = fid_score_grad if "clip" in args.metrics and clip_score_grad is not None: results["gradient_ascent"]["clip_score"] = clip_score_grad if "aesthetic" in args.metrics and aesthetic_score_grad is not None: results["gradient_ascent"]["aesthetic_score"] = aesthetic_score_grad if "pickscore" in args.metrics and pick_score_grad is not None: results["gradient_ascent"]["pickscore"] = pick_score_grad if "hpsv2" in args.metrics and hpsv2_score_grad is not None: results["gradient_ascent"]["hpsv2_score"] = hpsv2_score_grad if "hpsv21" in args.metrics and hpsv21_score_grad is not None: results["gradient_ascent"]["hpsv21_score"] = hpsv21_score_grad if "imagereward" in args.metrics and imagereward_score_grad is not None: results["gradient_ascent"]["imagereward_score"] = imagereward_score_grad if grad_stats: results["gradient_ascent"]["stats"] = grad_stats if avg_reward_baseline is not None and avg_reward_grad is not None: results["comparison"] = { "reward_difference": avg_reward_grad - avg_reward_baseline } if fid_score_baseline is not None and fid_score_grad is not None: results["comparison"]["fid_difference"] = fid_score_grad - fid_score_baseline if "clip" in args.metrics and clip_score_baseline is not None and clip_score_grad is not None: results["comparison"]["clip_difference"] = clip_score_grad - clip_score_baseline if "aesthetic" in args.metrics and aesthetic_score_baseline is not None and aesthetic_score_grad is not None: results["comparison"]["aesthetic_difference"] = aesthetic_score_grad - aesthetic_score_baseline if "pickscore" in args.metrics and pick_score_baseline is not None and pick_score_grad is not None: results["comparison"]["pickscore_difference"] = pick_score_grad - pick_score_baseline if "hpsv2" in args.metrics and hpsv2_score_baseline is not None and hpsv2_score_grad is not None: results["comparison"]["hpsv2_difference"] = hpsv2_score_grad - hpsv2_score_baseline if "hpsv21" in args.metrics and hpsv21_score_baseline is not None and hpsv21_score_grad is not None: results["comparison"]["hpsv21_difference"] = hpsv21_score_grad - hpsv21_score_baseline if "imagereward" in args.metrics and imagereward_score_baseline is not None and imagereward_score_grad is not None: results["comparison"]["imagereward_difference"] = imagereward_score_grad - imagereward_score_baseline # Save results to output directory output_path = Path(args.output_dir) output_path.mkdir(parents=True, exist_ok=True) results_path = output_path / "evaluation_results.txt" with open(results_path, "w") as f: for k, v in results.items(): f.write(f"{k}: {v}\n") print(f"\n✓ Results saved to: {results_path}") if args.save_images: print(f"✓ Generated images saved to: {output_path}/baseline/ and {output_path}/gradient_ascent/") print("\n" + "="*70) # Close logger tee_logger.close() sys.stdout = tee_logger.terminal if __name__ == "__main__": main()