Reward_sd15_idealized/eval.py

"""
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()