Add files using upload-large-folder tool

Reward_sd15_idealized/timestep_convergence_analysis.ipynb

@@ -0,0 +1,1105 @@
+{
+ "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
+}

Reward_sd15_idealized/tune_parallel.sh

@@ -0,0 +1,253 @@
+#!/bin/bash
+
+# Parallel hyperparameter tuning across 8 GPUs
+# This script distributes experiments evenly across all available GPUs
+
+clear
+
+# Activate conda environment
+source ~/miniconda3/etc/profile.d/conda.sh
+conda activate /home/ec2-user/aev
+
+# Configuration
+DATASET_TYPE="pickapic"  # "coco" or "pickapic"
+MODEL_VARIANT="lpo"      # "origin", "spo", "diffusion_dpo", or "lpo"
+MAX_SAMPLES=500           # Number of samples for tuning
+NUM_STEPS=50             # Fixed inference steps
+SEARCH_TYPE="grid"       # "grid" or "random"
+OUTPUT_DIR="RESULTS_TURNING/run_2"
+NUM_GPUS=8               # Number of GPUs to use
+
+echo "=============================================="
+echo "  PARALLEL HYPERPARAMETER TUNING"
+echo "=============================================="
+echo ""
+echo "Configuration:"
+echo "  Dataset: $DATASET_TYPE"
+echo "  Model: $MODEL_VARIANT"
+echo "  Samples: $MAX_SAMPLES"
+echo "  Inference Steps: $NUM_STEPS"
+echo "  Search Type: $SEARCH_TYPE"
+echo "  GPUs: $NUM_GPUS"
+echo "  Output: $OUTPUT_DIR"
+echo ""
+
+# First, calculate total number of experiments
+echo "Calculating total experiments..."
+TOTAL_CONFIGS=$(python -c "
+from tune_hyperparams import HyperparameterTuner
+import sys
+tuner = HyperparameterTuner()
+configs = tuner.define_search_space()
+sys.stderr.write(f'Generated {len(configs)} configurations\n')
+print(len(configs))
+" 2>&1 | tail -1)
+
+echo "Total configurations: $TOTAL_CONFIGS"
+echo ""
+
+# Calculate experiments per GPU
+CONFIGS_PER_GPU=$((TOTAL_CONFIGS / NUM_GPUS))
+REMAINDER=$((TOTAL_CONFIGS % NUM_GPUS))
+
+echo "Distributing work:"
+echo "  Base configs per GPU: $CONFIGS_PER_GPU"
+echo "  Extra configs for first GPUs: $REMAINDER"
+echo ""
+
+# Create output directory
+mkdir -p "$OUTPUT_DIR"
+
+# Array to store background process IDs
+PIDS=()
+
+# Launch parallel processes on each GPU
+for GPU_ID in $(seq 0 $((NUM_GPUS - 1))); do
+    # Calculate start and end indices for this GPU
+    START_IDX=$((GPU_ID * CONFIGS_PER_GPU))
+    
+    # Give extra configs to first GPUs
+    if [ $GPU_ID -lt $REMAINDER ]; then
+        START_IDX=$((START_IDX + GPU_ID))
+        END_IDX=$((START_IDX + CONFIGS_PER_GPU + 1))
+    else
+        START_IDX=$((START_IDX + REMAINDER))
+        END_IDX=$((START_IDX + CONFIGS_PER_GPU))
+    fi
+    
+    # Create GPU-specific output directory
+    GPU_OUTPUT_DIR="${OUTPUT_DIR}/gpu_${GPU_ID}"
+    mkdir -p "$GPU_OUTPUT_DIR"
+    
+    echo "GPU $GPU_ID: configs $START_IDX to $END_IDX"
+    
+    # Launch tuning process in background
+    nohup python tune_hyperparams.py \
+        --output_dir "$GPU_OUTPUT_DIR" \
+        --max_samples $MAX_SAMPLES \
+        --num_steps $NUM_STEPS \
+        --dataset_type "$DATASET_TYPE" \
+        --model_variant "$MODEL_VARIANT" \
+        --cuda $GPU_ID \
+        --search_type "$SEARCH_TYPE" \
+        --start_idx $START_IDX \
+        --end_idx $END_IDX \
+        --metrics clip aesthetic pickscore hpsv2 imagereward \
+        > "${GPU_OUTPUT_DIR}/tuning.log" 2>&1 &
+    
+    # Store PID
+    PIDS+=($!)
+    
+    echo "  Launched with PID: ${PIDS[$GPU_ID]}"
+    
+    # Small delay to avoid race conditions
+    sleep 2
+done
+
+echo ""
+echo "=============================================="
+echo "  ALL PROCESSES LAUNCHED"
+echo "=============================================="
+echo ""
+echo "Background processes running:"
+for GPU_ID in $(seq 0 $((NUM_GPUS - 1))); do
+    echo "  GPU $GPU_ID: PID ${PIDS[$GPU_ID]} -> ${OUTPUT_DIR}/gpu_${GPU_ID}/tuning.log"
+done
+echo ""
+echo "To monitor progress:"
+echo "  tail -f ${OUTPUT_DIR}/gpu_0/tuning.log"
+echo "  tail -f ${OUTPUT_DIR}/gpu_1/tuning.log"
+echo "  ... etc"
+echo ""
+echo "To check all GPU processes:"
+echo "  ps aux | grep tune_hyperparams.py"
+echo ""
+echo "To monitor GPU usage:"
+echo "  watch -n 1 nvidia-smi"
+echo ""
+echo "To kill all processes:"
+echo "  kill ${PIDS[@]}"
+echo ""
+echo "Waiting for all processes to complete..."
+echo "(Press Ctrl+C to stop waiting, processes will continue in background)"
+echo ""
+
+# Wait for all background processes
+for PID in "${PIDS[@]}"; do
+    wait $PID
+done
+
+echo ""
+echo "=============================================="
+echo "  ALL TUNING PROCESSES COMPLETE"
+echo "=============================================="
+echo ""
+
+# Merge results from all GPUs
+echo "Merging results from all GPUs..."
+
+# Activate conda environment for Python script
+source ~/miniconda3/etc/profile.d/conda.sh
+conda activate /home/ec2-user/aev
+
+python - <<'EOF'
+import json
+from pathlib import Path
+import sys
+
+output_dir = Path("RESULTS_TURNING")
+all_results = []
+baseline_result = None
+
+# Collect results from each GPU
+for gpu_id in range(8):
+    gpu_dir = output_dir / f"gpu_{gpu_id}"
+    results_file = gpu_dir / "tuning_results.json"
+    
+    if results_file.exists():
+        with open(results_file, 'r') as f:
+            data = json.load(f)
+            
+        # Get baseline (should be same from all)
+        if baseline_result is None and "baseline" in data:
+            baseline_result = data["baseline"]
+        
+        # Collect experiments
+        if "experiments" in data:
+            all_results.extend(data["experiments"])
+        
+        print(f"GPU {gpu_id}: {len(data.get('experiments', []))} results")
+
+# Merge all results
+merged_data = {
+    "baseline": baseline_result,
+    "experiments": all_results,
+    "num_gpus": 8,
+    "total_experiments": len(all_results)
+}
+
+# Save merged results
+merged_file = output_dir / "merged_results.json"
+with open(merged_file, 'w') as f:
+    json.dump(merged_data, f, indent=2)
+
+print(f"\nMerged {len(all_results)} total results")
+print(f"Saved to: {merged_file}")
+
+# Find best configuration
+successful = [r for r in all_results if "metrics" in r]
+if successful:
+    # Compute aggregate scores
+    def compute_score(metrics):
+        weights = {
+            "reward": 1.0, "clip": 0.8, "aesthetic": 0.8,
+            "pickscore": 1.0, "hpsv2": 1.0, "imagereward": 1.0,
+            "fid": -0.5
+        }
+        score = sum(weights.get(k, 0) * v for k, v in metrics.items())
+        return score / sum(abs(w) for w in weights.values())
+    
+    for r in successful:
+        r["aggregate_score"] = compute_score(r["metrics"])
+    
+    successful.sort(key=lambda x: x["aggregate_score"], reverse=True)
+    
+    best = successful[0]
+    best_file = output_dir / "best_config.json"
+    with open(best_file, 'w') as f:
+        json.dump({
+            "config": best["config"],
+            "metrics": best["metrics"],
+            "aggregate_score": best["aggregate_score"],
+            "improvements": best.get("improvements", {})
+        }, f, indent=2)
+    
+    print(f"\n{'='*60}")
+    print("BEST CONFIGURATION:")
+    print(f"{'='*60}")
+    print(json.dumps(best["config"], indent=2))
+    print(f"\nAggregate Score: {best['aggregate_score']:.4f}")
+    print(f"Saved to: {best_file}")
+else:
+    print("\nNo successful experiments found!")
+    sys.exit(1)
+EOF
+
+if [ $? -eq 0 ]; then
+    echo ""
+    echo "=============================================="
+    echo "  TUNING COMPLETE!"
+    echo "=============================================="
+    echo ""
+    echo "Results:"
+    echo "  Merged results: ${OUTPUT_DIR}/merged_results.json"
+    echo "  Best config: ${OUTPUT_DIR}/best_config.json"
+    echo ""
+    echo "View best configuration:"
+    echo "  cat ${OUTPUT_DIR}/best_config.json"
+    echo ""
+else
+    echo ""
+    echo "ERROR: Failed to merge results"
+    exit 1
+fi

evaluation/blip/blip.py

@@ -0,0 +1,70 @@
+'''
+ * Adapted from BLIP (https://github.com/salesforce/BLIP)
+'''
+
+import warnings
+warnings.filterwarnings("ignore")
+
+import torch
+import os
+from urllib.parse import urlparse
+from timm.models.hub import download_cached_file
+from transformers import BertTokenizer
+from .vit import VisionTransformer, interpolate_pos_embed
+
+
+def init_tokenizer():
+    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+    tokenizer.add_special_tokens({'bos_token':'[DEC]'})
+    tokenizer.add_special_tokens({'additional_special_tokens':['[ENC]']})       
+    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]  
+    return tokenizer
+
+
+def create_vit(vit, image_size, use_grad_checkpointing=False, ckpt_layer=0, drop_path_rate=0):
+        
+    assert vit in ['base', 'large'], "vit parameter must be base or large"
+    if vit=='base':
+        vision_width = 768
+        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=12, 
+                                           num_heads=12, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
+                                           drop_path_rate=0 or drop_path_rate
+                                          )   
+    elif vit=='large':
+        vision_width = 1024
+        visual_encoder = VisionTransformer(img_size=image_size, patch_size=16, embed_dim=vision_width, depth=24, 
+                                           num_heads=16, use_grad_checkpointing=use_grad_checkpointing, ckpt_layer=ckpt_layer,
+                                           drop_path_rate=0.1 or drop_path_rate
+                                          )   
+    return visual_encoder, vision_width
+
+
+def is_url(url_or_filename):
+    parsed = urlparse(url_or_filename)
+    return parsed.scheme in ("http", "https")
+
+def load_checkpoint(model,url_or_filename):
+    if is_url(url_or_filename):
+        cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
+        checkpoint = torch.load(cached_file, map_location='cpu') 
+    elif os.path.isfile(url_or_filename):        
+        checkpoint = torch.load(url_or_filename, map_location='cpu') 
+    else:
+        raise RuntimeError('checkpoint url or path is invalid')
+        
+    state_dict = checkpoint['model']
+    
+    state_dict['visual_encoder.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder.pos_embed'],model.visual_encoder) 
+    if 'visual_encoder_m.pos_embed' in model.state_dict().keys():
+        state_dict['visual_encoder_m.pos_embed'] = interpolate_pos_embed(state_dict['visual_encoder_m.pos_embed'],
+                                                                         model.visual_encoder_m)    
+    for key in model.state_dict().keys():
+        if key in state_dict.keys():
+            if state_dict[key].shape!=model.state_dict()[key].shape:
+                print(key, ": ", state_dict[key].shape, ', ', model.state_dict()[key].shape)
+                del state_dict[key]
+    
+    msg = model.load_state_dict(state_dict,strict=False)
+    print('load checkpoint from %s'%url_or_filename)  
+    return model,msg
+    

