{ "cells": [ { "cell_type": "markdown", "id": "513b682a", "metadata": {}, "source": [ "#### Setup and Imports" ] }, { "cell_type": "code", "execution_count": null, "id": "f9f44d61", "metadata": {}, "outputs": [], "source": [ "import os\n", "import sys\n", "import json\n", "import warnings\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib\n", "import torch\n", "import torch.nn as nn\n", "from pathlib import Path\n", "from PIL import Image\n", "from tqdm.auto import tqdm\n", "from diffusers import StableDiffusionPipeline, DDIMScheduler, UNet2DConditionModel\n", "from transformers import CLIPModel, CLIPProcessor\n", "from torchmetrics.image.fid import FrechetInceptionDistance\n", "from torchmetrics.multimodal import CLIPScore\n", "\n", "warnings.filterwarnings(\"ignore\")\n", "\n", "# Import local modules\n", "from models import LRMRewardModel\n", "from pipelines.sd15_gradient_ascent_pipeline import StableDiffusionGradientAscentPipeline\n", "from grad_ascent_configs import get_config, list_configs\n", "\n", "# Import evaluation metrics\n", "sys.path.append('../evaluation')\n", "from pick_score import PickScorer\n", "from hpsv2_score import HPSv2Scorer\n", "from imagereward_score import load_imagereward\n" ] }, { "cell_type": "markdown", "id": "1740dd7c", "metadata": {}, "source": [ "#### Configuration" ] }, { "cell_type": "code", "execution_count": null, "id": "f1bc2b07", "metadata": {}, "outputs": [], "source": [ "# ============ CONFIGURATION ============\n", "\n", "# Dataset\n", "DATA_DIR = \"./data\"\n", "DATASET_TYPE = \"coco\" # \"coco\" or \"pickapic\"\n", "NUM_SAMPLES = 20 # Number of samples to analyze\n", "\n", "# Model\n", "BASE_MODEL = \"runwayml/stable-diffusion-v1-5\"\n", "MODEL_VARIANT = \"lpo\" # \"origin\", \"spo\", \"diffusion_dpo\", \"lpo\"\n", "LRM_MODEL = \"casiatao/LRM\"\n", "\n", "# Generation\n", "NUM_INFERENCE_STEPS = 100\n", "CFG_SCALE = 5.0\n", "SEED = 42\n", "BATCH_SIZE = 1\n", "\n", "# Gradient Ascent Config\n", "GRAD_CONFIG = \"low_to_high_nesterov\" # Use None for manual config, or specify preset name\n", "GRAD_RANGE_START = 0\n", "GRAD_RANGE_END = 500\n", "GRAD_STEPS = 1\n", "GRAD_STEP_SIZE = 0.1\n", "\n", "# Metrics to compute\n", "METRICS = [\"reward\", \"clip\", \"aesthetic\", \"pickscore\", \"hpsv2\", \"fid\"] # Add/remove as needed\n", "\n", "# Device\n", "CUDA_DEVICE = 0\n", "device = f\"cuda:{CUDA_DEVICE}\" if torch.cuda.is_available() else \"cpu\"\n", "dtype = torch.float16 if torch.cuda.is_available() else torch.float32\n", "\n", "# Output\n", "OUTPUT_DIR = \"timestep_analysis_results\"\n", "os.makedirs(OUTPUT_DIR, exist_ok=True)\n", "\n", "print(f\"Device: {device}\")\n", "print(f\"Dataset: {DATASET_TYPE}\")\n", "print(f\"Samples to analyze: {NUM_SAMPLES}\")\n", "print(f\"Metrics: {METRICS}\")\n", "print(f\"Output directory: {OUTPUT_DIR}\")" ] }, { "cell_type": "markdown", "id": "1b1b6d02", "metadata": {}, "source": [ "#### Load Dataset" ] }, { "cell_type": "code", "execution_count": null, "id": "a74b2816", "metadata": {}, "outputs": [], "source": [ "def load_validation_data(data_dir, max_samples=None):\n", " \"\"\"Load COCO validation prompts and image paths.\"\"\"\n", " data_dir = Path(data_dir)\n", " val_json = data_dir / \"coco\" / \"caption_val.json\"\n", " \n", " if not val_json.exists():\n", " raise FileNotFoundError(f\"Validation data not found at {val_json}\")\n", " \n", " with open(val_json, 'r') as f:\n", " data = json.load(f)\n", " \n", " print(f\"Loaded JSON with {len(data)} entries\")\n", " \n", " # Validate that image folder exists\n", " val_img_dir = data_dir / \"coco\" / \"images\" / \"val\"\n", " if not val_img_dir.exists():\n", " print(f\"Warning: Standard validation directory not found: {val_img_dir}\")\n", " \n", " # Parse data - img_path already contains \"images/val/\" prefix\n", " prompts = []\n", " image_paths = []\n", " \n", " for img_path, caption in data.items():\n", " # Try the path as given (relative to data_dir/coco/)\n", " full_path = data_dir / \"coco\" / img_path\n", " if full_path.exists():\n", " prompts.append(caption)\n", " image_paths.append(str(full_path))\n", " \n", " print(f\"Found {len(prompts)} valid image-caption pairs\")\n", " \n", " if len(prompts) == 0:\n", " print(f\"\\n⚠ WARNING: No valid images found!\")\n", " print(f\"Debug information:\")\n", " print(f\" JSON file: {val_json}\")\n", " print(f\" JSON entries: {len(data)}\")\n", " print(f\" Sample keys from JSON: {list(data.keys())[:3]}\")\n", " \n", " # Check if images exist at all\n", " coco_dir = data_dir / \"coco\"\n", " if coco_dir.exists():\n", " print(f\" COCO dir exists: {coco_dir}\")\n", " # List subdirectories\n", " subdirs = [d.name for d in coco_dir.iterdir() if d.is_dir()]\n", " print(f\" Subdirectories in COCO: {subdirs}\")\n", " \n", " # Try to find images\n", " if val_img_dir.exists():\n", " img_files = list(val_img_dir.glob(\"*.jpg\"))[:5]\n", " print(f\" Sample images in val dir: {[f.name for f in img_files]}\")\n", " \n", " if max_samples and len(prompts) > 0:\n", " prompts = prompts[:max_samples]\n", " image_paths = image_paths[:max_samples]\n", " \n", " return prompts, image_paths\n", "\n", "# Load data\n", "prompts, image_paths = load_validation_data(DATA_DIR, NUM_SAMPLES)\n", "print(f\"\\n✓ Loaded {len(prompts)} samples\")\n", "\n", "if len(prompts) > 0:\n", " print(f\"\\nSample prompts:\")\n", " for i, prompt in enumerate(prompts[:3]):\n", " print(f\" {i+1}. {prompt[:80]}...\")\n", " print(f\"\\nSample image paths:\")\n", " for i, path in enumerate(image_paths[:3]):\n", " print(f\" {i+1}. {path}\")\n", "else:\n", " print(\"\\n❌ ERROR: No samples loaded! Please check your data directory structure.\")\n", " print(\"Expected structure:\")\n", " print(\" ./data/coco/caption_val.json\")\n", " print(\" ./data/coco/images/val/*.jpg\")" ] }, { "cell_type": "markdown", "id": "5ceae64a", "metadata": {}, "source": [ "#### Load Models and Scorers" ] }, { "cell_type": "code", "execution_count": null, "id": "43ad1f56", "metadata": {}, "outputs": [], "source": [ "# ============ MLP for Aesthetic Scoring ============\n", "class MLP(nn.Module):\n", " def __init__(self):\n", " super().__init__()\n", " self.layers = nn.Sequential(\n", " nn.Linear(768, 1024),\n", " nn.Dropout(0.2),\n", " nn.Linear(1024, 128),\n", " nn.Dropout(0.2),\n", " nn.Linear(128, 64),\n", " nn.Dropout(0.1),\n", " nn.Linear(64, 16),\n", " nn.Linear(16, 1),\n", " )\n", " \n", " @torch.no_grad()\n", " def forward(self, embed):\n", " return self.layers(embed)\n", "\n", "class AestheticScorer(torch.nn.Module):\n", " def __init__(self, dtype, device):\n", " super().__init__()\n", " self.clip = CLIPModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n", " self.processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-large-patch14\")\n", " self.mlp = MLP()\n", " \n", " aesthetic_path = \"../evaluation/sac+logos+ava1-l14-linearMSE.pth\"\n", " if os.path.exists(aesthetic_path):\n", " state_dict = torch.load(aesthetic_path, map_location='cpu')\n", " self.mlp.load_state_dict(state_dict)\n", " \n", " self.dtype = dtype\n", " self.to(device)\n", " self.eval()\n", " \n", " @torch.no_grad()\n", " def __call__(self, images):\n", " if not isinstance(images, list):\n", " images = [images]\n", " inputs = self.processor(images=images, return_tensors=\"pt\", padding=True)\n", " inputs = {k: v.to(self.clip.device) for k, v in inputs.items()}\n", " image_embeds = self.clip.get_image_features(**inputs)\n", " image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)\n", " scores = self.mlp(image_embeds.float())\n", " return scores.squeeze().cpu().numpy()\n", "\n", "print(\"Loading models...\")" ] }, { "cell_type": "code", "execution_count": null, "id": "36a70595", "metadata": {}, "outputs": [], "source": [ "# Load Reward Model\n", "print(\"Loading reward model...\")\n", "reward_model = LRMRewardModel(\n", " pretrained_model_name_or_path=BASE_MODEL,\n", " lrm_model_path=LRM_MODEL,\n", " guidance_scale=CFG_SCALE,\n", " device=device\n", ")\n", "if dtype == torch.float16:\n", " reward_model = reward_model.half()\n", "reward_model.eval()\n", "print(\"✓ Reward model loaded\")\n", "\n", "# Load Pipeline\n", "print(\"\\nLoading diffusion pipeline...\")\n", "if MODEL_VARIANT == \"origin\":\n", " base_pipeline = StableDiffusionPipeline.from_pretrained(\n", " BASE_MODEL, torch_dtype=dtype, safety_checker=None\n", " )\n", "elif MODEL_VARIANT == \"spo\":\n", " base_pipeline = StableDiffusionPipeline.from_pretrained(\n", " 'SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep',\n", " torch_dtype=dtype, safety_checker=None\n", " )\n", " CFG_SCALE = 5.0\n", "elif MODEL_VARIANT == \"diffusion_dpo\":\n", " unet = UNet2DConditionModel.from_pretrained(\n", " 'mhdang/dpo-sd1.5-text2image-v1', subfolder=\"unet\", torch_dtype=dtype\n", " )\n", " base_pipeline = StableDiffusionPipeline.from_pretrained(\n", " BASE_MODEL, torch_dtype=dtype, safety_checker=None, unet=unet\n", " )\n", "elif MODEL_VARIANT == \"lpo\":\n", " unet = UNet2DConditionModel.from_pretrained(\n", " 'casiatao/LPO', subfolder=\"lpo_sd15_merge/unet\", torch_dtype=dtype\n", " )\n", " base_pipeline = StableDiffusionPipeline.from_pretrained(\n", " BASE_MODEL, torch_dtype=dtype, safety_checker=None, unet=unet\n", " )\n", " CFG_SCALE = 5.0\n", "\n", "pipeline = StableDiffusionGradientAscentPipeline(**base_pipeline.components)\n", "pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)\n", "pipeline = pipeline.to(device)\n", "pipeline.set_reward_model(reward_model)\n", "print(\"✓ Pipeline loaded\")" ] }, { "cell_type": "code", "execution_count": null, "id": "4e18f075", "metadata": {}, "outputs": [], "source": [ "# Load Metric Scorers\n", "print(\"\\nLoading metric scorers...\")\n", "\n", "clip_scorer = None\n", "aesthetic_scorer = None\n", "pick_scorer = None\n", "hpsv2_scorer = None\n", "imagereward_scorer = None\n", "\n", "if \"clip\" in METRICS:\n", " print(\" Loading CLIP scorer...\")\n", " clip_scorer = CLIPScore(model_name_or_path=\"openai/clip-vit-base-patch16\").to(device)\n", " print(\" ✓ CLIP scorer loaded\")\n", "\n", "if \"aesthetic\" in METRICS:\n", " print(\" Loading Aesthetic scorer...\")\n", " aesthetic_scorer = AestheticScorer(dtype, device)\n", " print(\" ✓ Aesthetic scorer loaded\")\n", "\n", "if \"pickscore\" in METRICS:\n", " print(\" Loading PickScore scorer...\")\n", " try:\n", " pick_scorer = PickScorer(device=device, dtype=dtype)\n", " print(\" ✓ PickScore loaded\")\n", " except Exception as e:\n", " print(f\" ✗ PickScore failed: {e}\")\n", " METRICS.remove(\"pickscore\")\n", "\n", "if \"hpsv2\" in METRICS:\n", " print(\" Loading HPSv2 scorer...\")\n", " try:\n", " hpsv2_scorer = HPSv2Scorer(device=device, dtype=dtype)\n", " print(\" ✓ HPSv2 loaded\")\n", " except Exception as e:\n", " print(f\" ✗ HPSv2 failed: {e}\")\n", " METRICS.remove(\"hpsv2\")\n", "\n", "if \"imagereward\" in METRICS:\n", " print(\" Loading ImageReward scorer...\")\n", " try:\n", " imagereward_scorer = load_imagereward(device=device)\n", " print(\" ✓ ImageReward loaded\")\n", " except Exception as e:\n", " print(f\" ✗ ImageReward failed: {e}\")\n", " METRICS.remove(\"imagereward\")\n", "\n", "print(f\"\\n✓ Active metrics: {METRICS}\")" ] }, { "cell_type": "markdown", "id": "70ac047b", "metadata": {}, "source": [ "#### Configure Gradient Ascent" ] }, { "cell_type": "code", "execution_count": null, "id": "05996448", "metadata": {}, "outputs": [], "source": [ "# Configure gradient ascent\n", "if GRAD_CONFIG:\n", " print(f\"Loading gradient ascent config: {GRAD_CONFIG}\")\n", " grad_config = get_config(GRAD_CONFIG)\n", " print(f\"Config: {grad_config}\")\n", "else:\n", " grad_config = {\n", " \"grad_timestep_range\": (GRAD_RANGE_START, GRAD_RANGE_END),\n", " \"num_grad_steps\": GRAD_STEPS,\n", " \"grad_step_size\": GRAD_STEP_SIZE,\n", " }\n", " print(f\"Manual gradient ascent configuration: {grad_config}\")\n", "\n", "pipeline.enable_gradient_ascent(**grad_config)\n", "print(\"\\n✓ Gradient ascent enabled\")" ] }, { "cell_type": "markdown", "id": "1f82c3df", "metadata": {}, "source": [ "#### Timestep Analysis Functions" ] }, { "cell_type": "code", "execution_count": null, "id": "e836d8f2", "metadata": {}, "outputs": [], "source": [ "def latents_to_images(latents, vae):\n", " \"\"\"Convert latents to PIL images.\"\"\"\n", " latents = 1 / 0.18215 * latents\n", " with torch.no_grad():\n", " images = vae.decode(latents).sample\n", " images = (images / 2 + 0.5).clamp(0, 1)\n", " images = images.cpu().permute(0, 2, 3, 1).numpy()\n", " images = (images * 255).round().astype(\"uint8\")\n", " pil_images = [Image.fromarray(image) for image in images]\n", " return pil_images\n", "\n", "\n", "def compute_metrics_for_image(image, prompt, reference_image=None):\n", " \"\"\"Compute all metrics for a single image.\"\"\"\n", " metrics = {}\n", " \n", " # CLIP Score\n", " if clip_scorer is not None:\n", " img_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).unsqueeze(0).to(device)\n", " with torch.no_grad():\n", " clip_score = clip_scorer(img_tensor, prompt).item()\n", " metrics['clip'] = clip_score\n", " \n", " # Aesthetic Score\n", " if aesthetic_scorer is not None:\n", " aesthetic_score = aesthetic_scorer([image])\n", " if isinstance(aesthetic_score, np.ndarray):\n", " aesthetic_score = aesthetic_score.item()\n", " metrics['aesthetic'] = aesthetic_score\n", " \n", " # PickScore\n", " if pick_scorer is not None:\n", " pick_score = pick_scorer.score(prompt, [image])[0]\n", " metrics['pickscore'] = pick_score\n", " \n", " # HPSv2\n", " if hpsv2_scorer is not None:\n", " hpsv2_score = hpsv2_scorer.score(prompt, [image])[0]\n", " metrics['hpsv2'] = hpsv2_score\n", " \n", " # ImageReward\n", " if imagereward_scorer is not None:\n", " imagereward_score = imagereward_scorer.score(prompt, [image])[0]\n", " metrics['imagereward'] = imagereward_score\n", " \n", " # FID (if reference image provided)\n", " if reference_image is not None:\n", " try:\n", " fid_metric = FrechetInceptionDistance(normalize=True).to(device)\n", " \n", " # Process reference image\n", " ref_img = Image.open(reference_image).convert('RGB').resize((299, 299))\n", " ref_tensor = torch.from_numpy(np.array(ref_img)).permute(2, 0, 1).unsqueeze(0).to(device)\n", " \n", " # Process generated image\n", " gen_img = image.resize((299, 299))\n", " gen_tensor = torch.from_numpy(np.array(gen_img)).permute(2, 0, 1).unsqueeze(0).to(device)\n", " \n", " if ref_tensor.size(0) == 1:\n", " ref_tensor = ref_tensor.repeat(2, 1, 1, 1)\n", " if gen_tensor.size(0) == 1:\n", " gen_tensor = gen_tensor.repeat(2, 1, 1, 1)\n", " \n", " fid_metric.update(ref_tensor, real=True)\n", " fid_metric.update(gen_tensor, real=False)\n", " \n", " fid_score = fid_metric.compute().item()/10\n", " metrics['fid'] = fid_score\n", " except Exception as e:\n", " print(f\"FID computation failed: {e}\")\n", " \n", " return metrics\n", "\n", "\n", "def analyze_sample_timesteps(prompt, reference_image, sample_idx):\n", " \"\"\"\n", " Generate images and track metrics at each timestep.\n", " Returns timestep-wise metrics and intermediate images.\n", " \"\"\"\n", " print(f\"\\n{'='*70}\")\n", " print(f\"Analyzing Sample {sample_idx + 1}\")\n", " print(f\"Prompt: {prompt[:80]}...\")\n", " print(f\"{'='*70}\")\n", " \n", " # Storage for results\n", " timestep_metrics = {\n", " 'timesteps': [],\n", " 'reward': [],\n", " 'clip': [],\n", " 'aesthetic': [],\n", " 'pickscore': [],\n", " 'hpsv2': [],\n", " 'imagereward': [],\n", " 'fid': []\n", " }\n", " intermediate_images = []\n", " \n", " # Reset gradient stats\n", " if hasattr(pipeline, 'grad_guidance'):\n", " pipeline.grad_guidance.reset_statistics()\n", " \n", " # Modified pipeline call to capture intermediate latents\n", " generator = torch.Generator(device=device).manual_seed(SEED + sample_idx)\n", " \n", " # We'll manually step through the denoising process\n", " pipeline.set_progress_bar_config(disable=True)\n", " \n", " # Prepare inputs\n", " height = pipeline.unet.config.sample_size * pipeline.vae_scale_factor\n", " width = pipeline.unet.config.sample_size * pipeline.vae_scale_factor\n", " \n", " # Encode prompt\n", " text_embeddings = pipeline._encode_prompt(\n", " prompt, device, 1, True, None\n", " )\n", " \n", " # Prepare timesteps\n", " pipeline.scheduler.set_timesteps(NUM_INFERENCE_STEPS, device=device)\n", " timesteps = pipeline.scheduler.timesteps\n", " \n", " # Prepare latents\n", " shape = (1, pipeline.unet.config.in_channels, height // 8, width // 8)\n", " latents = torch.randn(shape, generator=generator, device=device, dtype=dtype)\n", " latents = latents * pipeline.scheduler.init_noise_sigma\n", " \n", " # Denoising loop with metric tracking\n", " for i, t in enumerate(tqdm(timesteps, desc=\"Denoising steps\")):\n", " # Apply gradient ascent if enabled\n", " if hasattr(pipeline, 'grad_guidance') and pipeline.grad_guidance:\n", " if pipeline.grad_guidance.should_apply_gradient(t.item()):\n", " latents, grad_stats = pipeline.grad_guidance.apply_gradient_ascent(\n", " latents, prompt, t.item(), verbose=False,\n", " total_denoising_steps=len(timesteps)\n", " )\n", " \n", " # Expand latents for classifier free guidance\n", " latent_model_input = torch.cat([latents] * 2)\n", " latent_model_input = pipeline.scheduler.scale_model_input(latent_model_input, t)\n", " \n", " # Predict noise\n", " with torch.no_grad():\n", " noise_pred = pipeline.unet(\n", " latent_model_input,\n", " t,\n", " encoder_hidden_states=text_embeddings,\n", " ).sample\n", " \n", " # Perform guidance\n", " noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n", " noise_pred = noise_pred_uncond + CFG_SCALE * (noise_pred_text - noise_pred_uncond)\n", " \n", " # Compute previous noisy sample\n", " latents = pipeline.scheduler.step(noise_pred, t, latents).prev_sample\n", " \n", " # Decode latents to image every few steps\n", " if i % 5 == 0 or i == len(timesteps) - 1:\n", " # Convert to image\n", " images = latents_to_images(latents, pipeline.vae)\n", " image = images[0]\n", " \n", " # Compute reward\n", " with torch.no_grad():\n", " reward = reward_model.get_reward_score(latents, prompt, t.item())\n", " reward_val = reward.mean().item() if reward.numel() > 1 else reward.item()\n", " \n", " # Compute other metrics\n", " metrics = compute_metrics_for_image(image, prompt, reference_image)\n", " \n", " # Store results\n", " timestep_metrics['timesteps'].append(t.item())\n", " timestep_metrics['reward'].append(reward_val)\n", " \n", " for metric_name in ['clip', 'aesthetic', 'pickscore', 'hpsv2', 'imagereward', 'fid']:\n", " if metric_name in metrics:\n", " timestep_metrics[metric_name].append(metrics[metric_name])\n", " else:\n", " timestep_metrics[metric_name].append(None)\n", " \n", " intermediate_images.append(image)\n", " \n", " print(f\" Step {i}/{len(timesteps)} | t={t.item():.0f} | Reward={reward_val:.4f}\")\n", " \n", " # Final image\n", " final_images = latents_to_images(latents, pipeline.vae)\n", " final_image = final_images[0]\n", " \n", " pipeline.set_progress_bar_config(disable=False)\n", " \n", " return timestep_metrics, intermediate_images, final_image\n", "\n", "print(\"✓ Analysis functions defined\")" ] }, { "cell_type": "markdown", "id": "fc089bfd", "metadata": {}, "source": [ "#### Run Timestep Analysis" ] }, { "cell_type": "code", "execution_count": null, "id": "67c25164", "metadata": {}, "outputs": [], "source": [ "# Run analysis for all samples\n", "all_results = []\n", "\n", "for idx in range(len(prompts)):\n", " prompt = prompts[idx]\n", " reference_image = image_paths[idx]\n", " \n", " # Analyze this sample\n", " metrics, images, final_image = analyze_sample_timesteps(prompt, reference_image, idx)\n", " \n", " # Store results\n", " all_results.append({\n", " 'prompt': prompt,\n", " 'reference_image': reference_image,\n", " 'metrics': metrics,\n", " 'intermediate_images': images,\n", " 'final_image': final_image\n", " })\n", " \n", " # Save intermediate results\n", " sample_dir = Path(OUTPUT_DIR) / f\"sample_{idx+1}\"\n", " sample_dir.mkdir(exist_ok=True)\n", " \n", " # Save final image\n", " final_image.save(sample_dir / \"final_image.png\")\n", " \n", " # Save all intermediate images\n", " images_dir = sample_dir / \"intermediate_images\"\n", " images_dir.mkdir(exist_ok=True)\n", " for img_idx, img in enumerate(images):\n", " t_val = metrics['timesteps'][img_idx]\n", " img.save(images_dir / f\"step_{img_idx:03d}_t{int(t_val)}.png\")\n", " \n", " # Save metrics\n", " with open(sample_dir / \"metrics.json\", 'w') as f:\n", " json.dump(metrics, f, indent=2)\n", " \n", " print(f\"✓ Saved {len(images)} intermediate images for sample {idx+1}\")\n", "\n", "print(\"\\n✓ Analysis complete for all samples\")" ] }, { "cell_type": "markdown", "id": "10dd749d", "metadata": {}, "source": [ "#### Visualization: Intermediate Images" ] }, { "cell_type": "code", "execution_count": null, "id": "bb32eaa1", "metadata": {}, "outputs": [], "source": [ "def plot_intermediate_images(results, sample_idx, max_images=8):\n", " \"\"\"Display intermediate images for a sample showing evolution over timesteps.\"\"\"\n", " result = results[sample_idx]\n", " images = result['intermediate_images']\n", " metrics = result['metrics']\n", " timesteps = metrics['timesteps']\n", " rewards = metrics['reward']\n", " \n", " # Select evenly spaced images if too many\n", " if len(images) > max_images:\n", " indices = np.linspace(0, len(images)-1, max_images, dtype=int)\n", " selected_images = [images[i] for i in indices]\n", " selected_timesteps = [timesteps[i] for i in indices]\n", " selected_rewards = [rewards[i] for i in indices]\n", " else:\n", " selected_images = images\n", " selected_timesteps = timesteps\n", " selected_rewards = rewards\n", " \n", " n_images = len(selected_images)\n", " cols = 5\n", " rows = (n_images + cols - 1) // cols\n", " \n", " fig, axes = plt.subplots(rows, cols, figsize=(4*cols, 4*rows))\n", " axes = axes.flatten() if n_images > 1 else [axes]\n", " \n", " fig.suptitle(f\"Sample {sample_idx + 1}: Image Evolution Over Timesteps\\n\"\n", " f\"Prompt: {result['prompt'][:80]}...\", \n", " fontsize=12, fontweight='bold')\n", " \n", " for idx, (img, t, r) in enumerate(zip(selected_images, selected_timesteps, selected_rewards)):\n", " ax = axes[idx]\n", " ax.imshow(img)\n", " ax.axis('off')\n", " ax.set_title(f\"t={t:.0f}\\nReward={r:.3f}\", fontsize=10)\n", " \n", " # Hide unused subplots\n", " for idx in range(n_images, len(axes)):\n", " axes[idx].axis('off')\n", " \n", " plt.tight_layout()\n", " \n", " # Save plot\n", " sample_dir = Path(OUTPUT_DIR) / f\"sample_{sample_idx+1}\"\n", " plt.savefig(sample_dir / \"image_evolution.png\", dpi=150, bbox_inches='tight')\n", " plt.show()\n", "\n", "# Plot intermediate images for all samples\n", "for idx in range(len(all_results)):\n", " plot_intermediate_images(all_results, idx)" ] }, { "cell_type": "code", "execution_count": null, "id": "878b7686", "metadata": {}, "outputs": [], "source": [ "def plot_final_images_grid(results):\n", " \"\"\"Display all final images in a grid for comparison.\"\"\"\n", " n_samples = len(results)\n", " cols = min(10, n_samples)\n", " rows = (n_samples + cols - 1) // cols\n", " \n", " fig, axes = plt.subplots(rows, cols, figsize=(5*cols, 5*rows))\n", " if n_samples == 1:\n", " axes = [axes]\n", " else:\n", " axes = axes.flatten()\n", " \n", " fig.suptitle(\"Final Generated Images: All Samples\", fontsize=14, fontweight='bold')\n", " \n", " for idx, result in enumerate(results):\n", " ax = axes[idx]\n", " ax.imshow(result['final_image'])\n", " ax.axis('off')\n", " \n", " # Get final metrics\n", " metrics = result['metrics']\n", " reward = metrics['reward'][-1] if metrics['reward'] else 0\n", " clip_score = metrics['clip'][-1] if 'clip' in metrics and metrics['clip'] and metrics['clip'][-1] is not None else 0\n", " \n", " ax.set_title(f\"Sample {idx+1}\\nReward: {reward:.3f} | CLIP: {clip_score:.3f}\\n{result['prompt'][:40]}...\", \n", " fontsize=9)\n", " \n", " # Hide unused subplots\n", " for idx in range(n_samples, len(axes)):\n", " axes[idx].axis('off')\n", " \n", " plt.tight_layout()\n", " plt.savefig(Path(OUTPUT_DIR) / \"final_images_grid.png\", dpi=150, bbox_inches='tight')\n", " plt.show()\n", "\n", "# Display final images\n", "plot_final_images_grid(all_results)" ] }, { "cell_type": "markdown", "id": "bc5a96a6", "metadata": {}, "source": [ "#### Debug: Check Data" ] }, { "cell_type": "code", "execution_count": null, "id": "40008ed4", "metadata": {}, "outputs": [], "source": [ "# Check if data was collected properly\n", "print(\"Data Collection Summary:\")\n", "print(\"=\"*70)\n", "\n", "for idx, result in enumerate(all_results):\n", " print(f\"\\nSample {idx+1}:\")\n", " print(f\" Prompt: {result['prompt'][:60]}...\")\n", " \n", " metrics = result['metrics']\n", " print(f\" Number of timesteps tracked: {len(metrics['timesteps'])}\")\n", " print(f\" Number of intermediate images: {len(result['intermediate_images'])}\")\n", " \n", " # Check which metrics have data\n", " for metric_name in ['reward', 'clip', 'aesthetic', 'pickscore', 'hpsv2', 'fid']:\n", " if metric_name in metrics:\n", " non_none = [v for v in metrics[metric_name] if v is not None]\n", " if non_none:\n", " print(f\" {metric_name.upper()}: {len(non_none)} values | \"\n", " f\"Range: [{min(non_none):.3f}, {max(non_none):.3f}]\")\n", " else:\n", " print(f\" {metric_name.upper()}: No valid data\")\n", " \n", " # Check timestep range\n", " if metrics['timesteps']:\n", " print(f\" Timestep range: [{max(metrics['timesteps']):.0f}, {min(metrics['timesteps']):.0f}]\")\n", "\n", "print(\"\\n\" + \"=\"*70)" ] }, { "cell_type": "markdown", "id": "5bdacf18", "metadata": {}, "source": [ "#### Visualization: Metrics Evolution" ] }, { "cell_type": "code", "execution_count": null, "id": "be2fc746", "metadata": {}, "outputs": [], "source": [ "def plot_metrics_evolution(results, sample_idx):\n", " \"\"\"Plot all metrics evolution in a single row for one sample.\"\"\"\n", " result = results[sample_idx]\n", " metrics = result['metrics']\n", " timesteps = metrics['timesteps']\n", " \n", " # Filter metrics to plot (exclude None values)\n", " metrics_to_plot = []\n", " for metric_name in ['reward', 'clip', 'aesthetic', 'pickscore', 'hpsv2', 'imagereward', 'fid']:\n", " if metric_name in metrics and any(v is not None for v in metrics[metric_name]):\n", " metrics_to_plot.append(metric_name)\n", " \n", " n_metrics = len(metrics_to_plot)\n", " \n", " # Create figure with subplots in a row\n", " fig, axes = plt.subplots(1, n_metrics, figsize=(5*n_metrics, 4))\n", " if n_metrics == 1:\n", " axes = [axes]\n", " \n", " fig.suptitle(f\"Sample {sample_idx + 1}: Metrics Evolution Across Timesteps\\n\"\n", " f\"Prompt: {result['prompt'][:80]}...\", fontsize=12, fontweight='bold')\n", " \n", " colors = ['blue', 'green', 'red', 'purple', 'orange', 'brown', 'pink']\n", " \n", " for idx, metric_name in enumerate(metrics_to_plot):\n", " ax = axes[idx]\n", " values = [v for v in metrics[metric_name] if v is not None]\n", " valid_timesteps = [t for t, v in zip(timesteps, metrics[metric_name]) if v is not None]\n", " \n", " if values:\n", " ax.plot(valid_timesteps, values, marker='o', linewidth=2, \n", " color=colors[idx % len(colors)], label=metric_name.upper())\n", " ax.set_xlabel('Timestep', fontsize=10)\n", " ax.set_ylabel(metric_name.upper(), fontsize=10)\n", " ax.set_title(f\"{metric_name.upper()}\\n{values[0]:.3f} → {values[-1]:.3f}\", fontsize=10)\n", " ax.grid(True, alpha=0.3)\n", " ax.invert_xaxis() # Timesteps go from high to low\n", " \n", " # Add improvement annotation\n", " improvement = values[-1] - values[0]\n", " color = 'green' if improvement > 0 else 'red'\n", " if metric_name == 'fid': # Lower is better for FID\n", " color = 'green' if improvement < 0 else 'red'\n", " ax.text(0.05, 0.95, f\"Δ: {improvement:+.3f}\", \n", " transform=ax.transAxes, fontsize=9, verticalalignment='top',\n", " bbox=dict(boxstyle='round', facecolor=color, alpha=0.3))\n", " \n", " plt.tight_layout()\n", " \n", " # Save plot\n", " sample_dir = Path(OUTPUT_DIR) / f\"sample_{sample_idx+1}\"\n", " plt.savefig(sample_dir / \"metrics_evolution.png\", dpi=150, bbox_inches='tight')\n", " plt.show()\n", "\n", "# Plot for all samples\n", "for idx in range(len(all_results)):\n", " plot_metrics_evolution(all_results, idx)" ] }, { "cell_type": "markdown", "id": "45b7abb7", "metadata": {}, "source": [ "#### Visualization: Compare All Samples" ] }, { "cell_type": "code", "execution_count": null, "id": "7d3df7e0", "metadata": {}, "outputs": [], "source": [ "def plot_all_samples_comparison(results):\n", " \"\"\"Plot metric evolution for all samples in a grid.\"\"\"\n", " # Choose key metrics to compare\n", " key_metrics = ['reward', 'clip', 'aesthetic', 'fid']\n", " n_metrics = len(key_metrics)\n", " n_samples = len(results)\n", " \n", " fig, axes = plt.subplots(n_metrics, 1, figsize=(14, 4*n_metrics))\n", " if n_metrics == 1:\n", " axes = [axes]\n", " \n", " fig.suptitle(\"Convergence Analysis: All Samples Comparison\", fontsize=14, fontweight='bold')\n", " \n", " colors = plt.cm.tab10(np.linspace(0, 1, n_samples))\n", " \n", " for metric_idx, metric_name in enumerate(key_metrics):\n", " ax = axes[metric_idx]\n", " \n", " for sample_idx, result in enumerate(results):\n", " metrics = result['metrics']\n", " timesteps = metrics['timesteps']\n", " values = [v for v in metrics[metric_name] if v is not None]\n", " valid_timesteps = [t for t, v in zip(timesteps, metrics[metric_name]) if v is not None]\n", " \n", " if values:\n", " ax.plot(valid_timesteps, values, marker='o', linewidth=2, \n", " color=colors[sample_idx], label=f\"Sample {sample_idx+1}\", alpha=0.7)\n", " \n", " ax.set_xlabel('Timestep', fontsize=11)\n", " ax.set_ylabel(metric_name.upper(), fontsize=11)\n", " ax.set_title(f\"{metric_name.upper()} Evolution\", fontsize=12, fontweight='bold')\n", " ax.grid(True, alpha=0.3)\n", " ax.invert_xaxis()\n", " ax.legend(loc='best', fontsize=9)\n", " \n", " plt.tight_layout()\n", " plt.savefig(Path(OUTPUT_DIR) / \"all_samples_comparison.png\", dpi=150, bbox_inches='tight')\n", " plt.show()\n", "\n", "# Plot comparison\n", "plot_all_samples_comparison(all_results)" ] }, { "cell_type": "markdown", "id": "44f97581", "metadata": {}, "source": [ "#### Convergence Analysis" ] }, { "cell_type": "code", "execution_count": null, "id": "9dffd878", "metadata": {}, "outputs": [], "source": [ "def analyze_convergence(results):\n", " \"\"\"Analyze convergence behavior across samples.\"\"\"\n", " print(\"\\n\" + \"=\"*70)\n", " print(\"CONVERGENCE ANALYSIS\")\n", " print(\"=\"*70)\n", " \n", " for metric_name in ['reward', 'clip', 'aesthetic', 'pickscore', 'hpsv2']:\n", " print(f\"\\n{metric_name.upper()} Convergence:\")\n", " print(\"-\" * 50)\n", " \n", " improvements = []\n", " initial_values = []\n", " final_values = []\n", " \n", " for idx, result in enumerate(results):\n", " metrics = result['metrics']\n", " if metric_name in metrics:\n", " values = [v for v in metrics[metric_name] if v is not None]\n", " if values:\n", " initial = values[0]\n", " final = values[-1]\n", " improvement = final - initial\n", " \n", " initial_values.append(initial)\n", " final_values.append(final)\n", " improvements.append(improvement)\n", " \n", " print(f\" Sample {idx+1}: {initial:.4f} → {final:.4f} ({improvement:+.4f})\")\n", " \n", " if improvements:\n", " avg_improvement = np.mean(improvements)\n", " std_improvement = np.std(improvements)\n", " print(f\"\\n Average Improvement: {avg_improvement:+.4f} (±{std_improvement:.4f})\")\n", " print(f\" Converged: {'YES' if std_improvement < 0.1 * abs(avg_improvement) else 'NO'}\")\n", " \n", " # Summary\n", " print(\"\\n\" + \"=\"*70)\n", " print(\"SUMMARY\")\n", " print(\"=\"*70)\n", " print(f\"Total samples analyzed: {len(results)}\")\n", " print(f\"Gradient ascent config: {grad_config}\")\n", " print(f\"\\nConclusion: Analyze the plots above to determine convergence behavior.\")\n", " print(f\"Look for:\")\n", " print(f\" 1. Metrics plateauing (flattening out)\")\n", " print(f\" 2. Consistent improvement across samples\")\n", " print(f\" 3. Low variance in final metric values\")\n", "\n", "analyze_convergence(all_results)" ] }, { "cell_type": "markdown", "id": "d263be5f", "metadata": {}, "source": [ "#### Save Results Summary" ] }, { "cell_type": "code", "execution_count": null, "id": "5434e7c0", "metadata": {}, "outputs": [], "source": [ "# Save comprehensive summary\n", "summary = {\n", " 'config': {\n", " 'num_samples': NUM_SAMPLES,\n", " 'num_inference_steps': NUM_INFERENCE_STEPS,\n", " 'cfg_scale': CFG_SCALE,\n", " 'grad_config': grad_config,\n", " 'metrics': METRICS,\n", " 'model_variant': MODEL_VARIANT\n", " },\n", " 'samples': []\n", "}\n", "\n", "for idx, result in enumerate(all_results):\n", " metrics = result['metrics']\n", " sample_summary = {\n", " 'sample_id': idx + 1,\n", " 'prompt': result['prompt'],\n", " 'reference_image': result['reference_image']\n", " }\n", " \n", " for metric_name in ['reward', 'clip', 'aesthetic', 'pickscore', 'hpsv2']:\n", " if metric_name in metrics:\n", " values = [v for v in metrics[metric_name] if v is not None]\n", " if values:\n", " sample_summary[metric_name] = {\n", " 'initial': values[0],\n", " 'final': values[-1],\n", " 'improvement': values[-1] - values[0],\n", " 'all_values': values\n", " }\n", " \n", " summary['samples'].append(sample_summary)\n", "\n", "# Save summary\n", "with open(Path(OUTPUT_DIR) / \"convergence_summary.json\", 'w') as f:\n", " json.dump(summary, f, indent=2)\n", "\n", "print(f\"\\n✓ Results saved to: {OUTPUT_DIR}\")\n", "print(f\" - convergence_summary.json\")\n", "print(f\" - all_samples_comparison.png\")\n", "print(f\" - sample_X/ directories with individual results\")" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.18" } }, "nbformat": 4, "nbformat_minor": 5 }