evaluation/blip/blip_pretrain.py

@@ -0,0 +1,43 @@
+'''
+ * Adapted from BLIP (https://github.com/salesforce/BLIP)
+'''
+
+import transformers
+transformers.logging.set_verbosity_error()
+
+from torch import nn
+import os
+from .med import BertConfig, BertModel
+from .blip import create_vit, init_tokenizer
+
+class BLIP_Pretrain(nn.Module):
+    def __init__(self,                 
+                 med_config = "med_config.json",  
+                 image_size = 224,
+                 vit = 'base',
+                 vit_grad_ckpt = False,
+                 vit_ckpt_layer = 0,                    
+                 embed_dim = 256,     
+                 queue_size = 57600,
+                 momentum = 0.995,
+                 ):
+        """
+        Args:
+            med_config (str): path for the mixture of encoder-decoder model's configuration file
+            image_size (int): input image size
+            vit (str): model size of vision transformer
+        """               
+        super().__init__()
+        
+        self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, 0)
+        
+        self.tokenizer = init_tokenizer()   
+        encoder_config = BertConfig.from_json_file(med_config)
+        encoder_config.encoder_width = vision_width
+        self.text_encoder = BertModel(config=encoder_config, add_pooling_layer=False)
+
+        text_width = self.text_encoder.config.hidden_size
+        
+        self.vision_proj = nn.Linear(vision_width, embed_dim)
+        self.text_proj = nn.Linear(text_width, embed_dim)
+

upload.py

@@ -201,10 +201,11 @@ def main() -> None:
     api = HfApi()
     repo_id = resolve_repo_id(api, args)
     ignore_patterns = build_ignore_patterns(args.extra_ignore)
-    total_size = folder_size_bytes(source_dir)
-    total_size_gb = total_size / (1024 ** 3)
+    total_size_gb = None
 
     if args.method == "auto":
+        total_size = folder_size_bytes(source_dir)
+        total_size_gb = total_size / (1024 ** 3)
         use_large = total_size_gb >= args.large_threshold_gb
     else:
         use_large = args.method == "large"
@@ -215,7 +216,10 @@ def main() -> None:
     print("Private:", args.private)
     print("Revision:", args.revision)
     print("Ignore patterns:", ignore_patterns)
-    print(f"Folder size (excluding .back): {total_size_gb:.2f} GB")
+    if total_size_gb is None:
+        print("Folder size scan: skipped (set --method auto to enable size-based selection)")
+    else:
+        print(f"Folder size (excluding .back): {total_size_gb:.2f} GB")
     print("Upload method:", "large" if use_large else "folder")
 
     if args.dry_run: