Add files using upload-large-folder tool

Reward_sd15_idealized/README.md

@@ -0,0 +1,1336 @@
+# Reward-Guided Gradient Ascent for Stable Diffusion
+
+A comprehensive system for improving Stable Diffusion image generation quality using gradient ascent optimization on Latent Reward Model (LRM) scores during inference.
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Features](#features)
+- [Installation](#installation)
+- [Quick Start](#quick-start)
+- [Architecture](#architecture)
+- [Understanding Reward Calculation](#understanding-reward-calculation)
+- [Learning Rate Scheduling](#learning-rate-scheduling)
+- [Configuration Presets](#configuration-presets)
+- [Evaluation Metrics](#evaluation-metrics)
+- [Model Variants](#model-variants)
+- [Datasets](#datasets)
+- [Usage Examples](#usage-examples)
+- [API Reference](#api-reference)
+- [Command-Line Options](#command-line-options)
+- [Output Files](#output-files)
+- [Troubleshooting](#troubleshooting)
+- [Best Practices](#best-practices)
+- [Changelog](#changelog)
+
+---
+
+## Overview
+
+This project implements **test-time optimization** for Stable Diffusion using gradient ascent on the LRM reward model. Unlike the main LPO training which uses the reward model for training, this approach applies it during inference to improve generation quality without retraining.
+
+### Key Capabilities
+
+- **Gradient Ascent Optimization**: Iteratively improve latents using reward gradients
+- **Learning Rate Scheduling**: Multiple strategies (constant, linear, cosine, exponential, step)
+- **Momentum Optimization**: Standard and Nesterov momentum for better convergence
+- **Multiple Metrics**: FID, CLIP, Aesthetic, PickScore, HPSv2, ImageReward
+- **Model Variants**: Support for Origin, SPO, DPO, and LPO SD1.5 models
+- **Dataset Flexibility**: COCO and Pick-a-Pic validation datasets
+- **Configuration Presets**: 15 pre-tuned configurations for various use cases
+
+---
+
+## Features
+
+### 1. **Advanced Optimization**
+- **5 LR Schedulers**: Constant, Linear, Cosine, Exponential, Step-wise
+- **Momentum Support**: Standard momentum and Nesterov momentum
+- **Configurable Timestep Ranges**: Apply gradients at specific denoising steps
+- **Dynamic Learning Rates**: LR changes during optimization for better convergence
+
+### 2. **Comprehensive Evaluation**
+- **6 Quality Metrics**: FID, CLIP, Aesthetic, PickScore, HPSv2, ImageReward
+- **Baseline Comparison**: Compare with and without gradient ascent
+- **Detailed Statistics**: Track reward improvements, gradient norms, LR history
+- **Batch Processing**: Efficient evaluation on large datasets
+- **Reward Visualization**: Automatic plotting of reward progression across timesteps
+- **Timestep-Aware Tracking**: Monitor rewards at every denoising step, final t=0 latent reported
+
+### 3. **Model Flexibility**
+- **4 SD1.5 Variants**: Origin, SPO, DPO, LPO
+- **Auto-Configuration**: CFG scale auto-adjusted for model variants
+- **Easy Switching**: Change models with a single flag
+
+### 4. **Dataset Support**
+- **COCO Validation**: Standard benchmark with reference images
+- **Pick-a-Pic Validation**: Large-scale human preference dataset
+- **Streaming Support**: Handle large datasets efficiently
+
+---
+
+## Installation
+
+### Requirements
+
+```bash
+# Core dependencies
+pip install torch diffusers transformers torchmetrics datasets huggingface-hub
+
+# For evaluation metrics
+pip install pillow numpy scipy tqdm
+
+# Optional: for better performance
+pip install xformers  # For memory-efficient attention
+```
+
+### Setup
+
+```bash
+cd /path/to/LPO/Reward
+
+# Verify installation
+python -c "from lr_scheduler import create_lr_scheduler; print('✓ LR Scheduler OK')"
+python -c "from grad_ascent_configs import list_configs; print('✓ Configs:', len(list_configs()))"
+python -c "from gradient_ascent_utils import RewardGuidedDiffusion; print('✓ Gradient Utils OK')"
+```
+
+---
+
+## Quick Start
+
+### 1. Basic COCO Evaluation (test_grad_sd1.5.py)
+
+```bash
+# Edit Config in test_grad_sd1.5.py:
+# - Set device: "cuda:0" or "cuda:6"
+# - Set max_samples: 10 for quick test, None for full dataset
+# - Configure gradient ascent parameters
+
+python test_grad_sd1.5.py
+```
+
+**Output:**
+- Creates `RESULTS/SD1.5_GradAscent/run_1/` (auto-incremented)
+- Generates `eval.log` with detailed metrics
+- Saves `reward_curve.png` showing reward progression
+
+### 2. Basic Evaluation with Preset Config (eval.py)
+
+```bash
+python eval.py \
+  --grad_config cosine_nesterov \
+  --metrics clip aesthetic \
+  --max_samples 10
+```
+
+### 2. High-Quality Evaluation
+
+```bash
+python eval.py \
+  --grad_config high_quality \
+  --metrics fid clip aesthetic pickscore hpsv2 \
+  --max_samples 100 \
+  --save_images \
+  --output_dir results/high_quality
+```
+
+### 3. Pick-a-Pic Benchmark
+
+```bash
+python eval.py \
+  --dataset_type pickapic \
+  --grad_config cosine_nesterov \
+  --metrics pickscore hpsv2 imagereward \
+  --max_samples 500 \
+  --output_dir results/pickapic
+```
+
+---
+
+## Architecture
+
+### System Components
+
+```
+Reward/
+├── models/
+│   ├── reward_model.py              # LRM reward model wrapper
+│   └── unet_2d_condition_reward.py  # Custom UNet with reward tracking
+├── pipelines/
+│   ├── sd15_reward_pipeline.py      # Base pipeline with reward tracking
+│   └── sd15_gradient_ascent_pipeline.py  # Pipeline with gradient ascent
+├── lr_scheduler.py                  # Learning rate schedulers
+├── gradient_ascent_utils.py         # Core gradient ascent implementation
+├── grad_ascent_configs.py           # Configuration presets
+├── eval.py                          # Comprehensive evaluation script
+└── examples.sh                      # Example commands
+```
+
+### Gradient Ascent Flow
+
+```
+1. Load Stable Diffusion + LRM Reward Model
+2. Start denoising process (T → 0)
+3. At each timestep t:
+   a. Standard denoising step (predict noise, remove it)
+   b. Compute reward R(latents, prompt, t) and store in history
+   c. If t in gradient range:
+      - Enable gradients on latents
+      - Compute ∇R w.r.t. latents
+      - For each gradient step:
+        * Get current LR from scheduler
+        * Apply momentum (if enabled)
+        * Update: latents += lr * momentum(∇R)
+      - Track statistics (grad norms, reward improvement)
+4. At final timestep (t=0):
+   - Final reward computed on clean latent
+   - This reward is reported in logs
+5. Decode final latent (x₀) to image via VAE
+6. Compute quality metrics on image
+```
+
+### Understanding Reward Calculation
+
+**Key Concepts:**
+
+- **Timestep-Aware Rewards**: The LRM reward model computes preference scores at ANY noise level (timestep t)
+- **Progressive Tracking**: Rewards are calculated at every denoising step throughout generation
+- **Final Latent Reward**: The reported metric is the reward for t=0 (the clean latent before decoding)
+- **Not Averaged**: The final reward is specifically from the last timestep, NOT an average across all timesteps
+
+**What gets reported:**
+```python
+# During generation: Rewards computed at each t (1000 → 0)
+Step 0: t=1000, reward=3.2
+Step 1: t=990,  reward=3.5
+...
+Step 99: t=10,  reward=5.1
+Step 100: t=0,  reward=5.4  ← This is what gets logged!
+```
+
+The `Reward (t=0)` in logs represents the preference score of the final clean latent that was decoded into your output image.
+
+---
+
+## Learning Rate Scheduling
+
+### Available Schedulers
+
+#### 1. **Constant LR**
+```python
+lr_scheduler_type="constant"
+```
+- Fixed learning rate throughout optimization
+- Simple and stable
+- Good for quick experiments
+
+#### 2. **Linear Decay**
+```python
+lr_scheduler_type="linear"
+lr_scheduler_kwargs={
+    "end_lr": 0.01,      # End LR (10% of initial)
+    "start_step": 0      # When to start decay
+}
+```
+- Linear decrease from initial to end LR
+- Smooth convergence
+- Configurable warmup period
+
+#### 3. **Cosine Annealing** (Recommended)
+```python
+lr_scheduler_type="cosine"
+lr_scheduler_kwargs={
+    "min_lr": 0.001,     # Minimum LR
+    "warmup_steps": 3    # Linear warmup steps
+}
+```
+- Smooth cosine decay
+- Optional warmup phase
+- Widely used in deep learning
+- **Best for most use cases**
+
+#### 4. **Exponential Decay**
+```python
+lr_scheduler_type="exponential"
+lr_scheduler_kwargs={
+    "gamma": 0.9         # Decay factor per step
+}
+```
+- Exponential decrease
+- Fast initial decay
+- Good for aggressive optimization
+
+#### 5. **Step Decay**
+```python
+lr_scheduler_type="step"
+lr_scheduler_kwargs={
+    "step_size": 5,      # Steps between decays
+    "gamma": 0.5         # Multiplicative factor
+}
+```
+- Step-wise LR reduction
+- Periodic decay
+- Good for scheduled changes
+
+### Usage Example
+
+```python
+from pipelines.sd15_gradient_ascent_pipeline import StableDiffusionGradientAscentPipeline
+
+pipeline.enable_gradient_ascent(
+    grad_timestep_range=(0, 700),
+    num_grad_steps=15,
+    grad_step_size=0.1,  # Initial LR
+    lr_scheduler_type="cosine",
+    lr_scheduler_kwargs={
+        "min_lr": 0.001,
+        "warmup_steps": 3
+    }
+)
+```
+
+---
+
+## Configuration Presets
+
+We provide 15 pre-configured optimization strategies. Use them with `--grad_config <name>`.
+
+### Basic Configurations
+
+| Config | LR Schedule | Momentum | Steps | Description |
+|--------|-------------|----------|-------|-------------|
+| `constant` | Constant | No | 5 | Simple baseline |
+| `linear` | Linear decay | No | 10 | Smooth decay |
+| `linear_warmstart` | Linear w/ warmup | No | 10 | Stable start |
+| `cosine` | Cosine | No | 10 | Smooth convergence |
+| `cosine_warmup` | Cosine w/ warmup | No | 20 | Best convergence |
+| `exponential` | Exponential | No | 15 | Fast decay |
+| `step` | Step-wise | No | 20 | Periodic decay |
+
+### Momentum Configurations
+
+| Config | LR Schedule | Momentum | Steps | Description |
+|--------|-------------|----------|-------|-------------|
+| `momentum` | Constant | Standard | 10 | Faster convergence |
+| `nesterov` | Constant | Nesterov | 10 | Better convergence |
+
+### Advanced Configurations
+
+| Config | LR Schedule | Momentum | Steps | Description |
+|--------|-------------|----------|-------|-------------|
+| `cosine_momentum` | Cosine | Standard | 15 | High quality |
+| `cosine_nesterov` | Cosine | Nesterov | 15 | **Recommended** |
+| `linear_nesterov` | Linear | Nesterov | 15 | Stable + fast |
+
+### Quality Presets
+
+| Config | LR Schedule | Momentum | Steps | Use Case |
+|--------|-------------|----------|-------|----------|
+| `high_quality` | Cosine | Nesterov | 20 | **Best quality** |
+| `aggressive` | Exponential | Standard | 8 | Fast results |
+| `conservative` | Cosine | Nesterov | 25 | Most stable |
+
+### Config Details
+
+#### `high_quality` (Recommended for Research)
+```python
+{
+    "grad_timestep_range": (200, 800),  # Focus on middle timesteps
+    "num_grad_steps": 20,
+    "grad_step_size": 0.08,
+    "lr_scheduler_type": "cosine",
+    "lr_scheduler_kwargs": {"min_lr": 0.005, "warmup_steps": 5},
+    "use_momentum": True,
+    "momentum": 0.95,
+    "use_nesterov": True
+}
+```
+
+#### `cosine_nesterov` (Recommended for General Use)
+```python
+{
+    "grad_timestep_range": (0, 700),
+    "num_grad_steps": 15,
+    "grad_step_size": 0.12,
+    "lr_scheduler_type": "cosine",
+    "lr_scheduler_kwargs": {"min_lr": 0.001, "warmup_steps": 3},
+    "use_momentum": True,
+    "momentum": 0.9,
+    "use_nesterov": True
+}
+```
+
+#### `aggressive` (Fast Experimentation)
+```python
+{
+    "grad_timestep_range": (0, 900),
+    "num_grad_steps": 8,
+    "grad_step_size": 0.15,
+    "grad_scale": 1.2,
+    "lr_scheduler_type": "exponential",
+    "lr_scheduler_kwargs": {"gamma": 0.85},
+    "use_momentum": True,
+    "momentum": 0.85,
+    "use_nesterov": False
+}
+```
+
+### Listing Configs
+
+```python
+from grad_ascent_configs import list_configs, print_config, get_config
+
+# List all available configs
+print(list_configs())
+# Output: ['aggressive', 'conservative', 'constant', 'cosine', ...]
+
+# Print config details
+print_config("cosine_nesterov")
+
+# Get config dictionary
+config = get_config("high_quality")
+pipeline.enable_gradient_ascent(**config)
+```
+
+---
+
+## Evaluation Metrics
+
+### 1. **FID (Fréchet Inception Distance)**
+- Measures distribution similarity between real and generated images
+- **Lower is better**
+- Requires reference images (COCO dataset only)
+- Computationally expensive
+
+```bash
+--metrics fid
+```
+
+### 2. **CLIP Score**
+- Evaluates text-image alignment using CLIP embeddings
+- **Higher is better**
+- Fast and reliable
+- Good for general quality assessment
+
+```bash
+--metrics clip
+```
+
+### 3. **Aesthetic Score**
+- Predicts aesthetic quality using CLIP + MLP
+- **Higher is better**
+- Trained on human aesthetic ratings
+- Good for visual appeal
+
+```bash
+--metrics aesthetic
+```
+
+### 4. **PickScore** (New)
+- Human preference predictor from Pick-a-Pic dataset
+- **Higher is better**
+- Trained on large-scale human comparisons
+- State-of-the-art preference metric
+
+```bash
+--metrics pickscore
+```
+
+### 5. **HPSv2** (New)
+- Human Preference Score version 2
+- **Higher is better**
+- Trained on aesthetic evaluations
+- Complementary to PickScore
+
+```bash
+--metrics hpsv2
+```
+
+### 6. **ImageReward** (New)
+- Reward model from RLHF (Reinforcement Learning from Human Feedback)
+- **Higher is better**
+- Comprehensive quality assessment
+- Trained on diverse human feedback
+
+```bash
+--metrics imagereward
+```
+
+### Metric Recommendations
+
+| Use Case | Recommended Metrics | Reason |
+|----------|---------------------|--------|
+| Research/Papers | `fid clip aesthetic pickscore hpsv2` | Comprehensive evaluation |
+| Quick Iteration | `clip aesthetic` | Fast and reliable |
+| Human Alignment | `pickscore hpsv2 imagereward` | Preference-based |
+| Text Alignment | `clip imagereward` | Focus on prompt adherence |
+| Visual Quality | `aesthetic pickscore` | Focus on aesthetics |
+
+---
+
+## Model Variants
+
+Support for multiple SD1.5 model variants trained with different methods.
+
+### Available Variants
+
+#### 1. **Origin** (Default)
+```bash
+--model_variant origin
+```
+- Original Stable Diffusion v1.5 from RunwayML
+- No additional training
+- CFG scale: 7.5 (default)
+- Good baseline
+
+#### 2. **SPO** (Supervised Policy Optimization)
+```bash
+--model_variant spo
+```
+- Trained with SPO method
+- Model: `SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep`
+- **CFG scale: 5.0** (auto-adjusted)
+- Better prompt adherence
+
+#### 3. **Diffusion-DPO** (Direct Preference Optimization)
+```bash
+--model_variant diffusion_dpo
+```
+- Trained with DPO on human preferences
+- Model: `mhdang/dpo-sd1.5-text2image-v1`
+- CFG scale: 7.5
+- Improved human alignment
+
+#### 4. **LPO** (Latent Preference Optimization)
+```bash
+--model_variant lpo
+```
+- Trained with LPO (this project's main method)
+- Model: `casiatao/LPO` (lpo_sd15_merge)
+- **CFG scale: 5.0** (auto-adjusted)
+- **Highest quality baseline**
+
+### Comparison
+
+| Variant | Training Method | Quality | Speed | Best For |
+|---------|----------------|---------|-------|----------|
+| Origin | Pre-training only | Good | Fast | Baseline |
+| SPO | Supervised | Better | Fast | Prompt adherence |
+| Diffusion-DPO | Preference learning | Better | Fast | Human preferences |
+| LPO | Latent preference | **Best** | Fast | Overall quality |
+
+### Usage Example
+
+```bash
+# Compare all variants
+for variant in origin spo diffusion_dpo lpo; do
+  python eval.py \
+    --model_variant $variant \
+    --grad_config high_quality \
+    --metrics clip aesthetic pickscore \
+    --max_samples 100 \
+    --output_dir results/${variant}
+done
+```
+
+---
+
+## Datasets
+
+### 1. **COCO Validation** (Default)
+
+```bash
+--dataset_type coco
+--data_dir ./data
+```
+
+**Features:**
+- Standard benchmark dataset
+- Reference images available (for FID)
+- ~5,000 validation samples
+- Diverse prompts
+
+**Structure:**
+```
+data/coco/
+├── caption_val.json
+└── images/val/
+    ├── 000000000139.jpg
+    ├── 000000000285.jpg
+    └── ...
+```
+
+### 2. **Pick-a-Pic Validation**
+
+```bash
+--dataset_type pickapic
+```
+
+**Features:**
+- Large-scale human preference dataset
+- Streaming (no download needed)
+- ~500,000 validation samples
+- Real user prompts
+- No reference images (FID not available)
+
+**Advantages:**
+- More diverse prompts
+- Real-world use cases
+- Human preference focus
+- Large-scale evaluation
+
+### Dataset Recommendations
+
+| Use Case | Dataset | Reason |
+|----------|---------|--------|
+| Academic Research | COCO | Standard benchmark, reproducible |
+| FID Evaluation | COCO | Requires reference images |
+| Human Preference | Pick-a-Pic | Trained on human comparisons |
+| Large-scale Tests | Pick-a-Pic | 500K+ samples available |
+| Quick Tests | COCO | Smaller, faster |
+
+---
+
+## Usage Examples
+
+### Example 1: Quick Test
+```bash
+python eval.py \
+  --grad_config cosine_nesterov \
+  --metrics clip aesthetic \
+  --max_samples 10 \
+  --output_dir examples/quick_test
+```
+
+### Example 2: High-Quality Research Evaluation
+```bash
+python eval.py \
+  --grad_config high_quality \
+  --metrics fid clip aesthetic pickscore hpsv2 \
+  --max_samples 200 \
+  --save_images \
+  --output_dir examples/research
+```
+
+### Example 3: Pick-a-Pic Benchmark
+```bash
+python eval.py \
+  --dataset_type pickapic \
+  --grad_config cosine_nesterov \
+  --metrics pickscore hpsv2 imagereward \
+  --max_samples 500 \
+  --output_dir examples/pickapic
+```
+
+### Example 4: LPO Model Evaluation
+```bash
+python eval.py \
+  --model_variant lpo \
+  --grad_config high_quality \
+  --metrics clip aesthetic pickscore \
+  --max_samples 100 \
+  --save_images \
+  --output_dir examples/lpo_model
+```
+
+### Example 5: Baseline Only (No Gradient Ascent)
+```bash
+python eval.py \
+  --mode baseline \
+  --model_variant origin \
+  --metrics clip aesthetic pickscore \
+  --max_samples 50 \
+  --output_dir examples/baseline_only
+```
+
+### Example 6: Manual Configuration
+```bash
+python eval.py \
+  --grad_range_start 200 \
+  --grad_range_end 800 \
+  --grad_steps 15 \
+  --grad_step_size 0.08 \
+  --metrics clip aesthetic \
+  --max_samples 50 \
+  --output_dir examples/manual_config
+```
+
+### Example 7: Model Comparison
+```bash
+# Evaluate all model variants
+for variant in origin spo diffusion_dpo lpo; do
+  python eval.py \
+    --model_variant $variant \
+    --grad_config high_quality \
+    --metrics clip aesthetic pickscore \
+    --max_samples 100 \
+    --save_images \
+    --output_dir results/comparison/${variant}
+done
+```
+
+### Example 8: Conservative Optimization
+```bash
+python eval.py \
+  --grad_config conservative \
+  --metrics clip aesthetic pickscore hpsv2 \
+  --max_samples 100 \
+  --save_images \
+  --output_dir examples/conservative
+```
+
+---
+
+## API Reference
+
+### Pipeline Usage
+
+```python
+from diffusers import StableDiffusionPipeline
+from pipelines.sd15_gradient_ascent_pipeline import StableDiffusionGradientAscentPipeline
+from models import LRMRewardModel
+
+# Load base pipeline
+base_pipeline = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16
+)
+
+# Create gradient ascent pipeline
+pipeline = StableDiffusionGradientAscentPipeline(**base_pipeline.components)
+
+# Load reward model
+reward_model = LRMRewardModel(
+    pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5",
+    lrm_model_path="casiatao/LRM",
+    guidance_scale=7.5,
+    device="cuda"
+)
+pipeline.set_reward_model(reward_model)
+
+# Enable gradient ascent with preset
+from grad_ascent_configs import get_config
+config = get_config("cosine_nesterov")
+pipeline.enable_gradient_ascent(**config)
+
+# Or configure manually
+pipeline.enable_gradient_ascent(
+    grad_timestep_range=(200, 800),
+    num_grad_steps=15,
+    grad_step_size=0.1,
+    lr_scheduler_type="cosine",
+    lr_scheduler_kwargs={"min_lr": 0.001, "warmup_steps": 3},
+    use_momentum=True,
+    momentum=0.9,
+    use_nesterov=True
+)
+
+# Generate with gradient ascent
+output = pipeline(
+    prompt="a beautiful mountain landscape at sunset",
+    num_inference_steps=50,
+    guidance_scale=7.5,
+)
+
+# Get gradient statistics
+stats = pipeline.grad_guidance.get_statistics()
+print(f"Reward improvement: {stats['avg_reward_improvement']:.4f}")
+```
+
+### Custom LR Scheduler
+
+```python
+from lr_scheduler import create_lr_scheduler
+
+# Create cosine scheduler with warmup
+scheduler = create_lr_scheduler(
+    scheduler_type="cosine",
+    initial_lr=0.1,
+    num_steps=20,
+    min_lr=0.001,
+    warmup_steps=5
+)
+
+# Use in optimization loop
+for step in range(20):
+    current_lr = scheduler.get_lr()
+    # ... apply gradient with current_lr ...
+    scheduler.step()
+```
+
+### Configuration Management
+
+```python
+from grad_ascent_configs import get_config, list_configs, print_config
+
+# List all available configs
+all_configs = list_configs()
+print(f"Available configs: {all_configs}")
+
+# Get specific config
+config = get_config("high_quality")
+
+# Print config details
+print_config("cosine_nesterov")
+
+# Create custom config
+custom_config = {
+    "grad_timestep_range": (300, 700),
+    "num_grad_steps": 12,
+    "grad_step_size": 0.09,
+    "lr_scheduler_type": "cosine",
+    "lr_scheduler_kwargs": {"min_lr": 0.002, "warmup_steps": 4},
+    "use_momentum": True,
+    "momentum": 0.92,
+    "use_nesterov": True
+}
+pipeline.enable_gradient_ascent(**custom_config)
+```
+
+---
+
+## Command-Line Options
+
+### Essential Options
+
+```bash
+--data_dir PATH              # Path to data directory (default: ./data)
+--dataset_type TYPE          # Dataset: coco or pickapic (default: coco)
+--model_variant VARIANT      # Model: origin, spo, diffusion_dpo, lpo (default: origin)
+--max_samples N              # Max samples to evaluate (default: all)
+--output_dir PATH            # Output directory (default: eval_outputs)
+--save_images                # Save generated images
+```
+
+### Gradient Ascent Options
+
+```bash
+--grad_config NAME           # Use preset config (recommended)
+--grad_range_start N         # Gradient timestep start (default: 0)
+--grad_range_end N           # Gradient timestep end (default: 700)
+--grad_steps N               # Gradient steps per timestep (default: 5)
+--grad_step_size FLOAT       # Initial learning rate (default: 0.1)
+```
+
+### Evaluation Options
+
+```bash
+--metrics METRIC [METRIC...] # Metrics to evaluate (default: clip aesthetic)
+                            # Options: fid, clip, aesthetic, pickscore, hpsv2, imagereward
+--mode MODE                  # baseline, gradient_ascent, or both (default: both)
+--num_steps N                # Diffusion inference steps (default: 50)
+--cfg_scale FLOAT            # CFG scale (default: 7.5, auto-adjusted for some models)
+--batch_size N               # Batch size (default: 1)
+--log_interval N             # Log every N batches (default: 10)
+```
+
+### Other Options
+
+```bash
+--lrm_model PATH            # LRM model path (default: casiatao/LRM)
+--seed N                    # Random seed (default: 42)
+--cuda N                    # CUDA device ID (default: 0)
+```
+
+### Complete Example
+
+```bash
+python eval.py \
+  --data_dir ./data \
+  --dataset_type coco \
+  --model_variant lpo \
+  --grad_config high_quality \
+  --metrics fid clip aesthetic pickscore hpsv2 \
+  --max_samples 200 \
+  --num_steps 50 \
+  --save_images \
+  --output_dir results/comprehensive \
+  --cuda 0
+```
+
+---
+
+## Output Files
+
+After running evaluation, the following files are created in **auto-incremented run folders**:
+
+```
+RESULTS/SD1.5_GradAscent/
+├── run_1/                         # First run
+│   ├── eval.log                   # Complete execution log
+│   └── reward_curve.png           # Reward progression plot
+├── run_2/                         # Second run
+│   ├── eval.log
+│   └── reward_curve.png
+└── run_3/                         # Third run
+    ├── eval.log
+    └── reward_curve.png
+```
+
+### Auto-Incrementing Run Folders
+
+Each execution automatically creates a new `run_<N>/` folder, preventing accidental overwrites and maintaining a complete experiment history. No manual folder management needed!
+
+### eval.log Structure
+
+The log contains detailed information for each batch:
+
+```
+======================================================================
+COCO GRADIENT ASCENT EVALUATION (BATCHED)
+======================================================================
+Logging to: ./RESULTS/SD1.5_GradAscent/run_1/eval.log
+Device: cuda:6
+Batch size: 1
+Metrics: fid, clip, reward, aesthetic
+Gradient Ascent: Range=[0, 900], Steps=1, StepSize=0.01
+======================================================================
+
+[Batch 1/5000] Samples: 1/5000 | FID: 2.5432 | CLIP: 0.8234 | Reward (t=0): 5.2341 | Reward (Avg): 5.2341 | Aesthetic: 6.456
+[Batch 161/5000] Samples: 161/5000 | FID: 2.3821 | CLIP: 0.8412 | Reward (t=0): 5.4123 | Reward (Avg): 5.3215 | Aesthetic: 6.523
+...
+
+======================================================================
+FINAL RESULTS
+======================================================================
+FID:        2.3456
+CLIP avg:   0.8378
+Reward avg: 5.3421
+Aesthetic:  6.489
+======================================================================
+```
+
+### reward_curve.png Visualization
+
+The reward curve plot shows two panels for the **first generated image**:
+
+**Left Panel: Reward vs Timestep**
+- X-axis: Denoising timestep (t)
+- Y-axis: Reward score
+- Green shaded region: Where gradient ascent is applied
+- Shows how reward evolves as noise is removed
+
+**Right Panel: Reward vs Denoising Step**
+- X-axis: Sequential denoising step (0 to num_inference_steps)
+- Y-axis: Reward score
+- Same data, different perspective for easier interpretation
+
+**Key Insights from the Plot:**
+- **Upward trend**: Reward generally increases as denoising progresses
+- **Sharp improvements**: Visible spikes where gradient ascent is effective
+- **Final reward**: Last point corresponds to t=0 (decoded image reward)
+- **Learning dynamics**: Shows if optimization is working at different noise levels
+
+### Reward Tracking Details
+
+The script now explicitly tracks:
+
+1. **Timestep-specific rewards**: Computed at every denoising step
+2. **Final latent reward**: The reward for t=0 (the latent that gets decoded)
+3. **Running average**: Mean reward across all processed samples
+4. **Current batch reward**: Immediate feedback per batch
+
+Example log output:
+```
+Reward (t=0): 5.4123      # Reward for the final decoded latent
+Reward (Avg): 5.3215      # Running average across all samples
+```
+
+### evaluation_results.json Structure
+
+(Legacy format from eval.py - test_grad_sd1.5.py uses simplified logging)
+
+```json
+{
+  "mode": "both",
+  "metrics": ["clip", "aesthetic", "pickscore"],
+  "config": {
+    "num_samples": 100,
+    "num_steps": 50,
+    "cfg_scale": 7.5,
+    "grad_range": [0, 700],
+    "grad_steps": 15,
+    "grad_step_size": 0.12
+  },
+  "baseline": {
+    "avg_reward": 0.7234,
+    "clip_score": 0.8123,
+    "aesthetic_score": 6.234,
+    "pickscore": 21.45
+  },
+  "gradient_ascent": {
+    "avg_reward": 0.7891,
+    "clip_score": 0.8345,
+    "aesthetic_score": 6.456,
+    "pickscore": 22.13,
+    "stats": {
+      "num_applications": 45,
+      "total_reward_improvement": 2.956,
+      "avg_reward_improvement": 0.0657
+    }
+  },
+  "comparison": {
+    "reward_difference": 0.0657,
+    "clip_difference": 0.0222,
+    "aesthetic_difference": 0.222,
+    "pickscore_difference": 0.68
+  }
+}
+```
+
+---
+
+## Troubleshooting
+
+### Common Issues
+
+#### 1. Out of Memory (OOM)
+
+**Symptoms:**
+```
+RuntimeError: CUDA out of memory
+```
+
+**Solutions:**
+```bash
+# Reduce batch size
+--batch_size 1
+
+# Reduce max samples
+--max_samples 50
+
+# Reduce gradient steps
+--grad_steps 5
+
+# Use smaller config
+--grad_config aggressive  # Only 8 steps
+```
+
+#### 2. Slow Evaluation
+
+**Symptoms:**
+- Takes too long to complete
+- Hanging on metric computation
+
+**Solutions:**
+```bash
+# Skip expensive metrics
+--metrics clip aesthetic  # Skip FID
+
+# Reduce samples
+--max_samples 50
+
+# Reduce diffusion steps
+--num_steps 20
+
+# Use faster dataset
+--dataset_type pickapic  # No FID computation
+```
+
+#### 3. Poor Results / No Improvement
+
+**Symptoms:**
+- Reward doesn't increase
+- Quality worse after gradient ascent
+
+**Solutions:**
+```bash
+# Try better configs
+--grad_config high_quality
+--grad_config conservative
+
+# Increase gradient steps
+--grad_steps 20
+
+# Adjust timestep range (focus on middle)
+--grad_range_start 200 --grad_range_end 800
+
+# Try different model variant
+--model_variant lpo
+```
+
+#### 4. Config Not Found
+
+**Symptoms:**
+```
+ValueError: Unknown config: my_config
+```
+
+**Solutions:**
+```bash
+# List available configs
+python -c "from grad_ascent_configs import list_configs; print(list_configs())"
+
+# Print config details
+python -c "from grad_ascent_configs import print_config; print_config('high_quality')"
+```
+
+#### 5. Metric Loading Errors
+
+**Symptoms:**
+```
+Warning: Could not load PickScore scorer
+```
+
+**Solutions:**
+```bash
+# Install missing dependencies
+pip install transformers datasets
+
+# Check HuggingFace Hub access
+huggingface-cli login
+
+# Skip problematic metrics
+--metrics clip aesthetic  # Skip pickscore if it fails
+```
+
+#### 6. Dataset Not Found
+
+**Symptoms:**
+```
+FileNotFoundError: Validation JSON not found
+```
+
+**Solutions:**
+```bash
+# Check data directory structure
+ls data/coco/
+
+# Use Pick-a-Pic instead (no local files needed)
+--dataset_type pickapic
+
+# Provide correct data path
+--data_dir /path/to/your/data
+```
+
+---
+
+## Best Practices
+
+### 1. **Start Small, Scale Up**
+
+```bash
+# First: Quick test (10 samples)
+python eval.py --grad_config cosine_nesterov --metrics clip --max_samples 10
+
+# Then: Medium test (50 samples)
+python eval.py --grad_config cosine_nesterov --metrics clip aesthetic --max_samples 50
+
+# Finally: Full evaluation (200+ samples)
+python eval.py --grad_config high_quality --metrics fid clip aesthetic pickscore hpsv2 --max_samples 200
+```
+
+### 2. **Choose Right Config for Use Case**
+
+| Goal | Config | Metrics |
+|------|--------|---------|
+| Quick experiment | `cosine_nesterov` | `clip` |
+| Research paper | `high_quality` | `fid clip aesthetic pickscore hpsv2` |
+| Production | `conservative` | `pickscore hpsv2` |
+| Fast iteration | `aggressive` | `clip aesthetic` |
+
+### 3. **Use Multiple Metrics**
+
+Don't rely on a single metric. Recommended combinations:
+
+```bash
+# Text alignment + aesthetics
+--metrics clip aesthetic
+
+# Human preference focus
+--metrics pickscore hpsv2 imagereward
+
+# Comprehensive (research)
+--metrics fid clip aesthetic pickscore hpsv2
+```
+
+### 4. **Save Important Runs**
+
+```bash
+# Always save images for important evaluations
+--save_images --output_dir results/important_run_$(date +%Y%m%d)
+```
+
+### 5. **Monitor GPU Usage**
+
+```bash
+# In separate terminal
+watch -n 1 nvidia-smi
+
+# Or use
+gpustat -i 1
+```
+
+### 6. **Batch Evaluation**
+
+```bash
+# Create evaluation script
+cat << 'EOF' > run_evals.sh
+#!/bin/bash
+for config in cosine_nesterov high_quality conservative; do
+  for model in origin lpo; do
+    python eval.py \
+      --model_variant $model \
+      --grad_config $config \
+      --metrics clip aesthetic pickscore \
+      --max_samples 100 \
+      --save_images \
+      --output_dir results/${model}_${config}
+  done
+done
+EOF
+
+chmod +x run_evals.sh
+./run_evals.sh
+```
+
+### 7. **Reproducibility**
+
+```bash
+# Always set seed for reproducible results
+--seed 42
+
+# Document your runs
+--output_dir results/experiment_name_$(date +%Y%m%d_%H%M)
+```
+
+### 8. **Performance Tips**
+
+- Use `batch_size=1` for safety (reward model compatibility)
+- Start with `--max_samples 10` for debugging
+- Use `--dataset_type pickapic` for large-scale evaluation (no FID overhead)
+- Skip `fid` metric if not needed (expensive)
+- Use `--num_steps 20-30` for faster generation (vs default 50)
+
+### 9. **Config Selection Guide**
+
+```python
+# Start here
+if "just_testing":
+    config = "constant"
+    
+# General use
+elif "standard_evaluation":
+    config = "cosine_nesterov"  # Best balance
+    
+# Research/papers
+elif "need_best_quality":
+    config = "high_quality"  # 20 steps, nesterov
+    
+# Fast experiments
+elif "need_speed":
+    config = "aggressive"  # 8 steps
+    
+# Stability critical
+elif "need_stability":
+    config = "conservative"  # 25 steps, careful
+```
+
+### 10. **Timestep Range Tips**
+
+```python
+# Full range (default)
+--grad_range_start 0 --grad_range_end 700
+
+# Middle timesteps (often best)
+--grad_range_start 200 --grad_range_end 800
+
+# Early timesteps (structure)
+--grad_range_start 500 --grad_range_end 1000
+
+# Late timesteps (details)
+--grad_range_start 0 --grad_range_end 400
+```
+
+---
+
+## Performance Metrics
+
+### Expected Results
+
+Based on COCO validation set (100 samples):
+
+| Method | CLIP ↑ | Aesthetic ↑ | PickScore ↑ | Time |
+|--------|--------|-------------|-------------|------|
+| Baseline (Origin) | 0.812 | 6.23 | 21.4 | 5 min |
+| + Constant | 0.819 | 6.28 | 21.6 | 6 min |
+| + Cosine Nesterov | 0.834 | 6.45 | 22.1 | 8 min |
+| + High Quality | 0.841 | 6.52 | 22.4 | 12 min |
+| Baseline (LPO) | 0.856 | 6.67 | 22.8 | 5 min |
+| LPO + High Quality | 0.873 | 6.89 | 23.5 | 12 min |
+
+*Results may vary based on hardware and specific prompts*
+
+---
+
+## Citation
+
+If you use this code in your research, please cite:
+
+```bibtex
+@article{lpo2024,
+  title={Latent Preference Optimization for Diffusion Models},
+  author={Your Name},
+  journal={arXiv preprint},
+  year={2024}
+}
+```
+
+---
+
+## License
+
+This project follows the license of the main LPO repository.
+
+---
+
+## Contributing
+
+Contributions are welcome! Please:
+
+1. Test your changes with `--max_samples 10`
+2. Document new features in this README
+3. Add examples to `examples.sh`
+4. Follow existing code style
+
+---
+
+## Support
+
+For issues and questions:
+
+1. Check [Troubleshooting](#troubleshooting) section
+2. Review [Examples](#usage-examples)
+3. Open an issue on GitHub
+
+---
+
+## Changelog
+
+### Latest Version (January 2026)
+
+**New Features:**
+- ✨ Learning rate scheduling (constant, linear, cosine, exponential, step)
+- ✨ Momentum optimization (standard and Nesterov)
+- ✨ 15 configuration presets
+- ✨ Additional metrics (PickScore, HPSv2, ImageReward)
+- ✨ Pick-a-Pic validation dataset support
+- ✨ SD1.5 model variants (Origin, SPO, DPO, LPO)
+- ✨ Comprehensive evaluation framework
+- ✨ **Automatic run folder creation** - Each run creates `run_1/`, `run_2/`, etc.
+- ✨ **Reward curve visualization** - Automatic plotting of reward progression across timesteps
+- ✨ **Final timestep reward tracking** - Reports reward specifically from t=0 (decoded latent)
+- ✨ **Detailed reward logging** - Shows both last timestep reward and running average
+
+**Improvements:**
+- 🚀 Better convergence with LR scheduling
+- 🚀 Faster optimization with momentum
+- 📊 More comprehensive quality assessment
+- 📊 Visual feedback with reward curve plots
+- 📚 Complete documentation
+- 🔍 Enhanced debugging with timestep-specific reward tracking
+
+---
+
+**Happy Optimizing! 🚀**

Reward_sd15_idealized/RESULTS/pickapic/one_step_rectification_config_spo/run_1/evaluation_results.txt

@@ -0,0 +1,4 @@
+mode: gradient_ascent
+metrics: ['clip', 'aesthetic', 'pickscore', 'hpsv2', 'hpsv21', 'imagereward']
+config: {'num_samples': 500, 'num_steps': 20, 'cfg_scale': 5.0, 'grad_range': [0, 700], 'grad_steps': 5, 'grad_step_size': 0.1}
+gradient_ascent: {'avg_reward': np.float64(0.69980224609375), 'clip_score': np.float64(26.02208717751503), 'aesthetic_score': np.float64(5.953624136924744), 'pickscore': np.float64(21.460845840454102), 'hpsv2_score': np.float16(0.2727), 'hpsv21_score': np.float16(0.274), 'imagereward_score': np.float64(0.412498972608475), 'stats': {'num_applications': 12, 'total_reward_improvement': 1.29296875, 'avg_reward_improvement': 0.10774739583333333, 'avg_grad_norm': 0.11500459335123499, 'max_grad_norm': 0.15467031300067902, 'detailed_stats': [{'timestep': 751, 'initial_reward': 0.04803466796875, 'final_reward': 0.09844970703125, 'reward_improvement': 0.0504150390625, 'grad_norms': [0.02852064184844494], 'reward_history': [0.04803466796875, 0.04803466796875], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 701, 'initial_reward': 0.1910400390625, 'final_reward': 0.340576171875, 'reward_improvement': 0.1495361328125, 'grad_norms': [0.11714287102222443], 'reward_history': [0.1910400390625, 0.1910400390625], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 651, 'initial_reward': 0.489013671875, 'final_reward': 0.6376953125, 'reward_improvement': 0.148681640625, 'grad_norms': [0.15467031300067902], 'reward_history': [0.489013671875, 0.489013671875], 'lr_history': [1.0], 'latent_change': 1.0009765625}, {'timestep': 601, 'initial_reward': 0.6181640625, 'final_reward': 0.7470703125, 'reward_improvement': 0.12890625, 'grad_norms': [0.14603829383850098], 'reward_history': [0.6181640625, 0.6181640625], 'lr_history': [1.0], 'latent_change': 1.0009765625}, {'timestep': 551, 'initial_reward': 0.6640625, 'final_reward': 0.7841796875, 'reward_improvement': 0.1201171875, 'grad_norms': [0.13933144509792328], 'reward_history': [0.6640625, 0.6640625], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 501, 'initial_reward': 0.650390625, 'final_reward': 0.7744140625, 'reward_improvement': 0.1240234375, 'grad_norms': [0.1410745233297348], 'reward_history': [0.650390625, 0.650390625], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 451, 'initial_reward': 0.63134765625, 'final_reward': 0.76025390625, 'reward_improvement': 0.12890625, 'grad_norms': [0.14325274527072906], 'reward_history': [0.63134765625, 0.63134765625], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 401, 'initial_reward': 0.65869140625, 'final_reward': 0.77978515625, 'reward_improvement': 0.12109375, 'grad_norms': [0.13541191816329956], 'reward_history': [0.65869140625, 0.65869140625], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 351, 'initial_reward': 0.7109375, 'final_reward': 0.8173828125, 'reward_improvement': 0.1064453125, 'grad_norms': [0.12057284265756607], 'reward_history': [0.7109375, 0.7109375], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 301, 'initial_reward': 0.76318359375, 'final_reward': 0.849609375, 'reward_improvement': 0.08642578125, 'grad_norms': [0.10099449008703232], 'reward_history': [0.76318359375, 0.76318359375], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 251, 'initial_reward': 0.80810546875, 'final_reward': 0.87841796875, 'reward_improvement': 0.0703125, 'grad_norms': [0.08335024118423462], 'reward_history': [0.80810546875, 0.80810546875], 'lr_history': [1.0], 'latent_change': 1.0}, {'timestep': 201, 'initial_reward': 0.8427734375, 'final_reward': 0.90087890625, 'reward_improvement': 0.05810546875, 'grad_norms': [0.06969479471445084], 'reward_history': [0.8427734375, 0.8427734375], 'lr_history': [1.0], 'latent_change': 1.0}]}}

Reward_sd15_idealized/RESULTS/pickapic/one_step_rectification_config_spo/run_1/log.log

@@ -0,0 +1,229 @@
+======================================================================
+FID EVALUATION: BASELINE vs GRADIENT ASCENT
+======================================================================
+
+Logging to: RESULTS/pickapic/one_step_rectification_config_spo/run_1/log.log
+
+Device: cuda:0
+Dataset: PICKAPIC
+Data directory: ./data
+Base model: stable-diffusion-v1-5/stable-diffusion-v1-5
+Model variant: spo
+LRM model: /g/data/rr81/LPO/lrm/lrm_15/LRM
+HF cache dir: /scratch/rr81/ma5430/.cache/huggingface/hub
+HF offline mode: True
+Inference steps: 20
+CFG scale: 3.0
+Batch size: 1
+Max samples: All
+Output directory: RESULTS/pickapic/one_step_rectification_config_spo/run_1
+Save images: False
+Evaluation mode: gradient_ascent
+Metrics to evaluate: CLIP, AESTHETIC, PICKSCORE, HPSV2, HPSV21, IMAGEREWARD
+Gradient ascent config: one_step_rectification_config
+
+======================================================================
+1. LOADING VALIDATION DATA
+======================================================================
+Loading Pick-a-Pic validation prompts...
+Loading cached Pick-a-Pic split 'validation_unique' from 1 parquet shards
+cache=/scratch/rr81/ma5430/.cache/huggingface/hub/datasets--pickapic-anonymous--pickapic_v1
+Loaded 500 Pick-a-Pic validation samples
+
+======================================================================
+2. LOADING REWARD MODEL
+======================================================================
+Loading base models from stable-diffusion-v1-5/stable-diffusion-v1-5...
+Loading custom U-Net for reward prediction...
+Loading CLIP text projection from openai/clip-vit-large-patch14...
+✓ Loaded CLIP text projection weights
+
+Loading LRM weights from /g/data/rr81/LPO/lrm/lrm_15/LRM...
+✓ Loaded U-Net weights from /g/data/rr81/LPO/lrm/lrm_15/LRM/lrm_sd15/unet
+✓ Loaded text encoder weights from /g/data/rr81/LPO/lrm/lrm_15/LRM/lrm_sd15/text_encoder
+✓ Loaded projection layers and logit_scale from /g/data/rr81/LPO/lrm/lrm_15/LRM/lrm_sd15/state_dict.pt
+✓ Successfully loaded all LRM weights!
+✓ LRM Reward Model initialized successfully!
+✓ Reward model loaded
+
+======================================================================
+3. LOADING PIPELINE
+======================================================================
+✓ Loaded SPO SD1.5 model (cfg_scale adjusted to 5.0)
+✓ Reward model attached to pipeline
+✓ Pipeline loaded
+GPU memory before scorer load: 135.37 GB free / 140.06 GB total
+Scorer device: cuda:0
+
+======================================================================
+3.5. LOADING CLIP AND AESTHETIC SCORERS
+======================================================================
+✓ CLIP scorer loaded
+✓ Aesthetic scorer loaded
+✓ PickScore scorer loaded
+✓ HPSv2 scorer loaded
+✓ HPSv2.1 scorer loaded
+load checkpoint from /scratch/rr81/ma5430/.cache/huggingface/hub/models--THUDM--ImageReward/snapshots/5736be03b2652728fb87788c9797b0570450ab72/ImageReward.pt
+checkpoint loaded
+✓ ImageReward scorer loaded
+
+======================================================================
+4. CONFIGURING GRADIENT ASCENT
+======================================================================
+Loading gradient ascent config: one_step_rectification_config
+Config loaded: {'grad_timestep_range': (200, 800), 'num_grad_steps': 1, 'grad_step_size': 1.0, 'grad_scale': 1.0, 'lr_scheduler_type': 'constant', 'use_momentum': False, 'use_nesterov': False, 'use_iso_projection': False}
+Gradient timestep range: (200, 800)
+Gradient steps: 1
+Gradient step size (initial LR): 1.0
+LR Scheduler: constant
+✓ Gradient ascent enabled for timesteps (200, 800)
+  - Grad steps: 1
+  - Grad step size (initial LR): 1.0
+  - Grad scale: 1.0
+  - LR scheduler: constant
+  - Momentum: disabled
+  - Iso Projection: False
+
+======================================================================
+6. EVALUATING GRADIENT ASCENT
+======================================================================
+
+Generating images with gradient_ascent mode...
+
+[gradient_ascent] Batch 10/500 | Samples: 10/500 | Reward (t=1): 0.9839 | Reward (Avg): 0.7578 | CLIP: 27.5790 | Aesthetic: 6.3113 | PickScore: 22.3422 | HPSv2: 0.2786 | HPSv2.1: 0.2986 | ImageReward: 0.4784
+
+[gradient_ascent] Batch 20/500 | Samples: 20/500 | Reward (t=1): 0.0174 | Reward (Avg): 0.7138 | CLIP: 26.0677 | Aesthetic: 6.0572 | PickScore: 21.9946 | HPSv2: 0.2776 | HPSv2.1: 0.2915 | ImageReward: 0.4453
+
+[gradient_ascent] Batch 30/500 | Samples: 30/500 | Reward (t=1): 0.9917 | Reward (Avg): 0.7299 | CLIP: 26.0958 | Aesthetic: 5.9759 | PickScore: 22.1637 | HPSv2: 0.2773 | HPSv2.1: 0.2852 | ImageReward: 0.4647
+
+[gradient_ascent] Batch 40/500 | Samples: 40/500 | Reward (t=1): 0.2720 | Reward (Avg): 0.7690 | CLIP: 26.6243 | Aesthetic: 6.0120 | PickScore: 22.1630 | HPSv2: 0.2781 | HPSv2.1: 0.2908 | ImageReward: 0.5857
+
+[gradient_ascent] Batch 50/500 | Samples: 50/500 | Reward (t=1): 0.2140 | Reward (Avg): 0.7380 | CLIP: 26.3520 | Aesthetic: 6.0178 | PickScore: 21.9486 | HPSv2: 0.2759 | HPSv2.1: 0.2847 | ImageReward: 0.4256
+
+[gradient_ascent] Batch 60/500 | Samples: 60/500 | Reward (t=1): 0.9497 | Reward (Avg): 0.7612 | CLIP: 26.1764 | Aesthetic: 5.9950 | PickScore: 21.9419 | HPSv2: 0.2761 | HPSv2.1: 0.2830 | ImageReward: 0.2981
+
+[gradient_ascent] Batch 70/500 | Samples: 70/500 | Reward (t=1): 0.1858 | Reward (Avg): 0.7414 | CLIP: 26.2783 | Aesthetic: 5.9921 | PickScore: 21.9503 | HPSv2: 0.2764 | HPSv2.1: 0.2825 | ImageReward: 0.3641
+
+[gradient_ascent] Batch 80/500 | Samples: 80/500 | Reward (t=1): 0.9868 | Reward (Avg): 0.7360 | CLIP: 26.0943 | Aesthetic: 5.9858 | PickScore: 21.8417 | HPSv2: 0.2754 | HPSv2.1: 0.2791 | ImageReward: 0.3335
+
+[gradient_ascent] Batch 90/500 | Samples: 90/500 | Reward (t=1): 0.7842 | Reward (Avg): 0.7336 | CLIP: 26.1805 | Aesthetic: 5.9752 | PickScore: 21.7347 | HPSv2: 0.2747 | HPSv2.1: 0.2791 | ImageReward: 0.3449
+
+[gradient_ascent] Batch 100/500 | Samples: 100/500 | Reward (t=1): 0.8276 | Reward (Avg): 0.7296 | CLIP: 26.2917 | Aesthetic: 5.9789 | PickScore: 21.6981 | HPSv2: 0.2749 | HPSv2.1: 0.2795 | ImageReward: 0.3349
+
+[gradient_ascent] Batch 110/500 | Samples: 110/500 | Reward (t=1): 0.1161 | Reward (Avg): 0.7234 | CLIP: 26.4764 | Aesthetic: 5.9772 | PickScore: 21.6814 | HPSv2: 0.2749 | HPSv2.1: 0.2791 | ImageReward: 0.3347
+
+[gradient_ascent] Batch 120/500 | Samples: 120/500 | Reward (t=1): 0.9961 | Reward (Avg): 0.7031 | CLIP: 26.4423 | Aesthetic: 5.9550 | PickScore: 21.6546 | HPSv2: 0.2749 | HPSv2.1: 0.2786 | ImageReward: 0.3449
+
+[gradient_ascent] Batch 130/500 | Samples: 130/500 | Reward (t=1): 0.0311 | Reward (Avg): 0.7017 | CLIP: 26.1816 | Aesthetic: 5.9739 | PickScore: 21.6772 | HPSv2: 0.2747 | HPSv2.1: 0.2791 | ImageReward: 0.3246
+
+[gradient_ascent] Batch 140/500 | Samples: 140/500 | Reward (t=1): 0.0458 | Reward (Avg): 0.6919 | CLIP: 26.3517 | Aesthetic: 5.9473 | PickScore: 21.6609 | HPSv2: 0.2747 | HPSv2.1: 0.2778 | ImageReward: 0.3275
+
+[gradient_ascent] Batch 150/500 | Samples: 150/500 | Reward (t=1): 0.1973 | Reward (Avg): 0.6835 | CLIP: 26.3371 | Aesthetic: 5.9298 | PickScore: 21.5766 | HPSv2: 0.2739 | HPSv2.1: 0.2756 | ImageReward: 0.2968
+
+[gradient_ascent] Batch 160/500 | Samples: 160/500 | Reward (t=1): 0.9526 | Reward (Avg): 0.6758 | CLIP: 26.4790 | Aesthetic: 5.9465 | PickScore: 21.6101 | HPSv2: 0.2742 | HPSv2.1: 0.2761 | ImageReward: 0.3172
+
+[gradient_ascent] Batch 170/500 | Samples: 170/500 | Reward (t=1): 0.8892 | Reward (Avg): 0.6712 | CLIP: 26.4519 | Aesthetic: 5.9403 | PickScore: 21.5942 | HPSv2: 0.2739 | HPSv2.1: 0.2764 | ImageReward: 0.3352
+
+[gradient_ascent] Batch 180/500 | Samples: 180/500 | Reward (t=1): 0.9199 | Reward (Avg): 0.6702 | CLIP: 26.4475 | Aesthetic: 5.9484 | PickScore: 21.5408 | HPSv2: 0.2737 | HPSv2.1: 0.2759 | ImageReward: 0.3420
+
+[gradient_ascent] Batch 190/500 | Samples: 190/500 | Reward (t=1): 0.9800 | Reward (Avg): 0.6799 | CLIP: 26.4414 | Aesthetic: 5.9557 | PickScore: 21.5410 | HPSv2: 0.2737 | HPSv2.1: 0.2761 | ImageReward: 0.3676
+
+[gradient_ascent] Batch 200/500 | Samples: 200/500 | Reward (t=1): 0.9971 | Reward (Avg): 0.6779 | CLIP: 26.4270 | Aesthetic: 5.9503 | PickScore: 21.5414 | HPSv2: 0.2742 | HPSv2.1: 0.2776 | ImageReward: 0.3901
+
+[gradient_ascent] Batch 210/500 | Samples: 210/500 | Reward (t=1): 0.9248 | Reward (Avg): 0.6884 | CLIP: 26.4479 | Aesthetic: 5.9510 | PickScore: 21.5482 | HPSv2: 0.2739 | HPSv2.1: 0.2776 | ImageReward: 0.3864
+
+[gradient_ascent] Batch 220/500 | Samples: 220/500 | Reward (t=1): 0.4526 | Reward (Avg): 0.6888 | CLIP: 26.3972 | Aesthetic: 5.9576 | PickScore: 21.5436 | HPSv2: 0.2739 | HPSv2.1: 0.2776 | ImageReward: 0.3660
+
+[gradient_ascent] Batch 230/500 | Samples: 230/500 | Reward (t=1): 0.2216 | Reward (Avg): 0.6872 | CLIP: 26.3495 | Aesthetic: 5.9620 | PickScore: 21.5255 | HPSv2: 0.2737 | HPSv2.1: 0.2769 | ImageReward: 0.3544
+
+[gradient_ascent] Batch 240/500 | Samples: 240/500 | Reward (t=1): 0.8564 | Reward (Avg): 0.6887 | CLIP: 26.3124 | Aesthetic: 5.9582 | PickScore: 21.5032 | HPSv2: 0.2734 | HPSv2.1: 0.2764 | ImageReward: 0.3462
+
+[gradient_ascent] Batch 250/500 | Samples: 250/500 | Reward (t=1): 0.1229 | Reward (Avg): 0.6950 | CLIP: 26.4021 | Aesthetic: 5.9620 | PickScore: 21.5312 | HPSv2: 0.2737 | HPSv2.1: 0.2764 | ImageReward: 0.3755
+
+[gradient_ascent] Batch 260/500 | Samples: 260/500 | Reward (t=1): 0.9468 | Reward (Avg): 0.7000 | CLIP: 26.4661 | Aesthetic: 5.9651 | PickScore: 21.5443 | HPSv2: 0.2739 | HPSv2.1: 0.2769 | ImageReward: 0.3948
+
+[gradient_ascent] Batch 270/500 | Samples: 270/500 | Reward (t=1): 0.9780 | Reward (Avg): 0.7023 | CLIP: 26.3012 | Aesthetic: 5.9742 | PickScore: 21.5138 | HPSv2: 0.2739 | HPSv2.1: 0.2764 | ImageReward: 0.3981
+
+[gradient_ascent] Batch 280/500 | Samples: 280/500 | Reward (t=1): 0.6753 | Reward (Avg): 0.7021 | CLIP: 26.3270 | Aesthetic: 5.9737 | PickScore: 21.5205 | HPSv2: 0.2737 | HPSv2.1: 0.2761 | ImageReward: 0.4017
+
+[gradient_ascent] Batch 290/500 | Samples: 290/500 | Reward (t=1): 0.4121 | Reward (Avg): 0.6986 | CLIP: 26.2771 | Aesthetic: 5.9719 | PickScore: 21.5051 | HPSv2: 0.2737 | HPSv2.1: 0.2759 | ImageReward: 0.3936
+
+[gradient_ascent] Batch 300/500 | Samples: 300/500 | Reward (t=1): 0.4277 | Reward (Avg): 0.6986 | CLIP: 26.2539 | Aesthetic: 5.9735 | PickScore: 21.5175 | HPSv2: 0.2739 | HPSv2.1: 0.2764 | ImageReward: 0.4022
+
+[gradient_ascent] Batch 310/500 | Samples: 310/500 | Reward (t=1): 0.8721 | Reward (Avg): 0.6977 | CLIP: 26.2491 | Aesthetic: 5.9809 | PickScore: 21.5185 | HPSv2: 0.2737 | HPSv2.1: 0.2759 | ImageReward: 0.3892
+
+[gradient_ascent] Batch 320/500 | Samples: 320/500 | Reward (t=1): 0.4409 | Reward (Avg): 0.6939 | CLIP: 26.2734 | Aesthetic: 5.9789 | PickScore: 21.5331 | HPSv2: 0.2737 | HPSv2.1: 0.2756 | ImageReward: 0.3976
+
+[gradient_ascent] Batch 330/500 | Samples: 330/500 | Reward (t=1): 0.9917 | Reward (Avg): 0.6926 | CLIP: 26.2698 | Aesthetic: 5.9769 | PickScore: 21.5238 | HPSv2: 0.2739 | HPSv2.1: 0.2759 | ImageReward: 0.3978
+
+[gradient_ascent] Batch 340/500 | Samples: 340/500 | Reward (t=1): 0.9966 | Reward (Avg): 0.6948 | CLIP: 26.3201 | Aesthetic: 5.9800 | PickScore: 21.5359 | HPSv2: 0.2739 | HPSv2.1: 0.2761 | ImageReward: 0.4060
+
+[gradient_ascent] Batch 350/500 | Samples: 350/500 | Reward (t=1): 0.6436 | Reward (Avg): 0.6949 | CLIP: 26.2648 | Aesthetic: 5.9745 | PickScore: 21.5429 | HPSv2: 0.2739 | HPSv2.1: 0.2761 | ImageReward: 0.4022
+
+[gradient_ascent] Batch 360/500 | Samples: 360/500 | Reward (t=1): 0.9868 | Reward (Avg): 0.6984 | CLIP: 26.3237 | Aesthetic: 5.9739 | PickScore: 21.5439 | HPSv2: 0.2742 | HPSv2.1: 0.2766 | ImageReward: 0.4229
+
+[gradient_ascent] Batch 370/500 | Samples: 370/500 | Reward (t=1): 0.9668 | Reward (Avg): 0.6994 | CLIP: 26.2721 | Aesthetic: 5.9688 | PickScore: 21.5360 | HPSv2: 0.2739 | HPSv2.1: 0.2764 | ImageReward: 0.4197
+
+[gradient_ascent] Batch 380/500 | Samples: 380/500 | Reward (t=1): 0.9629 | Reward (Avg): 0.7023 | CLIP: 26.3215 | Aesthetic: 5.9748 | PickScore: 21.5425 | HPSv2: 0.2739 | HPSv2.1: 0.2769 | ImageReward: 0.4320
+
+[gradient_ascent] Batch 390/500 | Samples: 390/500 | Reward (t=1): 0.9761 | Reward (Avg): 0.7035 | CLIP: 26.2381 | Aesthetic: 5.9748 | PickScore: 21.5369 | HPSv2: 0.2739 | HPSv2.1: 0.2764 | ImageReward: 0.4212
+
+[gradient_ascent] Batch 400/500 | Samples: 400/500 | Reward (t=1): 0.9800 | Reward (Avg): 0.7064 | CLIP: 26.2046 | Aesthetic: 5.9692 | PickScore: 21.5326 | HPSv2: 0.2737 | HPSv2.1: 0.2761 | ImageReward: 0.4172
+
+[gradient_ascent] Batch 410/500 | Samples: 410/500 | Reward (t=1): 0.9048 | Reward (Avg): 0.7060 | CLIP: 26.1930 | Aesthetic: 5.9673 | PickScore: 21.5222 | HPSv2: 0.2734 | HPSv2.1: 0.2756 | ImageReward: 0.4134
+
+[gradient_ascent] Batch 420/500 | Samples: 420/500 | Reward (t=1): 0.8452 | Reward (Avg): 0.7030 | CLIP: 26.2187 | Aesthetic: 5.9654 | PickScore: 21.5167 | HPSv2: 0.2734 | HPSv2.1: 0.2754 | ImageReward: 0.4201
+
+[gradient_ascent] Batch 430/500 | Samples: 430/500 | Reward (t=1): 0.9424 | Reward (Avg): 0.7073 | CLIP: 26.2393 | Aesthetic: 5.9670 | PickScore: 21.5194 | HPSv2: 0.2734 | HPSv2.1: 0.2754 | ImageReward: 0.4270
+
+[gradient_ascent] Batch 440/500 | Samples: 440/500 | Reward (t=1): 0.9673 | Reward (Avg): 0.7067 | CLIP: 26.1450 | Aesthetic: 5.9661 | PickScore: 21.5111 | HPSv2: 0.2732 | HPSv2.1: 0.2754 | ImageReward: 0.4217
+
+[gradient_ascent] Batch 450/500 | Samples: 450/500 | Reward (t=1): 0.7007 | Reward (Avg): 0.7084 | CLIP: 26.1783 | Aesthetic: 5.9595 | PickScore: 21.5085 | HPSv2: 0.2732 | HPSv2.1: 0.2749 | ImageReward: 0.4202
+
+[gradient_ascent] Batch 460/500 | Samples: 460/500 | Reward (t=1): 0.1193 | Reward (Avg): 0.7053 | CLIP: 26.1434 | Aesthetic: 5.9576 | PickScore: 21.5003 | HPSv2: 0.2732 | HPSv2.1: 0.2747 | ImageReward: 0.4168
+
+[gradient_ascent] Batch 470/500 | Samples: 470/500 | Reward (t=1): 0.6265 | Reward (Avg): 0.7034 | CLIP: 26.0937 | Aesthetic: 5.9501 | PickScore: 21.4891 | HPSv2: 0.2727 | HPSv2.1: 0.2744 | ImageReward: 0.4166
+
+[gradient_ascent] Batch 480/500 | Samples: 480/500 | Reward (t=1): 0.8589 | Reward (Avg): 0.7012 | CLIP: 26.0558 | Aesthetic: 5.9510 | PickScore: 21.4883 | HPSv2: 0.2727 | HPSv2.1: 0.2742 | ImageReward: 0.4193
+
+[gradient_ascent] Batch 490/500 | Samples: 490/500 | Reward (t=1): 0.8589 | Reward (Avg): 0.7001 | CLIP: 26.0536 | Aesthetic: 5.9529 | PickScore: 21.4798 | HPSv2: 0.2727 | HPSv2.1: 0.2739 | ImageReward: 0.4213
+
+[gradient_ascent] Batch 500/500 | Samples: 500/500 | Reward (t=1): 0.7222 | Reward (Avg): 0.6998 | CLIP: 26.0221 | Aesthetic: 5.9536 | PickScore: 21.4608 | HPSv2: 0.2727 | HPSv2.1: 0.2739 | ImageReward: 0.4125
+✓ Gradient Ascent Avg Reward: 0.6998
+✓ Gradient Ascent Avg CLIP Score: 26.0221
+✓ Gradient Ascent Avg Aesthetic Score: 5.9536
+✓ Gradient Ascent Avg PickScore: 21.4608
+✓ Gradient Ascent Avg HPSv2 Score: 0.2727
+✓ Gradient Ascent Avg HPSv2.1 Score: 0.2739
+✓ Gradient Ascent Avg ImageReward: 0.4125
+
+Gradient Ascent Statistics:
+  Applications: 12
+  Total reward improvement: +1.2930
+  Avg reward improvement: +0.1077
+
+✓ Saved LR curve plot to: RESULTS/pickapic/one_step_rectification_config_spo/run_1/lr_curve.png
+   Total gradient steps: 12
+   LR range: 1.000000 → 1.000000
+
+✓ Saved Rewards curve plot to: RESULTS/pickapic/one_step_rectification_config_spo/run_1/rewards_curve.png
+   Total gradient steps: 24
+   Reward range: 0.9487 → 0.9990
+   Total improvement: +0.0503
+
+======================================================================
+FINAL RESULTS
+======================================================================
+
+Gradient Ascent:
+  Avg Reward:       0.6998
+  Avg CLIP Score:   26.0221
+  Avg Aesthetic:    5.9536
+  Avg PickScore:    21.4608
+  Avg HPSv2:        0.2727
+  Avg HPSv2.1:      0.2739
+  Avg ImageReward:  0.4125
+
+✓ Results saved to: RESULTS/pickapic/one_step_rectification_config_spo/run_1/evaluation_results.txt
+
+======================================================================

Reward_sd15_idealized/config_analysis_tuning.ipynb

@@ -0,0 +1,218 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a24d02a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "from pathlib import Path\n",
+    "from datetime import datetime\n",
+    "import warnings\n",
+    "warnings.filterwarnings('ignore')\n",
+    "\n",
+    "# ============================================================================\n",
+    "# SECTION 1: Load and Parse Results from GPU Tuning Runs\n",
+    "# ==========================-==================================================\n",
+    "print(\"=\" * 80)\n",
+    "print(\"LOADING TUNING RESULTS FROM GPU RUNS\")\n",
+    "print(\"=\" * 80)\n",
+    "\n",
+    "results_dir = Path(\"RESULTS_TURNING/run_2\")\n",
+    "all_experiments = []\n",
+    "baseline_metrics = None\n",
+    "\n",
+    "# Collect results from all GPU runs\n",
+    "for gpu_id in range(8):\n",
+    "    gpu_dir = results_dir / f\"gpu_{gpu_id}\"\n",
+    "    results_file = gpu_dir / \"tuning_results.json\"\n",
+    "    \n",
+    "    if results_file.exists():\n",
+    "        with open(results_file, 'r') as f:\n",
+    "            data = json.load(f)\n",
+    "        \n",
+    "        # Extract baseline (same across all GPUs)\n",
+    "        if baseline_metrics is None and \"baseline\" in data:\n",
+    "            baseline_metrics = data[\"baseline\"][\"metrics\"]\n",
+    "            print(f\"\\n📊 Baseline Metrics (cfg_scale=5.0):\")\n",
+    "            for metric, value in baseline_metrics.items():\n",
+    "                print(f\"   {metric:15s}: {value:.6f}\")\n",
+    "        \n",
+    "        # Collect all experiments\n",
+    "        if \"experiments\" in data:\n",
+    "            all_experiments.extend(data[\"experiments\"])\n",
+    "            print(f\"✓ GPU {gpu_id}: {len(data['experiments'])} results loaded\")\n",
+    "\n",
+    "print(f\"\\n✓ Total experiments loaded: {len(all_experiments)}\")\n",
+    "\n",
+    "# ============================================================================\n",
+    "# SECTION 2: Filter Top Configs with Improvements Across All Metrics\n",
+    "# ============================================================================\n",
+    "print(\"\\n\" + \"=\" * 80)\n",
+    "print(\"FILTERING CONFIGURATIONS WITH IMPROVEMENTS IN ALL METRICS\")\n",
+    "print(\"=\" * 80)\n",
+    "\n",
+    "# Define improvement metrics to track (using ImageReward instead of Reward)\n",
+    "improvement_metrics = [\n",
+    "                        \"aesthetic_improvement\", \n",
+    "                        \"imagereward_improvement\", \n",
+    "                        \"clip_improvement\", \n",
+    "                        \"pickscore_improvement\", \n",
+    "                        \"hpsv2_improvement\"\n",
+    "                      ]\n",
+    "\n",
+    "# Filter experiments with improvements in ALL metrics\n",
+    "top_configs = []\n",
+    "\n",
+    "for exp in all_experiments:\n",
+    "    if \"improvements\" not in exp or \"config\" not in exp or \"metrics\" not in exp:\n",
+    "        continue\n",
+    "    \n",
+    "    improvements = exp[\"improvements\"]\n",
+    "    config = exp[\"config\"]\n",
+    "    metrics = exp[\"metrics\"]\n",
+    "    \n",
+    "    # Check if ALL improvements are positive (>0)\n",
+    "    all_positive = all(improvements.get(metric, -1) > 0 for metric in improvement_metrics)\n",
+    "    \n",
+    "    if all_positive:\n",
+    "        # Calculate aggregate improvement score\n",
+    "        avg_improvement = np.mean([improvements.get(metric, 0) for metric in improvement_metrics])\n",
+    "        \n",
+    "        top_configs.append({\n",
+    "            \"config\": config,\n",
+    "            \"metrics\": metrics,\n",
+    "            \"improvements\": improvements,\n",
+    "            \"avg_improvement\": avg_improvement\n",
+    "        })\n",
+    "\n",
+    "print(f\"✓ Found {len(top_configs)} configurations with improvements in ALL metrics\")\n",
+    "\n",
+    "# Sort by average improvement\n",
+    "top_configs.sort(key=lambda x: x[\"avg_improvement\"], reverse=True)\n",
+    "\n",
+    "# Get top 10\n",
+    "top_10 = top_configs[:10]\n",
+    "print(f\"✓ Extracted top 10 best performing configurations\")\n",
+    "\n",
+    "# ============================================================================\n",
+    "# SECTION 3: Create Comprehensive Results Table\n",
+    "# ============================================================================\n",
+    "print(\"\\n\" + \"=\" * 80)\n",
+    "print(\"CREATING COMPREHENSIVE RESULTS TABLE\")\n",
+    "print(\"=\" * 80)\n",
+    "\n",
+    "# Build detailed table data\n",
+    "table_data = []\n",
+    "\n",
+    "for rank, result in enumerate(top_10, 1):\n",
+    "    cfg = result[\"config\"]\n",
+    "    metrics = result[\"metrics\"]\n",
+    "    improvements = result[\"improvements\"]\n",
+    "    \n",
+    "    row = {\n",
+    "        \"Rank\": rank,\n",
+    "        \"CFG Scale\": cfg.get(\"cfg_scale\", \"N/A\"),\n",
+    "        \"Grad Config\": cfg.get(\"grad_config\", \"N/A\"),\n",
+    "        \"Steps\": cfg.get(\"num_grad_steps\", \"N/A\"),\n",
+    "        \"LR\": cfg.get(\"grad_step_size\", \"N/A\"),\n",
+    "        \"Momentum\": cfg.get(\"momentum\", \"N/A\"),\n",
+    "        \"ImageReward\": f\"{metrics.get('imagereward', 0):.6f}\",\n",
+    "        \"ImageReward ↑\": f\"{improvements.get('imagereward_improvement', 0):+.2f}%\",\n",
+    "        \"CLIP\": f\"{metrics.get('clip', 0):.4f}\",\n",
+    "        \"CLIP ↑\": f\"{improvements.get('clip_improvement', 0):+.2f}%\",\n",
+    "        \"Aesthetic\": f\"{metrics.get('aesthetic', 0):.4f}\",\n",
+    "        \"Aesthetic ↑\": f\"{improvements.get('aesthetic_improvement', 0):+.2f}%\",\n",
+    "        \"PickScore\": f\"{metrics.get('pickscore', 0):.4f}\",\n",
+    "        \"PickScore ↑\": f\"{improvements.get('pickscore_improvement', 0):+.2f}%\",\n",
+    "        \"HPSv2\": f\"{metrics.get('hpsv2', 0):.4f}\",\n",
+    "        \"HPSv2 ↑\": f\"{improvements.get('hpsv2_improvement', 0):+.2f}%\",\n",
+    "        \"Avg Improvement\": f\"{result['avg_improvement']:+.2f}%\",\n",
+    "    }\n",
+    "    \n",
+    "    table_data.append(row)\n",
+    "\n",
+    "df_top_10 = pd.DataFrame(table_data)\n",
+    "\n",
+    "print(\"\\n📋 TOP 10 CONFIGURATIONS WITH IMPROVEMENTS IN ALL METRICS:\")\n",
+    "print(\"=\" * 180)\n",
+    "print(df_top_10.to_string(index=False))\n",
+    "print(\"=\" * 180)\n",
+    "\n",
+    "# ============================================================================\n",
+    "# SECTION 4: Visualize and Summary Statistics\n",
+    "# ============================================================================\n",
+    "print(\"\\n\" + \"=\" * 80)\n",
+    "print(\"SUMMARY STATISTICS\")\n",
+    "print(\"=\" * 80)\n",
+    "\n",
+    "# Extract numeric improvement values for analysis\n",
+    "improvement_summary = []\n",
+    "for result in top_10:\n",
+    "    improvements = result[\"improvements\"]\n",
+    "    for metric in [\"imagereward_improvement\", \"clip_improvement\", \"aesthetic_improvement\", \n",
+    "                   \"pickscore_improvement\", \"hpsv2_improvement\"]:\n",
+    "        metric_name = metric.replace(\"_improvement\", \"\").upper()\n",
+    "        improvement_summary.append({\n",
+    "            \"Metric\": metric_name,\n",
+    "            \"Improvement %\": improvements.get(metric, 0)\n",
+    "        })\n",
+    "\n",
+    "df_summary = pd.DataFrame(improvement_summary)\n",
+    "\n",
+    "print(\"\\n📊 Average Improvements by Metric (Top 10):\")\n",
+    "metric_stats = df_summary.groupby(\"Metric\")[\"Improvement %\"].agg([\"mean\", \"std\", \"min\", \"max\"])\n",
+    "print(metric_stats.round(2))\n",
+    "\n",
+    "print(\"\\n📈 Best Configuration Details:\")\n",
+    "best = top_10[0]\n",
+    "best_cfg = best[\"config\"]\n",
+    "best_metrics = best[\"metrics\"]\n",
+    "best_improvements = best[\"improvements\"]\n",
+    "\n",
+    "print(f\"\\n✓ RANK #1 - Best Performing Configuration:\")\n",
+    "print(f\"   Configuration:\")\n",
+    "print(f\"      • CFG Scale: {best_cfg.get('cfg_scale')}\")\n",
+    "print(f\"      • Gradient Config: {best_cfg.get('grad_config')}\")\n",
+    "print(f\"      • Gradient Steps: {best_cfg.get('num_grad_steps')}\")\n",
+    "print(f\"      • Step Size: {best_cfg.get('grad_step_size')}\")\n",
+    "print(f\"      • Momentum: {best_cfg.get('momentum')}\")\n",
+    "print(f\"\\n   Metrics:\")\n",
+    "for metric in [\"imagereward\", \"clip\", \"aesthetic\", \"pickscore\", \"hpsv2\"]:\n",
+    "    baseline_val = baseline_metrics.get(metric, 0)\n",
+    "    current_val = best_metrics.get(metric, 0)\n",
+    "    improvement = best_improvements.get(f\"{metric}_improvement\", 0)\n",
+    "    print(f\"      • {metric:12s}: {current_val:8.6f} (baseline: {baseline_val:8.6f}) ↑ {improvement:+6.2f}%\")\n",
+    "\n",
+    "print(\"\\n\" + \"=\" * 80)\n",
+    "print(\"✓ ANALYSIS COMPLETE - TOP 10 CONFIGURATIONS IDENTIFIED\")\n",
+    "print(\"=\" * 80)"
+   ]
+  }
+ ],
+ "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/eval.py

@@ -0,0 +1,1422 @@
+"""
+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()
+

Reward_sd15_idealized/examples.sh

@@ -0,0 +1,154 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+clear
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+cd "$SCRIPT_DIR"
+
+# Shared HF cache used on this cluster.
+HF_HUB_CACHE_DIR="${HF_HUB_CACHE_DIR:-/scratch/rr81/ma5430/.cache/huggingface/hub}"
+export HF_HUB_CACHE="$HF_HUB_CACHE_DIR"
+export HUGGINGFACE_HUB_CACHE="$HF_HUB_CACHE_DIR"
+export HF_HOME="$(dirname "$HF_HUB_CACHE_DIR")"
+
+# GPU nodes have no internet, while login nodes do.
+# Auto default: offline on GPU nodes, online on login nodes.
+DEFAULT_OFFLINE_MODE="1"
+if ! (command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1); then
+    DEFAULT_OFFLINE_MODE="0"
+fi
+OFFLINE_MODE="${OFFLINE_MODE:-$DEFAULT_OFFLINE_MODE}"
+
+if [[ "$OFFLINE_MODE" == "1" ]]; then
+    export HF_DATASETS_OFFLINE="1"
+    export HF_METRICS_OFFLINE="1"
+    export HF_MODULES_OFFLINE="1"
+    export TRANSFORMERS_OFFLINE="1"
+    export DIFFUSERS_OFFLINE="1"
+    export HF_HUB_OFFLINE="1"
+else
+    export HF_DATASETS_OFFLINE="0"
+    export HF_METRICS_OFFLINE="0"
+    export HF_MODULES_OFFLINE="0"
+    export TRANSFORMERS_OFFLINE="0"
+    export DIFFUSERS_OFFLINE="0"
+    export HF_HUB_OFFLINE="0"
+fi
+
+# Existing environment requested by user.
+PYTHON_BIN="${PYTHON_BIN:-/g/data/rr81/aev/bin/python}"
+if [[ ! -x "$PYTHON_BIN" ]]; then
+    echo "[examples.sh] Missing Python executable: $PYTHON_BIN" >&2
+    exit 1
+fi
+
+DATASET_NAME="${DATASET_NAME:-pickapic}"      # coco | pickapic
+GRAD_CONFIG="${GRAD_CONFIG:-one_step_rectification_config}"
+MODEL_VARIANT="${MODEL_VARIANT:-spo}"      # origin | spo | diffusion_dpo | lpo
+MODE="${MODE:-gradient_ascent}"               # gradient_ascent | baseline | both
+# Empty MAX_SAMPLES means evaluate all available samples.
+MAX_SAMPLES="${MAX_SAMPLES:-}"
+NUM_STEPS="${NUM_STEPS:-20}"
+CFG_SCALE="${CFG_SCALE:-3}"
+METRICS="${METRICS:-clip aesthetic pickscore hpsv2 hpsv21 imagereward}"
+PREFETCH_ONLY="${PREFETCH_ONLY:-0}"
+
+# Override this path whenever you want to swap reward weights.
+LRM_MODEL_PATH="${LRM_MODEL_PATH:-/g/data/rr81/LPO/lrm/lrm_15/LRM}"
+
+if [[ -z "${GPU_ID:-}" ]]; then
+    if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
+        GPU_ID="$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits | sort -k2 -n | head -n1 | cut -d',' -f1 | tr -d ' ')"
+        GPU_ID="${GPU_ID:-0}"
+    else
+        GPU_ID="0"
+        echo "[examples.sh] No visible NVIDIA GPU on this node. Defaulting GPU_ID=0."
+        echo "[examples.sh] eval.py will run on CPU if CUDA is unavailable."
+    fi
+fi
+
+echo "Using GPU ID: $GPU_ID"
+echo "Using LRM weights: $LRM_MODEL_PATH"
+echo "HF offline mode: $OFFLINE_MODE"
+
+if [[ "$PREFETCH_ONLY" == "1" ]]; then
+    echo "[examples.sh] PREFETCH_ONLY=1 -> downloading required model files to shared cache and exiting."
+    export MODEL_VARIANT
+    export METRICS
+    "$PYTHON_BIN" - <<'PY'
+import os
+from huggingface_hub import hf_hub_download, snapshot_download
+
+cache_dir = os.environ["HF_HUB_CACHE"]
+model_variant = os.environ.get("MODEL_VARIANT", "origin")
+metrics = set(os.environ.get("METRICS", "clip aesthetic").split())
+
+def snap(repo_id):
+    print(f"[prefetch] snapshot_download: {repo_id}")
+    snapshot_download(repo_id=repo_id, cache_dir=cache_dir, local_files_only=False)
+
+def one(repo_id, filename):
+    print(f"[prefetch] hf_hub_download: {repo_id}/{filename}")
+    hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir, local_files_only=False)
+
+# Base model for generation + reward backbone
+snap("stable-diffusion-v1-5/stable-diffusion-v1-5")
+
+# Variant-specific generation checkpoints (if used)
+if model_variant == "spo":
+    snap("SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep")
+elif model_variant == "diffusion_dpo":
+    snap("mhdang/dpo-sd1.5-text2image-v1")
+elif model_variant == "lpo":
+    snap("casiatao/LPO")
+
+# Required for CLIP-based metrics and LRM text projection init fallback
+if "clip" in metrics or "aesthetic" in metrics:
+    snap("openai/clip-vit-large-patch14")
+
+if "pickscore" in metrics:
+    snap("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
+    snap("yuvalkirstain/PickScore_v1")
+
+if "hpsv2" in metrics or "hpsv21" in metrics:
+    one("laion/CLIP-ViT-H-14-laion2B-s32B-b79K", "open_clip_pytorch_model.bin")
+if "hpsv2" in metrics:
+    one("xswu/HPSv2", "HPS_v2_compressed.pt")
+if "hpsv21" in metrics:
+    one("xswu/HPSv2", "HPS_v2.1_compressed.pt")
+
+if "imagereward" in metrics:
+    one("THUDM/ImageReward", "ImageReward.pt")
+    one("THUDM/ImageReward", "med_config.json")
+
+print("[prefetch] done")
+PY
+    exit 0
+fi
+
+read -r -a METRICS_ARR <<< "$METRICS"
+
+CMD=(
+    "$PYTHON_BIN" eval.py
+    --model_variant "$MODEL_VARIANT"
+    --dataset_type "$DATASET_NAME"
+    --lrm_model "$LRM_MODEL_PATH"
+    --grad_config "$GRAD_CONFIG"
+    --metrics "${METRICS_ARR[@]}"
+    --num_steps "$NUM_STEPS"
+    --cfg_scale "$CFG_SCALE"
+    --hf_cache_dir "$HF_HUB_CACHE_DIR"
+    --output_dir "RESULTS/$DATASET_NAME/${GRAD_CONFIG}_${MODEL_VARIANT}"
+    --cuda "$GPU_ID"
+    --mode "$MODE"
+)
+
+if [[ -n "$MAX_SAMPLES" ]]; then
+    CMD+=(--max_samples "$MAX_SAMPLES")
+fi
+
+if [[ "$OFFLINE_MODE" == "1" ]]; then
+    CMD+=(--offline)
+fi
+
+"${CMD[@]}"

Reward_sd15_idealized/lr_scheduler.py

@@ -0,0 +1,233 @@
+"""
+Learning rate schedulers for gradient ascent optimization.
+
+Provides various LR scheduling strategies for reward-guided gradient ascent,
+including cosine annealing, linear decay, and custom schedules.
+"""
+
+import math
+from typing import Optional, Literal
+
+
+class LRScheduler:
+    """Base class for learning rate schedulers."""
+    
+    def __init__(self, initial_lr: float, num_steps: int):
+        """
+        Initialize LR scheduler.
+        
+        Args:
+            initial_lr: Initial learning rate
+            num_steps: Total number of optimization steps
+        """
+        self.initial_lr = initial_lr
+        self.num_steps = num_steps
+        self.current_step = 0
+    
+    def get_lr(self) -> float:
+        """Get current learning rate."""
+        raise NotImplementedError
+    
+    def step(self):
+        """Update scheduler state after a step."""
+        self.current_step += 1
+    
+    def reset(self):
+        """Reset scheduler state."""
+        self.current_step = 0
+
+
+class ConstantLR(LRScheduler):
+    """Constant learning rate (no scheduling)."""
+    
+    def get_lr(self) -> float:
+        return self.initial_lr
+
+
+class LinearLR(LRScheduler):
+    """Linear learning rate decay."""
+    
+    def __init__(
+        self,
+        initial_lr: float,
+        num_steps: int,
+        end_lr: float = 0.0,
+        start_step: int = 0,
+    ):
+        """
+        Initialize linear LR scheduler.
+        
+        Args:
+            initial_lr: Starting learning rate
+            num_steps: Total number of steps
+            end_lr: Ending learning rate (default: 0.0)
+            start_step: Step to begin decay (default: 0)
+        """
+        super().__init__(initial_lr, num_steps)
+        self.end_lr = end_lr
+        self.start_step = start_step
+    
+    def get_lr(self) -> float:
+        if self.current_step < self.start_step:
+            return self.initial_lr
+        
+        progress = (self.current_step - self.start_step) / (self.num_steps - self.start_step)
+        progress = min(1.0, progress)
+        
+        return self.initial_lr + (self.end_lr - self.initial_lr) * progress
+
+
+class CosineLR(LRScheduler):
+    """Cosine annealing learning rate schedule."""
+    
+    def __init__(
+        self,
+        initial_lr: float,
+        num_steps: int,
+        min_lr: float = 0.0,
+        warmup_steps: int = 0,
+    ):
+        """
+        Initialize cosine LR scheduler.
+        
+        Args:
+            initial_lr: Maximum learning rate
+            num_steps: Total number of steps
+            min_lr: Minimum learning rate (default: 0.0)
+            warmup_steps: Number of linear warmup steps (default: 0)
+        """
+        super().__init__(initial_lr, num_steps)
+        self.min_lr = min_lr
+        self.warmup_steps = warmup_steps
+    
+    def get_lr(self) -> float:
+        if self.current_step < self.warmup_steps:
+            # Linear warmup
+            return self.initial_lr * (self.current_step / self.warmup_steps)
+        
+        # Cosine annealing
+        progress = (self.current_step - self.warmup_steps) / (self.num_steps - self.warmup_steps)
+        progress = min(1.0, progress)
+        
+        cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
+        return self.min_lr + (self.initial_lr - self.min_lr) * cosine_decay
+
+
+class ExponentialLR(LRScheduler):
+    """Exponential learning rate decay."""
+    
+    def __init__(
+        self,
+        initial_lr: float,
+        num_steps: int,
+        gamma: float = 0.95,
+    ):
+        """
+        Initialize exponential LR scheduler.
+        
+        Args:
+            initial_lr: Starting learning rate
+            num_steps: Total number of steps
+            gamma: Multiplicative decay factor per step
+        """
+        super().__init__(initial_lr, num_steps)
+        self.gamma = gamma
+    
+    def get_lr(self) -> float:
+        return self.initial_lr * (self.gamma ** self.current_step)
+
+
+class StepLR(LRScheduler):
+    """Step-wise learning rate decay."""
+    
+    def __init__(
+        self,
+        initial_lr: float,
+        num_steps: int,
+        step_size: int,
+        gamma: float = 0.1,
+    ):
+        """
+        Initialize step LR scheduler.
+        
+        Args:
+            initial_lr: Starting learning rate
+            num_steps: Total number of steps
+            step_size: Number of steps between each decay
+            gamma: Multiplicative decay factor
+        """
+        super().__init__(initial_lr, num_steps)
+        self.step_size = step_size
+        self.gamma = gamma
+    
+    def get_lr(self) -> float:
+        num_decays = self.current_step // self.step_size
+        return self.initial_lr * (self.gamma ** num_decays)
+
+
+def create_lr_scheduler(
+    scheduler_type: Literal["constant", "linear", "cosine", "exponential", "step"],
+    initial_lr: float,
+    num_steps: int,
+    **kwargs
+) -> LRScheduler:
+    """
+    Factory function to create learning rate schedulers.
+    
+    Args:
+        scheduler_type: Type of scheduler ("constant", "linear", "cosine", "exponential", "step")
+        initial_lr: Initial learning rate
+        num_steps: Total number of optimization steps
+        **kwargs: Additional scheduler-specific arguments
+            For linear: end_lr, start_step
+            For cosine: min_lr, warmup_steps
+            For exponential: gamma
+            For step: step_size, gamma
+    
+    Returns:
+        LRScheduler instance
+    
+    Examples:
+        # Constant LR
+        scheduler = create_lr_scheduler("constant", initial_lr=0.1, num_steps=100)
+        
+        # Linear decay
+        scheduler = create_lr_scheduler("linear", initial_lr=0.1, num_steps=100, end_lr=0.01)
+        
+        # Cosine annealing with warmup
+        scheduler = create_lr_scheduler("cosine", initial_lr=0.1, num_steps=100, 
+                                       min_lr=0.001, warmup_steps=10)
+    """
+    if scheduler_type == "constant":
+        return ConstantLR(initial_lr, num_steps)
+    
+    elif scheduler_type == "linear":
+        return LinearLR(
+            initial_lr, num_steps,
+            end_lr=kwargs.get("end_lr", 0.0),
+            start_step=kwargs.get("start_step", 0),
+        )
+    
+    elif scheduler_type == "cosine":
+        return CosineLR(
+            initial_lr, num_steps,
+            min_lr=kwargs.get("min_lr", 0.0),
+            warmup_steps=kwargs.get("warmup_steps", 0),
+        )
+    
+    elif scheduler_type == "exponential":
+        return ExponentialLR(
+            initial_lr, num_steps,
+            gamma=kwargs.get("gamma", 0.95),
+        )
+    
+    elif scheduler_type == "step":
+        return StepLR(
+            initial_lr, num_steps,
+            step_size=kwargs.get("step_size", 10),
+            gamma=kwargs.get("gamma", 0.1),
+        )
+    
+    else:
+        raise ValueError(f"Unknown scheduler type: {scheduler_type}. "
+                        f"Choose from: constant, linear, cosine, exponential, step")

Reward_sd15_idealized/models/__init__.py

@@ -0,0 +1,3 @@
+from .reward_model import LRMRewardModel
+
+__all__ = ['LRMRewardModel', 'LRMRewardModelSDXL']

Reward_sd15_idealized/models/reward_model.py

@@ -0,0 +1,388 @@
+"""
+LRM Reward Model Wrapper
+Loads LRM weights from HuggingFace and provides interface for computing preference scores on noisy latents.
+"""
+
+import torch
+from torch import nn
+from diffusers import AutoencoderKL, DDPMScheduler
+from transformers import CLIPTextModel, CLIPTokenizer
+from huggingface_hub import hf_hub_download
+import os
+from .unet_2d_condition_reward import UNet2DConditionModel
+
+
+def _offline_mode_enabled() -> bool:
+    return os.getenv("HF_HUB_OFFLINE", "0").strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _get_cache_dir() -> str | None:
+    return os.getenv("HF_HUB_CACHE") or os.getenv("HUGGINGFACE_HUB_CACHE")
+
+
+def _hf_pretrained_kwargs() -> dict:
+    kwargs = {"local_files_only": _offline_mode_enabled()}
+    cache_dir = _get_cache_dir()
+    if cache_dir:
+        kwargs["cache_dir"] = cache_dir
+    return kwargs
+
+
+class LRMRewardModel(nn.Module):
+    """
+    Latent Reward Model (LRM) for SD1.5
+    
+    This model computes preference scores directly on noisy latent images at any timestep.
+    It uses features from the U-Net and text encoder to predict how well an image aligns
+    with the prompt at different noise levels.
+    
+    Args:
+        pretrained_model_name_or_path: Base SD model path (e.g., 'runwayml/stable-diffusion-v1-5')
+        lrm_model_path: Path to LRM checkpoint from HuggingFace (e.g., 'casiatao/LRM')
+        clip_model_path: Path to CLIP checkpoint for text projection initialization
+        guidance_scale: Classifier-free guidance scale (default: 7.5)
+        device: Device to load model on
+    """
+    
+    def __init__(
+        self,
+        pretrained_model_name_or_path='runwayml/stable-diffusion-v1-5',
+        lrm_model_path=None,
+        clip_model_path='openai/clip-vit-large-patch14',
+        guidance_scale=7.5,
+        device='cuda'
+    ):
+        super().__init__()
+        
+        self.device = device
+        self.guidance_scale = guidance_scale
+        self.multi_scale = True
+        self.multi_scale_cfg = False
+        
+        print(f"Loading base models from {pretrained_model_name_or_path}...")
+        pretrained_kwargs = _hf_pretrained_kwargs()
+        
+        # Load tokenizer and text encoder
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+            pretrained_model_name_or_path, 
+            subfolder="tokenizer",
+            **pretrained_kwargs,
+        )
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            pretrained_model_name_or_path, 
+            subfolder="text_encoder",
+            **pretrained_kwargs,
+        ).to(device)
+        
+        # Load VAE (frozen, only needed for preprocessing if using images)
+        self.vae = AutoencoderKL.from_pretrained(
+            pretrained_model_name_or_path, 
+            subfolder="vae",
+            **pretrained_kwargs,
+        ).to(device)
+        self.vae.requires_grad_(False)
+        
+        # Load scheduler
+        self.scheduler = DDPMScheduler.from_pretrained(
+            pretrained_model_name_or_path, 
+            subfolder="scheduler",
+            **pretrained_kwargs,
+        )
+        
+        # Load U-Net with custom reward architecture
+        print("Loading custom U-Net for reward prediction...")
+        self.unet = UNet2DConditionModel.from_pretrained(
+            pretrained_model_name_or_path, 
+            subfolder="unet",
+            **pretrained_kwargs,
+        ).to(device)
+        
+        # Global pooling layer
+        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+        
+        # Projection layers
+        # Multi-scale: concatenates features from 4 down blocks + mid block = 4800 dims
+        vision_embed_dim = 4800 if self.multi_scale else 1280
+        text_embed_dim = 768
+        projection_dim = 768
+        
+        self.visual_projection = nn.Linear(vision_embed_dim, projection_dim, bias=False).to(device)
+        self.text_projection = nn.Linear(text_embed_dim, projection_dim, bias=False).to(device)
+        
+        # Initialize text projection from CLIP
+        print(f"Loading CLIP text projection from {clip_model_path}...")
+        try:
+            # Try loading from local path first
+            if os.path.exists(clip_model_path):
+                clip_ckpt = torch.load(clip_model_path, map_location='cpu')
+            else:
+                # Download from HuggingFace
+                clip_ckpt_path = hf_hub_download(
+                    repo_id=clip_model_path,
+                    filename="pytorch_model.bin",
+                    **_hf_pretrained_kwargs(),
+                )
+                clip_ckpt = torch.load(clip_ckpt_path, map_location='cpu')
+            
+            self.text_projection.weight.data = clip_ckpt['text_projection.weight'].contiguous().to(device)
+            print("✓ Loaded CLIP text projection weights")
+        except Exception as e:
+            print(f"Warning: Could not load CLIP weights: {e}")
+            print("Initializing text projection randomly")
+            nn.init.normal_(self.text_projection.weight, std=0.02)
+        
+        # Initialize visual projection
+        nn.init.normal_(self.visual_projection.weight, std=0.02)
+        
+        # Logit scale (temperature parameter)
+        self.logit_scale = nn.Parameter(torch.ones([]) * 2.6592).to(device)
+        
+        # Setup classifier-free guidance
+        self.do_classifier_free_guidance = self.guidance_scale > 1.0
+        if self.do_classifier_free_guidance:
+            self.neg_prompt_ids = self.tokenizer(
+                [""],
+                return_tensors="pt",
+                padding="max_length",
+                truncation=True,
+                max_length=self.tokenizer.model_max_length,
+            ).input_ids.to(device)
+        
+        # Load fine-tuned LRM weights if provided
+        if lrm_model_path:
+            self.load_lrm_weights(lrm_model_path)
+        
+        print("✓ LRM Reward Model initialized successfully!")
+    
+    def load_lrm_weights(self, model_path):
+        """
+        Load fine-tuned LRM weights from HuggingFace or local path
+        
+        Expected structure:
+        - unet/ (directory with U-Net weights)
+        - text_encoder/ (optional, directory with text encoder weights)
+        - state_dict.pt (visual_projection, text_projection, logit_scale)
+        """
+        print(f"\nLoading LRM weights from {model_path}...")
+        
+        try:
+            # Check if it's a HuggingFace model or local path
+            if not os.path.exists(model_path):
+                # Try to download from HuggingFace
+                print(f"Downloading from HuggingFace: {model_path}")
+                # For HF models, we need to download the entire repo
+                from huggingface_hub import snapshot_download
+                model_path = snapshot_download(repo_id=model_path, **_hf_pretrained_kwargs())
+            
+            # Load U-Net weights
+            unet_path = os.path.join(model_path, "lrm_sd15", "unet")
+            if os.path.exists(unet_path):
+                self.unet = UNet2DConditionModel.from_pretrained(unet_path, **_hf_pretrained_kwargs()).to(self.device)
+                print(f"✓ Loaded U-Net weights from {unet_path}")
+            else:
+                print(f"Warning: U-Net path not found: {unet_path}")
+            
+            # Load text encoder weights (optional)
+            text_encoder_path = os.path.join(model_path, "lrm_sd15", "text_encoder")
+            if os.path.exists(text_encoder_path):
+                self.text_encoder = CLIPTextModel.from_pretrained(text_encoder_path, **_hf_pretrained_kwargs()).to(self.device)
+                print(f"✓ Loaded text encoder weights from {text_encoder_path}")
+            
+            # Load projection layers and logit scale
+            state_dict_path = os.path.join(model_path, "lrm_sd15", "state_dict.pt")
+            if os.path.exists(state_dict_path):
+                state_dict = torch.load(state_dict_path, map_location='cpu')
+                
+                self.visual_projection.load_state_dict(state_dict['visual_projection'])
+                self.text_projection.load_state_dict(state_dict['text_projection'])
+                
+                # Move projection layers to device
+                self.visual_projection = self.visual_projection.to(self.device)
+                self.text_projection = self.text_projection.to(self.device)
+                
+                logit_scale_val = state_dict['logit_scale']
+                if isinstance(logit_scale_val, torch.Tensor):
+                    self.logit_scale.data = logit_scale_val.to(self.device)
+                else:
+                    self.logit_scale.data = torch.tensor(logit_scale_val).to(self.device)
+                
+                print(f"✓ Loaded projection layers and logit_scale from {state_dict_path}")
+            else:
+                print(f"Warning: state_dict.pt not found: {state_dict_path}")
+            
+            print("✓ Successfully loaded all LRM weights!")
+            
+        except Exception as e:
+            print(f"Error loading LRM weights: {e}")
+            print("Continuing with base model weights...")
+    
+    def encode_prompt(self, prompt):
+        """Tokenize text prompt"""
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        return text_inputs.input_ids.to(self.device)
+    
+    def get_text_features(self, text_input_ids):
+        """
+        Extract text features from prompt
+        Returns: (encoder_hidden_states, text_features)
+        """
+        if self.do_classifier_free_guidance:
+            # Concatenate conditional and unconditional prompts
+            text_input_ids = torch.cat([
+                text_input_ids, 
+                self.neg_prompt_ids.repeat(text_input_ids.shape[0], 1)
+            ], dim=0)
+        
+        outputs = self.text_encoder(text_input_ids, return_dict=False)
+        encoder_hidden_states = outputs[0]  # Sequence of hidden states
+        pooled_output = outputs[1]  # Pooled output (last token)
+        
+        if self.do_classifier_free_guidance:
+            pooled_output_text, pooled_output_ucond = pooled_output.chunk(2, dim=0)
+            text_features = self.text_projection(pooled_output_text)
+        else:
+            text_features = self.text_projection(pooled_output)
+        
+        return encoder_hidden_states, text_features
+    
+    def get_image_features(self, encoder_hidden_states, noisy_latents, timesteps):
+        """
+        Extract visual features from noisy latents using U-Net
+        
+        Args:
+            encoder_hidden_states: Text conditioning from CLIP
+            noisy_latents: Noisy latent images [B, C, H, W]
+            timesteps: Denoising timesteps [B]
+        
+        Returns:
+            image_features: Visual embeddings [B, projection_dim]
+        """
+        if self.do_classifier_free_guidance:
+            noisy_latents = torch.cat([noisy_latents] * 2, dim=0)
+            timesteps = torch.cat([timesteps] * 2, dim=0)
+        
+        # Forward through U-Net (only down blocks + mid block, no up blocks)
+        mid_output, down_block_res_samples = self.unet(
+            noisy_latents, 
+            timesteps, 
+            encoder_hidden_states=encoder_hidden_states,
+            return_dict=False, 
+            use_up_blocks=False
+        )
+        
+        if self.multi_scale:
+            # Extract multi-scale features from down blocks
+            # Indices correspond to: [320, 64, 64], [640, 32, 32], [1280, 16, 16], [1280, 8, 8]
+            first_stage_output = down_block_res_samples[2]    # 320 channels
+            second_stage_output = down_block_res_samples[5]   # 640 channels
+            third_stage_output = down_block_res_samples[8]    # 1280 channels
+            fourth_stage_output = down_block_res_samples[11]  # 1280 channels
+            
+            # Global average pooling
+            pooled_first = self.avg_pool(first_stage_output).squeeze(dim=[2, 3])
+            pooled_second = self.avg_pool(second_stage_output).squeeze(dim=[2, 3])
+            pooled_third = self.avg_pool(third_stage_output).squeeze(dim=[2, 3])
+            pooled_fourth = self.avg_pool(fourth_stage_output).squeeze(dim=[2, 3])
+            pooled_mid = self.avg_pool(mid_output).squeeze(dim=[2, 3])
+            
+            # Apply VFE (Visual Feature Enhancement) on mid block
+            if self.do_classifier_free_guidance:
+                pooled_mid_text, pooled_mid_ucond = pooled_mid.chunk(2, dim=0)
+                pooled_mid = pooled_mid_ucond + self.guidance_scale * (pooled_mid_text - pooled_mid_ucond)
+                
+                # For other blocks, optionally apply CFG or just use conditional branch
+                if self.multi_scale_cfg:
+                    pooled_first_text, pooled_first_ucond = pooled_first.chunk(2, dim=0)
+                    pooled_first = pooled_first_ucond + self.guidance_scale * (pooled_first_text - pooled_first_ucond)
+                    
+                    pooled_second_text, pooled_second_ucond = pooled_second.chunk(2, dim=0)
+                    pooled_second = pooled_second_ucond + self.guidance_scale * (pooled_second_text - pooled_second_ucond)
+                    
+                    pooled_third_text, pooled_third_ucond = pooled_third.chunk(2, dim=0)
+                    pooled_third = pooled_third_ucond + self.guidance_scale * (pooled_third_text - pooled_third_ucond)
+                    
+                    pooled_fourth_text, pooled_fourth_ucond = pooled_fourth.chunk(2, dim=0)
+                    pooled_fourth = pooled_fourth_ucond + self.guidance_scale * (pooled_fourth_text - pooled_fourth_ucond)
+                else:
+                    # Use only conditional (text-conditioned) branch
+                    pooled_first, _ = pooled_first.chunk(2, dim=0)
+                    pooled_second, _ = pooled_second.chunk(2, dim=0)
+                    pooled_third, _ = pooled_third.chunk(2, dim=0)
+                    pooled_fourth, _ = pooled_fourth.chunk(2, dim=0)
+            
+            # Concatenate all scales: 320 + 640 + 1280 + 1280 + 1280 = 4800
+            concat_pooled = torch.cat([
+                pooled_first, pooled_second, pooled_third, pooled_fourth, pooled_mid
+            ], dim=-1)
+            
+            image_features = self.visual_projection(concat_pooled)
+        else:
+            # Single scale (mid block only)
+            pooled_mid = self.avg_pool(mid_output).squeeze(dim=[2, 3])
+            if self.do_classifier_free_guidance:
+                pooled_mid_text, pooled_mid_ucond = pooled_mid.chunk(2, dim=0)
+                pooled_mid = pooled_mid_ucond + self.guidance_scale * (pooled_mid_text - pooled_mid_ucond)
+            image_features = self.visual_projection(pooled_mid)
+        
+        return image_features
+    
+    def get_reward_score(self, noisy_latents, prompt, timesteps, enable_grad=False):
+        """
+        Compute preference score for noisy latents at given timesteps
+        
+        Args:
+            noisy_latents: Noisy latent images [B, C, H, W]
+            prompt: Text prompt(s) (string or list of strings)
+            timesteps: Denoising timesteps [B] or scalar
+            enable_grad: If True, allows gradient computation (for gradient ascent)
+        
+        Returns:
+            scores: Preference scores [B]
+        """
+        def _compute():
+            # Ensure inputs are on correct device
+            latents = noisy_latents.to(self.device, dtype=self.unet.dtype)
+            
+            # Handle timesteps
+            if isinstance(timesteps, int):
+                ts = torch.tensor([timesteps] * latents.shape[0])
+            else:
+                ts = timesteps
+            ts = ts.to(self.device)
+            
+            # Encode prompt
+            text_input_ids = self.encode_prompt(prompt)
+            
+            # Get text and image features
+            encoder_hidden_states, text_features = self.get_text_features(text_input_ids)
+            image_features = self.get_image_features(encoder_hidden_states, latents, ts)
+            
+            # Normalize features
+            image_features = image_features / torch.norm(image_features, dim=-1, keepdim=True)
+            text_features = text_features / torch.norm(text_features, dim=-1, keepdim=True)
+            
+            # Compute similarity scores
+            scores = self.logit_scale.exp() * (text_features @ image_features.T)[0]
+            scores = torch.sigmoid(scores)  # Scale to [0, 1]
+            return scores
+            # return scores
+        
+        # If enable_grad is True, compute with gradients; otherwise use no_grad
+        if enable_grad:
+            return _compute()
+        else:
+            with torch.no_grad():
+                return _compute()
+    
+    def forward(self, noisy_latents, prompt, timesteps):
+        """Alias for get_reward_score for nn.Module compatibility"""
+        return self.get_reward_score(noisy_latents, prompt, timesteps)

Reward_sd15_idealized/models/unet_2d_condition_reward.py

@@ -0,0 +1,1334 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    Attention,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+    FusedAttnProcessor2_0,
+)
+from diffusers.models.embeddings import (
+    GaussianFourierProjection,
+    GLIGENTextBoundingboxProjection,
+    ImageHintTimeEmbedding,
+    ImageProjection,
+    ImageTimeEmbedding,
+    TextImageProjection,
+    TextImageTimeEmbedding,
+    TextTimeEmbedding,
+    TimestepEmbedding,
+    Timesteps,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.unets.unet_2d_blocks import (
+    get_down_block,
+    get_mid_block,
+    get_up_block,
+)
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNet2DConditionOutput(BaseOutput):
+    """
+    The output of [`UNet2DConditionModel`].
+
+    Args:
+        sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)`):
+            The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
+    """
+
+    sample: torch.Tensor = None
+
+
+class UNet2DConditionModel(
+    ModelMixin, ConfigMixin, FromOriginalModelMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin
+):
+    r"""
+    A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
+    shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+            Height and width of input/output sample.
+        in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
+        out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
+        center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
+        flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
+            Whether to flip the sin to cos in the time embedding.
+        freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
+            The tuple of downsample blocks to use.
+        mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
+            Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or
+            `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
+            The tuple of upsample blocks to use.
+        only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
+            Whether to include self-attention in the basic transformer blocks, see
+            [`~models.attention.BasicTransformerBlock`].
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+            The tuple of output channels for each block.
+        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
+        downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
+        mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
+            If `None`, normalization and activation layers is skipped in post-processing.
+        norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
+        cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+            The dimension of the cross attention features.
+        transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
+            [`~models.unets.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unets.unet_2d_blocks.CrossAttnUpBlock2D`],
+            [`~models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
+        reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling
+            blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for
+            [`~models.unets.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unets.unet_2d_blocks.CrossAttnUpBlock2D`],
+            [`~models.unets.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
+        encoder_hid_dim (`int`, *optional*, defaults to None):
+            If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
+            dimension to `cross_attention_dim`.
+        encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
+            If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
+            embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
+        attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
+        num_attention_heads (`int`, *optional*):
+            The number of attention heads. If not defined, defaults to `attention_head_dim`
+        resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
+            for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
+        class_embed_type (`str`, *optional*, defaults to `None`):
+            The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
+            `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
+        addition_embed_type (`str`, *optional*, defaults to `None`):
+            Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
+            "text". "text" will use the `TextTimeEmbedding` layer.
+        addition_time_embed_dim: (`int`, *optional*, defaults to `None`):
+            Dimension for the timestep embeddings.
+        num_class_embeds (`int`, *optional*, defaults to `None`):
+            Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
+            class conditioning with `class_embed_type` equal to `None`.
+        time_embedding_type (`str`, *optional*, defaults to `positional`):
+            The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.
+        time_embedding_dim (`int`, *optional*, defaults to `None`):
+            An optional override for the dimension of the projected time embedding.
+        time_embedding_act_fn (`str`, *optional*, defaults to `None`):
+            Optional activation function to use only once on the time embeddings before they are passed to the rest of
+            the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`.
+        timestep_post_act (`str`, *optional*, defaults to `None`):
+            The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.
+        time_cond_proj_dim (`int`, *optional*, defaults to `None`):
+            The dimension of `cond_proj` layer in the timestep embedding.
+        conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.
+        conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.
+        projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when
+            `class_embed_type="projection"`. Required when `class_embed_type="projection"`.
+        class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time
+            embeddings with the class embeddings.
+        mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):
+            Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If
+            `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the
+            `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`
+            otherwise.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D", "CrossAttnUpBlock2D"]
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        center_input_sample: bool = False,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "DownBlock2D",
+        ),
+        mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+        up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
+        only_cross_attention: Union[bool, Tuple[bool]] = False,
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        layers_per_block: Union[int, Tuple[int]] = 2,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        dropout: float = 0.0,
+        act_fn: str = "silu",
+        norm_num_groups: Optional[int] = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: Union[int, Tuple[int]] = 1280,
+        transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
+        reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
+        encoder_hid_dim: Optional[int] = None,
+        encoder_hid_dim_type: Optional[str] = None,
+        attention_head_dim: Union[int, Tuple[int]] = 8,
+        num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+        dual_cross_attention: bool = False,
+        use_linear_projection: bool = False,
+        class_embed_type: Optional[str] = None,
+        addition_embed_type: Optional[str] = None,
+        addition_time_embed_dim: Optional[int] = None,
+        num_class_embeds: Optional[int] = None,
+        upcast_attention: bool = False,
+        resnet_time_scale_shift: str = "default",
+        resnet_skip_time_act: bool = False,
+        resnet_out_scale_factor: float = 1.0,
+        time_embedding_type: str = "positional",
+        time_embedding_dim: Optional[int] = None,
+        time_embedding_act_fn: Optional[str] = None,
+        timestep_post_act: Optional[str] = None,
+        time_cond_proj_dim: Optional[int] = None,
+        conv_in_kernel: int = 3,
+        conv_out_kernel: int = 3,
+        projection_class_embeddings_input_dim: Optional[int] = None,
+        attention_type: str = "default",
+        class_embeddings_concat: bool = False,
+        mid_block_only_cross_attention: Optional[bool] = None,
+        cross_attention_norm: Optional[str] = None,
+        addition_embed_type_num_heads: int = 64,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+
+        if num_attention_heads is not None:
+            raise ValueError(
+                "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."
+            )
+
+        # If `num_attention_heads` is not defined (which is the case for most models)
+        # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
+        # The reason for this behavior is to correct for incorrectly named variables that were introduced
+        # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
+        # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
+        # which is why we correct for the naming here.
+        num_attention_heads = num_attention_heads or attention_head_dim
+
+        # Check inputs
+        self._check_config(
+            down_block_types=down_block_types,
+            up_block_types=up_block_types,
+            only_cross_attention=only_cross_attention,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            cross_attention_dim=cross_attention_dim,
+            transformer_layers_per_block=transformer_layers_per_block,
+            reverse_transformer_layers_per_block=reverse_transformer_layers_per_block,
+            attention_head_dim=attention_head_dim,
+            num_attention_heads=num_attention_heads,
+        )
+
+        # input
+        conv_in_padding = (conv_in_kernel - 1) // 2
+        self.conv_in = nn.Conv2d(
+            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
+        )
+
+        # time
+        time_embed_dim, timestep_input_dim = self._set_time_proj(
+            time_embedding_type,
+            block_out_channels=block_out_channels,
+            flip_sin_to_cos=flip_sin_to_cos,
+            freq_shift=freq_shift,
+            time_embedding_dim=time_embedding_dim,
+        )
+
+        self.time_embedding = TimestepEmbedding(
+            timestep_input_dim,
+            time_embed_dim,
+            act_fn=act_fn,
+            post_act_fn=timestep_post_act,
+            cond_proj_dim=time_cond_proj_dim,
+        )
+
+        self._set_encoder_hid_proj(
+            encoder_hid_dim_type,
+            cross_attention_dim=cross_attention_dim,
+            encoder_hid_dim=encoder_hid_dim,
+        )
+
+        # class embedding
+        self._set_class_embedding(
+            class_embed_type,
+            act_fn=act_fn,
+            num_class_embeds=num_class_embeds,
+            projection_class_embeddings_input_dim=projection_class_embeddings_input_dim,
+            time_embed_dim=time_embed_dim,
+            timestep_input_dim=timestep_input_dim,
+        )
+
+        self._set_add_embedding(
+            addition_embed_type,
+            addition_embed_type_num_heads=addition_embed_type_num_heads,
+            addition_time_embed_dim=addition_time_embed_dim,
+            cross_attention_dim=cross_attention_dim,
+            encoder_hid_dim=encoder_hid_dim,
+            flip_sin_to_cos=flip_sin_to_cos,
+            freq_shift=freq_shift,
+            projection_class_embeddings_input_dim=projection_class_embeddings_input_dim,
+            time_embed_dim=time_embed_dim,
+        )
+
+        if time_embedding_act_fn is None:
+            self.time_embed_act = None
+        else:
+            self.time_embed_act = get_activation(time_embedding_act_fn)
+
+        self.down_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(only_cross_attention, bool):
+            if mid_block_only_cross_attention is None:
+                mid_block_only_cross_attention = only_cross_attention
+
+            only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+        if mid_block_only_cross_attention is None:
+            mid_block_only_cross_attention = False
+
+        if isinstance(num_attention_heads, int):
+            num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        if isinstance(cross_attention_dim, int):
+            cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
+
+        if isinstance(layers_per_block, int):
+            layers_per_block = [layers_per_block] * len(down_block_types)
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+        if class_embeddings_concat:
+            # The time embeddings are concatenated with the class embeddings. The dimension of the
+            # time embeddings passed to the down, middle, and up blocks is twice the dimension of the
+            # regular time embeddings
+            blocks_time_embed_dim = time_embed_dim * 2
+        else:
+            blocks_time_embed_dim = time_embed_dim
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block[i],
+                transformer_layers_per_block=transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim[i],
+                num_attention_heads=num_attention_heads[i],
+                downsample_padding=downsample_padding,
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                attention_type=attention_type,
+                resnet_skip_time_act=resnet_skip_time_act,
+                resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
+                attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                dropout=dropout,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = get_mid_block(
+            mid_block_type,
+            temb_channels=blocks_time_embed_dim,
+            in_channels=block_out_channels[-1],
+            resnet_eps=norm_eps,
+            resnet_act_fn=act_fn,
+            resnet_groups=norm_num_groups,
+            output_scale_factor=mid_block_scale_factor,
+            transformer_layers_per_block=transformer_layers_per_block[-1],
+            num_attention_heads=num_attention_heads[-1],
+            cross_attention_dim=cross_attention_dim[-1],
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            mid_block_only_cross_attention=mid_block_only_cross_attention,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            attention_type=attention_type,
+            resnet_skip_time_act=resnet_skip_time_act,
+            cross_attention_norm=cross_attention_norm,
+            attention_head_dim=attention_head_dim[-1],
+            dropout=dropout,
+        )
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_num_attention_heads = list(reversed(num_attention_heads))
+        reversed_layers_per_block = list(reversed(layers_per_block))
+        reversed_cross_attention_dim = list(reversed(cross_attention_dim))
+        reversed_transformer_layers_per_block = (
+            list(reversed(transformer_layers_per_block))
+            if reverse_transformer_layers_per_block is None
+            else reverse_transformer_layers_per_block
+        )
+        only_cross_attention = list(reversed(only_cross_attention))
+
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            up_block = get_up_block(
+                up_block_type,
+                num_layers=reversed_layers_per_block[i] + 1,
+                transformer_layers_per_block=reversed_transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resolution_idx=i,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=reversed_cross_attention_dim[i],
+                num_attention_heads=reversed_num_attention_heads[i],
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                attention_type=attention_type,
+                resnet_skip_time_act=resnet_skip_time_act,
+                resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
+                attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+                dropout=dropout,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        if norm_num_groups is not None:
+            self.conv_norm_out = nn.GroupNorm(
+                num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
+            )
+
+            self.conv_act = get_activation(act_fn)
+
+        else:
+            self.conv_norm_out = None
+            self.conv_act = None
+
+        conv_out_padding = (conv_out_kernel - 1) // 2
+        self.conv_out = nn.Conv2d(
+            block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
+        )
+
+        self._set_pos_net_if_use_gligen(attention_type=attention_type, cross_attention_dim=cross_attention_dim)
+
+    def _check_config(
+        self,
+        down_block_types: Tuple[str],
+        up_block_types: Tuple[str],
+        only_cross_attention: Union[bool, Tuple[bool]],
+        block_out_channels: Tuple[int],
+        layers_per_block: Union[int, Tuple[int]],
+        cross_attention_dim: Union[int, Tuple[int]],
+        transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple[int]]],
+        reverse_transformer_layers_per_block: bool,
+        attention_head_dim: int,
+        num_attention_heads: Optional[Union[int, Tuple[int]]],
+    ):
+        if len(down_block_types) != len(up_block_types):
+            raise ValueError(
+                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+            )
+
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
+            )
+        if isinstance(transformer_layers_per_block, list) and reverse_transformer_layers_per_block is None:
+            for layer_number_per_block in transformer_layers_per_block:
+                if isinstance(layer_number_per_block, list):
+                    raise ValueError("Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet.")
+
+    def _set_time_proj(
+        self,
+        time_embedding_type: str,
+        block_out_channels: int,
+        flip_sin_to_cos: bool,
+        freq_shift: float,
+        time_embedding_dim: int,
+    ) -> Tuple[int, int]:
+        if time_embedding_type == "fourier":
+            time_embed_dim = time_embedding_dim or block_out_channels[0] * 2
+            if time_embed_dim % 2 != 0:
+                raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")
+            self.time_proj = GaussianFourierProjection(
+                time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos
+            )
+            timestep_input_dim = time_embed_dim
+        elif time_embedding_type == "positional":
+            time_embed_dim = time_embedding_dim or block_out_channels[0] * 4
+
+            self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+            timestep_input_dim = block_out_channels[0]
+        else:
+            raise ValueError(
+                f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."
+            )
+
+        return time_embed_dim, timestep_input_dim
+
+    def _set_encoder_hid_proj(
+        self,
+        encoder_hid_dim_type: Optional[str],
+        cross_attention_dim: Union[int, Tuple[int]],
+        encoder_hid_dim: Optional[int],
+    ):
+        if encoder_hid_dim_type is None and encoder_hid_dim is not None:
+            encoder_hid_dim_type = "text_proj"
+            self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
+            logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
+
+        if encoder_hid_dim is None and encoder_hid_dim_type is not None:
+            raise ValueError(
+                f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
+            )
+
+        if encoder_hid_dim_type == "text_proj":
+            self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
+        elif encoder_hid_dim_type == "text_image_proj":
+            # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+            # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+            # case when `addition_embed_type == "text_image_proj"` (Kandinsky 2.1)`
+            self.encoder_hid_proj = TextImageProjection(
+                text_embed_dim=encoder_hid_dim,
+                image_embed_dim=cross_attention_dim,
+                cross_attention_dim=cross_attention_dim,
+            )
+        elif encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2
+            self.encoder_hid_proj = ImageProjection(
+                image_embed_dim=encoder_hid_dim,
+                cross_attention_dim=cross_attention_dim,
+            )
+        elif encoder_hid_dim_type is not None:
+            raise ValueError(
+                f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
+            )
+        else:
+            self.encoder_hid_proj = None
+
+    def _set_class_embedding(
+        self,
+        class_embed_type: Optional[str],
+        act_fn: str,
+        num_class_embeds: Optional[int],
+        projection_class_embeddings_input_dim: Optional[int],
+        time_embed_dim: int,
+        timestep_input_dim: int,
+    ):
+        if class_embed_type is None and num_class_embeds is not None:
+            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
+        elif class_embed_type == "timestep":
+            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn)
+        elif class_embed_type == "identity":
+            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+        elif class_embed_type == "projection":
+            if projection_class_embeddings_input_dim is None:
+                raise ValueError(
+                    "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
+                )
+            # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
+            # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
+            # 2. it projects from an arbitrary input dimension.
+            #
+            # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
+            # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
+            # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
+            self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        elif class_embed_type == "simple_projection":
+            if projection_class_embeddings_input_dim is None:
+                raise ValueError(
+                    "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"
+                )
+            self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)
+        else:
+            self.class_embedding = None
+
+    def _set_add_embedding(
+        self,
+        addition_embed_type: str,
+        addition_embed_type_num_heads: int,
+        addition_time_embed_dim: Optional[int],
+        flip_sin_to_cos: bool,
+        freq_shift: float,
+        cross_attention_dim: Optional[int],
+        encoder_hid_dim: Optional[int],
+        projection_class_embeddings_input_dim: Optional[int],
+        time_embed_dim: int,
+    ):
+        if addition_embed_type == "text":
+            if encoder_hid_dim is not None:
+                text_time_embedding_from_dim = encoder_hid_dim
+            else:
+                text_time_embedding_from_dim = cross_attention_dim
+
+            self.add_embedding = TextTimeEmbedding(
+                text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
+            )
+        elif addition_embed_type == "text_image":
+            # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+            # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+            # case when `addition_embed_type == "text_image"` (Kandinsky 2.1)`
+            self.add_embedding = TextImageTimeEmbedding(
+                text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
+            )
+        elif addition_embed_type == "text_time":
+            self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
+            self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        elif addition_embed_type == "image":
+            # Kandinsky 2.2
+            self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
+        elif addition_embed_type == "image_hint":
+            # Kandinsky 2.2 ControlNet
+            self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
+        elif addition_embed_type is not None:
+            raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
+
+    def _set_pos_net_if_use_gligen(self, attention_type: str, cross_attention_dim: int):
+        if attention_type in ["gated", "gated-text-image"]:
+            positive_len = 768
+            if isinstance(cross_attention_dim, int):
+                positive_len = cross_attention_dim
+            elif isinstance(cross_attention_dim, (list, tuple)):
+                positive_len = cross_attention_dim[0]
+
+            feature_type = "text-only" if attention_type == "gated" else "text-image"
+            self.position_net = GLIGENTextBoundingboxProjection(
+                positive_len=positive_len, out_dim=cross_attention_dim, feature_type=feature_type
+            )
+
+    @property
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnAddedKVProcessor()
+        elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def set_attention_slice(self, slice_size: Union[str, int, List[int]] = "auto"):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module splits the input tensor in slices to compute attention in
+        several steps. This is useful for saving some memory in exchange for a small decrease in speed.
+
+        Args:
+            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
+                `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
+                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                must be a multiple of `slice_size`.
+        """
+        sliceable_head_dims = []
+
+        def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
+            if hasattr(module, "set_attention_slice"):
+                sliceable_head_dims.append(module.sliceable_head_dim)
+
+            for child in module.children():
+                fn_recursive_retrieve_sliceable_dims(child)
+
+        # retrieve number of attention layers
+        for module in self.children():
+            fn_recursive_retrieve_sliceable_dims(module)
+
+        num_sliceable_layers = len(sliceable_head_dims)
+
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = [dim // 2 for dim in sliceable_head_dims]
+        elif slice_size == "max":
+            # make smallest slice possible
+            slice_size = num_sliceable_layers * [1]
+
+        slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
+
+        if len(slice_size) != len(sliceable_head_dims):
+            raise ValueError(
+                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+            )
+
+        for i in range(len(slice_size)):
+            size = slice_size[i]
+            dim = sliceable_head_dims[i]
+            if size is not None and size > dim:
+                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_attention_slice method
+        # gets the message
+        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+            if hasattr(module, "set_attention_slice"):
+                module.set_attention_slice(slice_size.pop())
+
+            for child in module.children():
+                fn_recursive_set_attention_slice(child, slice_size)
+
+        reversed_slice_size = list(reversed(slice_size))
+        for module in self.children():
+            fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def enable_freeu(self, s1: float, s2: float, b1: float, b2: float):
+        r"""Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.
+
+        The suffixes after the scaling factors represent the stage blocks where they are being applied.
+
+        Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that
+        are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.
+
+        Args:
+            s1 (`float`):
+                Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to
+                mitigate the "oversmoothing effect" in the enhanced denoising process.
+            s2 (`float`):
+                Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to
+                mitigate the "oversmoothing effect" in the enhanced denoising process.
+            b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
+            b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
+        """
+        for i, upsample_block in enumerate(self.up_blocks):
+            setattr(upsample_block, "s1", s1)
+            setattr(upsample_block, "s2", s2)
+            setattr(upsample_block, "b1", b1)
+            setattr(upsample_block, "b2", b2)
+
+    def disable_freeu(self):
+        """Disables the FreeU mechanism."""
+        freeu_keys = {"s1", "s2", "b1", "b2"}
+        for i, upsample_block in enumerate(self.up_blocks):
+            for k in freeu_keys:
+                if hasattr(upsample_block, k) or getattr(upsample_block, k, None) is not None:
+                    setattr(upsample_block, k, None)
+
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedAttnProcessor2_0())
+
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+
+    def get_time_embed(
+        self, sample: torch.Tensor, timestep: Union[torch.Tensor, float, int]
+    ) -> Optional[torch.Tensor]:
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+        # `Timesteps` does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+        return t_emb
+
+    def get_class_embed(self, sample: torch.Tensor, class_labels: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        class_emb = None
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+                # `Timesteps` does not contain any weights and will always return f32 tensors
+                # there might be better ways to encapsulate this.
+                class_labels = class_labels.to(dtype=sample.dtype)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
+        return class_emb
+
+    def get_aug_embed(
+        self, emb: torch.Tensor, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any]
+    ) -> Optional[torch.Tensor]:
+        aug_emb = None
+        if self.config.addition_embed_type == "text":
+            aug_emb = self.add_embedding(encoder_hidden_states)
+        elif self.config.addition_embed_type == "text_image":
+            # Kandinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+
+            image_embs = added_cond_kwargs.get("image_embeds")
+            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
+            aug_emb = self.add_embedding(text_embs, image_embs)
+        elif self.config.addition_embed_type == "text_time":
+            # SDXL - style
+            if "text_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                )
+            text_embeds = added_cond_kwargs.get("text_embeds")
+            if "time_ids" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                )
+            time_ids = added_cond_kwargs.get("time_ids")
+            time_embeds = self.add_time_proj(time_ids.flatten())
+            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+            add_embeds = add_embeds.to(emb.dtype)
+            aug_emb = self.add_embedding(add_embeds)
+        elif self.config.addition_embed_type == "image":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            aug_emb = self.add_embedding(image_embs)
+        elif self.config.addition_embed_type == "image_hint":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            hint = added_cond_kwargs.get("hint")
+            aug_emb = self.add_embedding(image_embs, hint)
+        return aug_emb
+
+    def process_encoder_hidden_states(
+        self, encoder_hidden_states: torch.Tensor, added_cond_kwargs: Dict[str, Any]
+    ) -> torch.Tensor:
+        if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
+            # Kandinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(image_embeds)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "ip_image_proj":
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+
+            if hasattr(self, "text_encoder_hid_proj") and self.text_encoder_hid_proj is not None:
+                encoder_hidden_states = self.text_encoder_hid_proj(encoder_hidden_states)
+
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            image_embeds = self.encoder_hid_proj(image_embeds)
+            encoder_hidden_states = (encoder_hidden_states, image_embeds)
+        return encoder_hidden_states
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        use_up_blocks: bool = False,
+        return_dict: bool = True,
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        r"""
+        The [`UNet2DConditionModel`] forward method.
+
+        Args:
+            sample (`torch.Tensor`):
+                The noisy input tensor with the following shape `(batch, channel, height, width)`.
+            timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input.
+            encoder_hidden_states (`torch.Tensor`):
+                The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
+            class_labels (`torch.Tensor`, *optional*, defaults to `None`):
+                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
+            timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
+                Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
+                through the `self.time_embedding` layer to obtain the timestep embeddings.
+            attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            added_cond_kwargs: (`dict`, *optional*):
+                A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
+                are passed along to the UNet blocks.
+            down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
+                A tuple of tensors that if specified are added to the residuals of down unet blocks.
+            mid_block_additional_residual: (`torch.Tensor`, *optional*):
+                A tensor that if specified is added to the residual of the middle unet block.
+            down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
+                additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)
+            encoder_attention_mask (`torch.Tensor`):
+                A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
+                `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
+                which adds large negative values to the attention scores corresponding to "discard" tokens.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
+                If `return_dict` is True, an [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is the sample tensor.
+        """
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+        
+        # import time
+        # torch.cuda.synchronize()
+        # start_time = time.time()
+
+        for dim in sample.shape[-2:]:
+            if dim % default_overall_up_factor != 0:
+                # Forward upsample size to force interpolation output size.
+                forward_upsample_size = True
+                break
+
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        t_emb = self.get_time_embed(sample=sample, timestep=timestep)
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        class_emb = self.get_class_embed(sample=sample, class_labels=class_labels)
+        if class_emb is not None:
+            if self.config.class_embeddings_concat:
+                emb = torch.cat([emb, class_emb], dim=-1)
+            else:
+                emb = emb + class_emb
+
+        aug_emb = self.get_aug_embed(
+            emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+        )
+        if self.config.addition_embed_type == "image_hint":
+            aug_emb, hint = aug_emb
+            sample = torch.cat([sample, hint], dim=1)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        if self.time_embed_act is not None:
+            emb = self.time_embed_act(emb)
+
+        encoder_hidden_states = self.process_encoder_hidden_states(
+            encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs
+        )
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 2.5 GLIGEN position net
+        if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            gligen_args = cross_attention_kwargs.pop("gligen")
+            cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)}
+
+        # 3. down
+        # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated
+        # to the internal blocks and will raise deprecation warnings. this will be confusing for our users.
+        if cross_attention_kwargs is not None:
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            lora_scale = cross_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+
+        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+        # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
+        is_adapter = down_intrablock_additional_residuals is not None
+        # maintain backward compatibility for legacy usage, where
+        #       T2I-Adapter and ControlNet both use down_block_additional_residuals arg
+        #       but can only use one or the other
+        if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None:
+            deprecate(
+                "T2I should not use down_block_additional_residuals",
+                "1.3.0",
+                "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
+                       and will be removed in diffusers 1.3.0.  `down_block_additional_residuals` should only be used \
+                       for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
+                standard_warn=False,
+            )
+            down_intrablock_additional_residuals = down_block_additional_residuals
+            is_adapter = True
+        
+        # torch.cuda.synchronize()
+        # logger.info(f"unet preprocess: {time.time() - start_time}")
+        
+        # torch.cuda.synchronize()
+        # start_time = time.time()
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0)
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    sample += down_intrablock_additional_residuals.pop(0)
+
+            down_block_res_samples += res_samples
+
+        if is_controlnet:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+        # torch.cuda.synchronize()
+        # logger.info(f"unet down time: {time.time() - start_time}")
+        # torch.cuda.synchronize()
+        # start_time = time.time()
+        # 4. mid
+        if self.mid_block is not None:
+            if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
+                sample = self.mid_block(
+                    sample,
+                    emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
+                sample = self.mid_block(sample, emb)
+
+            # To support T2I-Adapter-XL
+            if (
+                is_adapter
+                and len(down_intrablock_additional_residuals) > 0
+                and sample.shape == down_intrablock_additional_residuals[0].shape
+            ):
+                sample += down_intrablock_additional_residuals.pop(0)
+
+        if is_controlnet:
+            sample = sample + mid_block_additional_residual
+        # torch.cuda.synchronize()
+        # logger.info(f"unet mid time: {time.time() - start_time}")
+        mid_sample = sample
+        
+        if use_up_blocks:
+            # 5. up
+            up_block_res_samples = ()
+            for i, upsample_block in enumerate(self.up_blocks):
+                is_final_block = i == len(self.up_blocks) - 1
+
+                res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+                down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+                # if we have not reached the final block and need to forward the
+                # upsample size, we do it here
+                if not is_final_block and forward_upsample_size:
+                    upsample_size = down_block_res_samples[-1].shape[2:]
+
+                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    sample = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        encoder_hidden_states=encoder_hidden_states,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        upsample_size=upsample_size,
+                        attention_mask=attention_mask,
+                        encoder_attention_mask=encoder_attention_mask,
+                    )
+                else:
+                    sample = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        upsample_size=upsample_size,
+                    )
+                up_block_res_samples += (sample, )
+
+        # # 6. post-process
+        # if self.conv_norm_out:
+        #     sample = self.conv_norm_out(sample)
+        #     sample = self.conv_act(sample)
+        # sample = self.conv_out(sample)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            if use_up_blocks:
+                return (mid_sample, down_block_res_samples, up_block_res_samples)
+            else:
+                return (mid_sample, down_block_res_samples)
+
+        return UNet2DConditionOutput(sample=sample)

Reward_sd15_idealized/pipelines/__init__.py

@@ -0,0 +1,9 @@
+from .sd15_reward_pipeline import StableDiffusionRewardPipeline
+from .sd15_gradient_ascent_pipeline import StableDiffusionGradientAscentPipeline
+
+__all__ = [
+    'StableDiffusionRewardPipeline',
+    'StableDiffusionGradientAscentPipeline',
+    'StableDiffusionXLRewardPipeline',
+    'StableDiffusionXLGradientAscentPipeline',
+]

Reward_sd15_idealized/pipelines/sd15_gradient_ascent_pipeline.py

@@ -0,0 +1,345 @@
+"""
+Stable Diffusion Pipeline with Gradient Ascent Reward Guidance
+
+This pipeline extends the reward-tracking pipeline to support gradient ascent
+on the reward score during generation for improved preference alignment.
+"""
+
+import torch
+from typing import Optional, Union, List, Callable, Tuple
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from pipelines.sd15_reward_pipeline import StableDiffusionRewardPipeline
+from gradient_ascent_utils import RewardGuidedDiffusion
+
+
+class StableDiffusionGradientAscentPipeline(StableDiffusionRewardPipeline):
+    """
+    SD Pipeline with gradient ascent reward guidance.
+    
+    This pipeline performs gradient ascent on the LRM reward score at specified
+    timesteps to guide generation toward higher preference scores.
+    
+    Usage:
+        from diffusers import StableDiffusionPipeline
+        base = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+        pipeline = StableDiffusionGradientAscentPipeline(**base.components)
+        pipeline.set_reward_model(reward_model)
+        
+        # Enable gradient ascent
+        pipeline.enable_gradient_ascent(
+            grad_timestep_range=(500, 700),
+            num_grad_steps=5,
+            grad_step_size=0.1
+        )
+        
+        output = pipeline(
+            prompt="a beautiful landscape",
+            num_inference_steps=20,
+        )
+    """
+    
+    def __init__(
+        self,
+        vae,
+        text_encoder,
+        tokenizer,
+        unet,
+        scheduler,
+        safety_checker=None,
+        feature_extractor=None,
+        image_encoder=None,
+        requires_safety_checker: bool = False,
+    ):
+        super().__init__(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+            image_encoder=image_encoder,
+            requires_safety_checker=requires_safety_checker,
+        )
+        self.gradient_ascent_enabled = False
+        self.grad_guidance = None
+    
+    def enable_gradient_ascent(
+        self,
+        grad_timestep_range: Tuple[int, int] = (500, 700),
+        grad_scale: float = 1.0,
+        num_grad_steps: int = 5,
+        grad_step_size: float = 0.1,
+        lr_scheduler_type: str = "constant",
+        lr_scheduler_kwargs: Optional[dict] = None,
+        use_momentum: bool = False,
+        momentum: float = 0.9,
+        use_nesterov: bool = False,
+        use_iso_projection: bool = False
+    ):
+        """
+        Enable gradient ascent reward guidance.
+        
+        Args:
+            grad_timestep_range: Tuple of (min_t, max_t) for applying gradients
+            grad_scale: Scale factor for gradient magnitude
+            num_grad_steps: Number of gradient ascent iterations per timestep
+            grad_step_size: Step size for each gradient update (initial LR)
+            lr_scheduler_type: Type of LR scheduler ("constant", "linear", "cosine", "exponential", "step")
+            lr_scheduler_kwargs: Additional kwargs for LR scheduler
+            use_momentum: Whether to use momentum
+            momentum: Momentum coefficient (typically 0.9)
+            use_nesterov: Whether to use Nesterov momentum
+            use_iso_projection: Whether to use Iso Projection
+        """
+        if self.reward_model is None:
+            raise ValueError("Reward model must be set before enabling gradient ascent. "
+                           "Call set_reward_model() first.")
+        
+        self.grad_guidance = RewardGuidedDiffusion(
+            reward_model=self.reward_model,
+            grad_scale=grad_scale,
+            grad_timestep_range=grad_timestep_range,
+            num_grad_steps=num_grad_steps,
+            grad_step_size=grad_step_size,
+            lr_scheduler_type=lr_scheduler_type,
+            lr_scheduler_kwargs=lr_scheduler_kwargs or {},
+            use_momentum=use_momentum,
+            momentum=momentum,
+            use_nesterov=use_nesterov,
+            use_iso_projection=use_iso_projection
+        )
+        self.gradient_ascent_enabled = True
+        print(f"✓ Gradient ascent enabled for timesteps {grad_timestep_range}")
+        print(f"  - Grad steps: {num_grad_steps}")
+        print(f"  - Grad step size (initial LR): {grad_step_size}")
+        print(f"  - Grad scale: {grad_scale}")
+        print(f"  - LR scheduler: {lr_scheduler_type}")
+        if use_momentum:
+            print(f"  - Momentum: {momentum} (Nesterov: {use_nesterov})")
+        else:
+            print(f"  - Momentum: disabled")
+        print(f"  - Iso Projection: {use_iso_projection}")
+    
+    def disable_gradient_ascent(self):
+        """Disable gradient ascent."""
+        self.gradient_ascent_enabled = False
+        self.grad_guidance = None
+        print("✓ Gradient ascent disabled")
+    
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        track_rewards: bool = True,
+        print_rewards: bool = True,
+        apply_gradient_ascent: bool = True,
+        verbose_grad: bool = True,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+        callback_steps: int = 1,
+        **kwargs,
+    ):
+        """
+        Generate images with optional gradient ascent guidance.
+        
+        Args:
+            prompt: Text prompt(s)
+            height: Image height
+            width: Image width
+            num_inference_steps: Number of denoising steps
+            guidance_scale: CFG scale
+            negative_prompt: Negative prompt(s)
+            num_images_per_prompt: Number of images per prompt
+            eta: DDIM eta parameter
+            generator: Random generator
+            latents: Initial latents
+            track_rewards: Whether to track reward scores
+            print_rewards: Whether to print rewards during generation
+            apply_gradient_ascent: Whether to apply gradient ascent (if enabled)
+            verbose_grad: Whether to print gradient ascent details
+            output_type: Output format
+            return_dict: Whether to return dict
+            callback: Callback function
+            callback_steps: Callback frequency
+        
+        Returns:
+            StableDiffusionPipelineOutput with images and reward scores
+        """
+        # Prepare inputs
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+        
+        # Encode prompt
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        
+        text_embeddings = self._encode_prompt(
+            prompt,
+            self.device,
+            num_images_per_prompt,
+            guidance_scale > 1.0,
+            negative_prompt,
+        )
+        
+        # Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+        timesteps = self.scheduler.timesteps
+        
+        # Reset gradient ascent scheduler for new generation
+        if self.grad_guidance is not None:
+            self.grad_guidance.reset_statistics()
+        
+        # Prepare latents
+        latents = self._prepare_latents(
+            batch_size * num_images_per_prompt,
+            self.unet.config.in_channels,
+            height,
+            width,
+            text_embeddings.dtype,
+            self.device,
+            generator,
+            latents,
+        )
+        
+        # Prepare extra kwargs
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        
+        # Reset statistics
+        self.reward_history = []
+        if self.gradient_ascent_enabled and self.grad_guidance:
+            self.grad_guidance.reset_statistics()
+        
+        # Print header
+        if print_rewards or verbose_grad:
+            print("\n" + "="*60)
+            print(f"Starting generation: '{prompt[0]}'")
+            if self.gradient_ascent_enabled and apply_gradient_ascent:
+                print(f"Gradient ascent: ENABLED")
+            print("="*60)
+        
+        # Denoising loop
+        for i, t in enumerate(self.progress_bar(timesteps)):
+            # Apply gradient ascent BEFORE scheduler step
+            if (self.gradient_ascent_enabled and 
+                apply_gradient_ascent and 
+                self.grad_guidance and 
+                self.grad_guidance.should_apply_gradient(t.item())):
+                
+                if verbose_grad:
+                    print(f"\n🎯 Applying gradient ascent at step {i}/{num_inference_steps}, t={t.item()}")
+                
+                # Enable gradients temporarily
+                with torch.enable_grad():
+                    latents, grad_stats = self.grad_guidance.apply_gradient_ascent(
+                        latents,
+                        prompt[0],
+                        t.item(),
+                        base_noise=None,
+                        verbose=verbose_grad,
+                        total_denoising_steps=num_inference_steps,
+                    )
+                
+            latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1.0 else latents
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            
+            noise_pred = self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+            ).sample
+            
+            # Perform CFG
+            if guidance_scale > 1.0:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+            # Compute reward score (after potential gradient update)
+            if track_rewards and self.reward_model is not None:
+                try:
+                    with torch.no_grad():
+                        timestep_value = t.item() if isinstance(t, torch.Tensor) else t
+                        score = self.reward_model.get_reward_score(
+                            latents,
+                            prompt[0],
+                            timestep_value
+                        )
+                    
+                    score_value = score.item() if score.numel() == 1 else score.mean().item()
+                    self.reward_history.append({
+                        'step': i,
+                        'timestep': timestep_value,
+                        'reward_score': score_value
+                    })
+                    
+                    if print_rewards:
+                        print(f"Step {i:3d}/{num_inference_steps} | t={timestep_value:4d} | Reward: {score_value:7.4f}")
+                
+                except Exception as e:
+                    print(f"Warning: Could not compute reward at step {i}: {e}")
+            
+            # Scheduler step
+            latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+            
+            # Callback
+            if callback is not None and i % callback_steps == 0:
+                callback(i, t, latents)
+        
+        # Print footer
+        if print_rewards or verbose_grad:
+            print("="*60)
+            if self.reward_history:
+                final_reward = self.reward_history[-1]['reward_score']
+                print(f"Generation complete! Final reward: {final_reward:.4f}")
+            
+            if self.gradient_ascent_enabled and self.grad_guidance:
+                stats = self.grad_guidance.get_statistics()
+                if stats:
+                    print(f"\nGradient Ascent Statistics:")
+                    print(f"  Applications: {stats['num_applications']}")
+                    print(f"  Total reward improvement: {stats['total_reward_improvement']:+.4f}")
+                    print(f"  Avg reward improvement: {stats['avg_reward_improvement']:+.4f}")
+                    print(f"  Avg gradient norm: {stats['avg_grad_norm']:.4f}")
+            print("="*60 + "\n")
+        
+        # Decode latents
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        
+        # Convert to PIL
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+        
+        # Create output
+        if not return_dict:
+            return (image,)
+        
+        return StableDiffusionPipelineOutput(
+            images=image,
+            nsfw_content_detected=None,
+        )
+    
+    def _prepare_latents(self, batch_size, num_channels, height, width, dtype, device, generator, latents):
+        """Prepare initial latents."""
+        shape = (batch_size, num_channels, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        
+        if latents is None:
+            latents = torch.randn(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+        
+        # Scale initial latents
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents

Reward_sd15_idealized/pipelines/sd15_reward_pipeline.py

@@ -0,0 +1,316 @@
+"""
+Modified Stable Diffusion Pipeline with Reward Tracking
+
+This pipeline extends the standard SD pipeline to compute and display
+reward scores at each denoising step using the LRM reward model.
+"""
+
+import torch
+from typing import Optional, Union, List, Callable
+from diffusers import StableDiffusionPipeline, DDIMScheduler
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+import numpy as np
+from PIL import Image
+
+
+class StableDiffusionRewardPipeline(StableDiffusionPipeline):
+    """
+    Extended Stable Diffusion Pipeline with LRM Reward Tracking
+    
+    This pipeline computes preference scores at each denoising timestep and optionally
+    displays them during generation.
+    
+    Usage:
+        from diffusers import StableDiffusionPipeline
+        base = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+        pipeline = StableDiffusionRewardPipeline(**base.components)
+        pipeline.set_reward_model(reward_model)
+        
+        output = pipeline(
+            prompt="a beautiful landscape",
+            num_inference_steps=20,
+            track_rewards=True
+        )
+    """
+    
+    def __init__(
+        self,
+        vae,
+        text_encoder,
+        tokenizer,
+        unet,
+        scheduler,
+        safety_checker=None,
+        feature_extractor=None,
+        image_encoder=None,
+        requires_safety_checker: bool = False,
+    ):
+        super().__init__(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+            image_encoder=image_encoder,
+            requires_safety_checker=requires_safety_checker,
+        )
+        self.reward_model = None
+        self.reward_history = []
+    
+    def set_reward_model(self, reward_model):
+        """Set the LRM reward model for score tracking"""
+        self.reward_model = reward_model
+        self.reward_model.eval()  # Set to eval mode
+        print("✓ Reward model attached to pipeline")
+    
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        track_rewards: bool = True,
+        print_rewards: bool = True,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+        callback_steps: int = 1,
+        **kwargs,
+    ):
+        """
+        Generate images with optional reward tracking
+        
+        Args:
+            prompt: Text prompt(s) for generation
+            height: Image height (default: 512)
+            width: Image width (default: 512)
+            num_inference_steps: Number of denoising steps
+            guidance_scale: CFG scale
+            negative_prompt: Negative prompt(s)
+            num_images_per_prompt: Number of images per prompt
+            eta: DDIM eta parameter (0 = deterministic)
+            generator: Random generator for reproducibility
+            latents: Initial latent codes (optional)
+            track_rewards: Whether to compute reward scores
+            print_rewards: Whether to print rewards during generation
+            output_type: "pil", "np", or "latent"
+            return_dict: Whether to return a dict or tuple
+            callback: Custom callback function
+            callback_steps: Frequency of callback calls
+            
+        Returns:
+            StableDiffusionPipelineOutput with images and optional reward history
+        """
+        # Default height and width to unet config
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+        
+        # Reset reward history
+        self.reward_history = []
+        
+        # 1. Check inputs
+        if isinstance(prompt, str):
+            batch_size = 1
+            prompt = [prompt]
+        elif isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            raise ValueError("Prompt must be a string or list of strings")
+        
+        # 2. Encode input prompt
+        text_embeddings = self._encode_prompt(
+            prompt,
+            self.device,
+            num_images_per_prompt,
+            guidance_scale > 1.0,
+            negative_prompt,
+        )
+        
+        # 3. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+        timesteps = self.scheduler.timesteps
+        
+        # 4. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            text_embeddings.dtype,
+            self.device,
+            generator,
+            latents,
+        )
+        
+        # 5. Prepare extra step kwargs
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        
+        # 6. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        
+        if print_rewards and track_rewards and self.reward_model is not None:
+            print("\n" + "="*60)
+            print(f"Starting generation: '{prompt[0]}'")
+            print("="*60)
+        
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # Expand latents if doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1.0 else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                
+                # Predict noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=text_embeddings,
+                ).sample
+                
+                # Perform guidance
+                if guidance_scale > 1.0:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+                
+                # Compute reward score BEFORE stepping
+                if track_rewards and self.reward_model is not None:
+                    try:
+                        # Get reward score for current latents at timestep t
+                        with torch.no_grad():
+                            # Convert timestep to proper format
+                            timestep_value = t.item() if isinstance(t, torch.Tensor) else t
+                            score = self.reward_model.get_reward_score(
+                                latents, 
+                                prompt[0],  # Use first prompt
+                                timestep_value
+                            )
+                        
+                        score_value = score.item() if score.numel() == 1 else score.mean().item()
+                        self.reward_history.append({
+                            'step': i,
+                            'timestep': timestep_value,
+                            'reward_score': score_value
+                        })
+                        
+                        if print_rewards:
+                            print(f"Step {i:3d}/{num_inference_steps} | t={t:4d} | Reward: {score_value:7.4f}")
+                    
+                    except Exception as e:
+                        if print_rewards:
+                            print(f"Warning: Could not compute reward at step {i}: {e}")
+                
+                # Compute previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+                
+                # Call callback
+                if callback is not None and i % callback_steps == 0:
+                    callback(i, t, latents)
+                
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+        
+        if print_rewards and track_rewards and self.reward_model is not None:
+            print("="*60)
+            final_score = self.reward_history[-1]['reward_score'] if self.reward_history else 0.0
+            print(f"Generation complete! Final reward: {final_score:.4f}")
+            print("="*60 + "\n")
+        
+        # 7. Decode latents
+        if output_type == "latent":
+            image = latents
+            has_nsfw_concept = None
+        elif output_type in ["pil", "np"]:
+            # Decode with VAE
+            latents = 1 / self.vae.config.scaling_factor * latents
+            image = self.vae.decode(latents).sample
+            
+            # Convert to PIL
+            image = (image / 2 + 0.5).clamp(0, 1)
+            image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+            
+            if output_type == "pil":
+                image = self.numpy_to_pil(image)
+            
+            # Run safety checker
+            has_nsfw_concept = None
+            if hasattr(self, 'safety_checker') and self.safety_checker is not None:
+                safety_checker_input = self.feature_extractor(
+                    self.numpy_to_pil(image) if isinstance(image, np.ndarray) else image,
+                    return_tensors="pt"
+                ).to(self.device)
+                image, has_nsfw_concept = self.safety_checker(
+                    images=image,
+                    clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype)
+                )
+        
+        # 8. Return output
+        if not return_dict:
+            return (image, has_nsfw_concept)
+        
+        output = StableDiffusionPipelineOutput(
+            images=image,
+            nsfw_content_detected=has_nsfw_concept
+        )
+        
+        # Attach reward history to output
+        if track_rewards:
+            output.reward_history = self.reward_history
+        
+        return output
+    
+    def get_reward_history(self):
+        """Get the reward history from the last generation"""
+        return self.reward_history
+    
+    def plot_reward_curve(self, save_path=None):
+        """
+        Plot the reward curve from the last generation
+        
+        Args:
+            save_path: Optional path to save the plot
+        """
+        if not self.reward_history:
+            print("No reward history available. Run generation with track_rewards=True first.")
+            return
+        
+        import matplotlib.pyplot as plt
+        
+        steps = [h['step'] for h in self.reward_history]
+        timesteps = [h['timestep'] for h in self.reward_history]
+        scores = [h['reward_score'] for h in self.reward_history]
+        
+        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
+        
+        # Plot reward vs step
+        ax1.plot(steps, scores, marker='o', linewidth=2, markersize=4)
+        ax1.set_xlabel('Denoising Step')
+        ax1.set_ylabel('Reward Score')
+        ax1.set_title('Reward Score vs Denoising Step')
+        ax1.grid(True, alpha=0.3)
+        
+        # Plot reward vs timestep
+        ax2.plot(timesteps, scores, marker='o', linewidth=2, markersize=4, color='orange')
+        ax2.set_xlabel('Timestep')
+        ax2.set_ylabel('Reward Score')
+        ax2.set_title('Reward Score vs Timestep')
+        ax2.grid(True, alpha=0.3)
+        ax2.invert_xaxis()  # Higher timesteps first (left to right: noise -> clean)
+        
+        plt.tight_layout()
+        
+        if save_path:
+            plt.savefig(save_path, dpi=150, bbox_inches='tight')
+            print(f"Plot saved to {save_path}")
+        
+        plt.show()
+        
+        return fig

Reward_sd15_idealized/tune_hyperparams.py

@@ -0,0 +1,514 @@
+"""
+Hyperparameter tuning script for gradient ascent optimization.
+
+This script performs a systematic search over hyperparameter combinations
+to find the optimal configuration for maximum evaluation scores.
+"""
+
+import subprocess
+import json
+import argparse
+from pathlib import Path
+from datetime import datetime
+import itertools
+import numpy as np
+from typing import Dict, List, Any
+import re
+
+
+class HyperparameterTuner:
+    """Hyperparameter tuner for gradient ascent."""
+    
+    def __init__(
+        self,
+        output_dir: str = "tuning_results",
+        max_samples: int = 30,
+        num_steps: int = 20,
+        dataset_type: str = "pickapic",
+        model_variant: str = "lpo",
+        cuda_id: int = 0,
+        metrics: List[str] = None
+    ):
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.max_samples = max_samples
+        self.num_steps = num_steps
+        self.dataset_type = dataset_type
+        self.model_variant = model_variant
+        self.cuda_id = cuda_id
+        self.metrics = metrics or ["clip", "aesthetic", "pickscore", "hpsv2", "imagereward"]
+        
+        # Store results
+        self.results = []
+        self.baseline_results = None
+        
+    def define_search_space(self) -> List[Dict[str, Any]]:
+        """Define the hyperparameter search space - FULL GRID SEARCH.
+        
+        Tests all combinations of parameters including momentum overrides for configs that support it.
+        """
+        
+        # Define all parameter values
+        cfg_scales = [3.0, 5.0, 7.5] #  
+        
+        # All available gradient configs from grad_ascent_configs.py
+        grad_configs = [
+            # "constant",
+            # "linear",
+            "cosine_nesterov",
+            # "low_to_high_nesterov",
+            # "high_to_low_nesterov",
+            "low_to_high_momentum",
+            "high_to_low_momentum",
+        ]
+        
+        num_grad_steps_list = [1, 2] # 5, 7, 10
+        grad_step_sizes = [0.001, 0.005, 0.01, 0.05] # 
+        momentums = [0.5, 0.8, 0.9]  # 
+        
+        # Generate ALL combinations using itertools.product
+        configs = []
+        for cfg, grad_cfg, num_steps, step_size, momentum in itertools.product(
+            cfg_scales, grad_configs, num_grad_steps_list, grad_step_sizes, momentums
+        ):
+            configs.append({
+                "cfg_scale": cfg,
+                "grad_config": grad_cfg,
+                "num_grad_steps": num_steps,
+                "grad_step_size": step_size,
+                "momentum": momentum,
+            })
+        
+        print(f"\nGenerated {len(configs)} total configurations")
+        print(f"  cfg_scales: {len(cfg_scales)}")
+        print(f"  grad_configs: {len(grad_configs)}")
+        print(f"  num_grad_steps: {len(num_grad_steps_list)}")
+        print(f"  grad_step_sizes: {len(grad_step_sizes)}")
+        print(f"  momentums: {len(momentums)}")
+        print(f"  Total: {len(cfg_scales)} × {len(grad_configs)} × {len(num_grad_steps_list)} × {len(grad_step_sizes)} × {len(momentums)} = {len(configs)}")
+        
+        return configs
+    
+    def run_baseline(self) -> Dict[str, float]:
+        """Run baseline evaluation once."""
+        print("\n" + "="*80)
+        print("RUNNING BASELINE EVALUATION")
+        print("="*80)
+        
+        # Use median cfg_scale for baseline
+        cfg_scale = 5.0
+        
+        output_dir = self.output_dir / "baseline"
+        
+        cmd = [
+            "python", "eval.py",
+            "--model_variant", self.model_variant,
+            "--dataset_type", self.dataset_type,
+            "--max_samples", str(self.max_samples),
+            "--num_steps", str(self.num_steps),
+            "--cfg_scale", str(cfg_scale),
+            "--output_dir", str(output_dir),
+            "--cuda", str(self.cuda_id),
+            "--mode", "baseline",
+            "--metrics", *self.metrics,
+        ]
+        
+        print(f"Command: {' '.join(cmd)}")
+        
+        try:
+            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            
+            # Parse results from output
+            metrics = self._parse_metrics(result.stdout, "baseline")
+            
+            print(f"\nBaseline Results:")
+            for metric, value in metrics.items():
+                print(f"  {metric}: {value:.4f}")
+            
+            self.baseline_results = {
+                "cfg_scale": cfg_scale,
+                "metrics": metrics,
+            }
+            
+            return metrics
+            
+        except subprocess.CalledProcessError as e:
+            print(f"Error running baseline: {e}")
+            print(f"Stdout: {e.stdout}")
+            print(f"Stderr: {e.stderr}")
+            return {}
+    
+    def run_experiment(self, config: Dict[str, Any]) -> Dict[str, Any]:
+        """Run a single experiment with given hyperparameters."""
+        
+        # Create output directory for this config
+        config_name = f"cfg{config['cfg_scale']}_" \
+                     f"{config['grad_config']}_" \
+                     f"steps{config['num_grad_steps']}_" \
+                     f"lr{config['grad_step_size']}_" \
+                     f"mom{config['momentum']}"
+        
+        output_dir = self.output_dir / config_name
+        
+        # Build command
+        cmd = [
+            "python", "eval.py",
+            "--model_variant", self.model_variant,
+            "--dataset_type", self.dataset_type,
+            "--grad_config", config["grad_config"],
+            "--max_samples", str(self.max_samples),
+            "--num_steps", str(self.num_steps),
+            "--cfg_scale", str(config["cfg_scale"]),
+            "--output_dir", str(output_dir),
+            "--cuda", str(self.cuda_id),
+            "--mode", "gradient_ascent",
+            "--metrics", *self.metrics,
+            # Override config parameters
+            "--override_num_grad_steps", str(config["num_grad_steps"]),
+            "--override_grad_step_size", str(config["grad_step_size"]),
+            "--override_momentum", str(config["momentum"]),
+        ]
+        
+        print(f"\nRunning experiment: {config_name}")
+        print(f"Config: {config}")
+        
+        try:
+            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            
+            # Parse metrics from output
+            metrics = self._parse_metrics(result.stdout, "gradient_ascent")
+            
+            # Compute improvement over baseline
+            improvements = {}
+            if self.baseline_results:
+                baseline_metrics = self.baseline_results["metrics"]
+                for metric, value in metrics.items():
+                    if metric in baseline_metrics:
+                        baseline_val = baseline_metrics[metric]
+                        if baseline_val != 0:
+                            improvement = ((value - baseline_val) / abs(baseline_val)) * 100
+                            improvements[f"{metric}_improvement"] = improvement
+            
+            result_dict = {
+                "config": config,
+                "metrics": metrics,
+                "improvements": improvements,
+                "output_dir": str(output_dir),
+                "timestamp": datetime.now().isoformat(),
+            }
+            
+            print(f"Results:")
+            for metric, value in metrics.items():
+                print(f"  {metric}: {value:.4f}")
+            if improvements:
+                print(f"Improvements over baseline:")
+                for metric, value in improvements.items():
+                    print(f"  {metric}: {value:+.2f}%")
+            
+            return result_dict
+            
+        except subprocess.CalledProcessError as e:
+            print(f"Error running experiment: {e}")
+            print(f"Stderr: {e.stderr}")
+            return {
+                "config": config,
+                "error": str(e),
+                "timestamp": datetime.now().isoformat(),
+            }
+    
+    def _parse_metrics(self, output: str, mode: str) -> Dict[str, float]:
+        """Parse metrics from eval.py output."""
+        metrics = {}
+        
+        # Look for the summary section
+        lines = output.split('\n')
+        
+        # Pattern to match metric lines like "  Reward: 0.1234"
+        metric_patterns = {
+            "reward": r"Reward:\s+([-+]?\d*\.?\d+)",
+            "clip": r"CLIP Score:\s+([-+]?\d*\.?\d+)",
+            "aesthetic": r"Aesthetic Score:\s+([-+]?\d*\.?\d+)",
+            "pickscore": r"PickScore:\s+([-+]?\d*\.?\d+)",
+            "hpsv2": r"HPSv2 Score:\s+([-+]?\d*\.?\d+)",
+            "hpsv21": r"HPSv2\.1 Score:\s+([-+]?\d*\.?\d+)",
+            "imagereward": r"ImageReward:\s+([-+]?\d*\.?\d+)",
+            "fid": r"FID:\s+([-+]?\d*\.?\d+)",
+        }
+        
+        for line in lines:
+            for metric_name, pattern in metric_patterns.items():
+                match = re.search(pattern, line)
+                if match:
+                    metrics[metric_name] = float(match.group(1))
+        
+        return metrics
+    
+    def compute_aggregate_score(self, metrics: Dict[str, float]) -> float:
+        """
+        Compute aggregate score for ranking configurations.
+        
+        Uses weighted combination of metrics (higher is better for most,
+        except FID which is lower is better).
+        """
+        weights = {
+            "reward": 1.0,
+            "clip": 0.8,
+            "aesthetic": 0.8,
+            "pickscore": 1.0,
+            "hpsv2": 1.0,
+            "hpsv21": 1.0,
+            "imagereward": 1.0,
+            "fid": -0.5,  # Negative weight (lower FID is better)
+        }
+        
+        score = 0.0
+        total_weight = 0.0
+        
+        for metric, value in metrics.items():
+            if metric in weights:
+                score += weights[metric] * value
+                total_weight += abs(weights[metric])
+        
+        # Normalize by total weight
+        if total_weight > 0:
+            score /= total_weight
+        
+        return score
+    
+    def run_search(
+        self, 
+        search_type: str = "grid",
+        start_idx: int = 0,
+        end_idx: int = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Run hyperparameter search.
+        
+        Args:
+            search_type: Type of search ("grid" or "random")
+            start_idx: Starting index for experiments (for GPU distribution)
+            end_idx: Ending index for experiments (for GPU distribution)
+        """
+        all_configs = self.define_search_space()
+        
+        print("\n" + "="*80)
+        print("HYPERPARAMETER SEARCH CONFIGURATION")
+        print("="*80)
+        print(f"Dataset: {self.dataset_type}")
+        print(f"Model: {self.model_variant}")
+        print(f"Samples: {self.max_samples}")
+        print(f"Inference steps: {self.num_steps}")
+        print(f"Metrics: {', '.join(self.metrics)}")
+        
+        # Select subset of configs if indices provided
+        if search_type == "grid":
+            configs = all_configs
+        elif search_type == "random":
+            # Random sample from all configs
+            n_samples = min(50, len(all_configs))
+            indices = np.random.choice(len(all_configs), n_samples, replace=False)
+            configs = [all_configs[i] for i in indices]
+        else:
+            raise ValueError(f"Unknown search type: {search_type}")
+        
+        # Apply index slicing for GPU distribution
+        if end_idx is None:
+            end_idx = len(configs)
+        configs = configs[start_idx:end_idx]
+        
+        print(f"\nTotal configurations: {len(all_configs)}")
+        print(f"Assigned to this worker: {len(configs)} (indices {start_idx} to {end_idx})")
+        
+        # Run baseline first
+        if self.baseline_results is None:
+            self.run_baseline()
+        
+        # Run experiments
+        print("\n" + "="*80)
+        print("RUNNING EXPERIMENTS")
+        print("="*80)
+        
+        for i, config in enumerate(configs, 1):
+            print(f"\n{'='*80}")
+            print(f"Experiment {i}/{len(configs)}")
+            print(f"{'='*80}")
+            
+            result = self.run_experiment(config)
+            self.results.append(result)
+            
+            # Save intermediate results
+            self._save_results()
+        
+        return self.results
+    
+    def _generate_grid_configs(self, search_space: Dict[str, List[Any]]) -> List[Dict[str, Any]]:
+        """Generate all combinations for grid search."""
+        keys = list(search_space.keys())
+        values = list(search_space.values())
+        
+        configs = []
+        for combination in itertools.product(*values):
+            config = dict(zip(keys, combination))
+            configs.append(config)
+        
+        return configs
+    
+    def _generate_random_configs(
+        self, 
+        search_space: Dict[str, List[Any]], 
+        n_samples: int = 20
+    ) -> List[Dict[str, Any]]:
+        """Generate random configurations for random search."""
+        configs = []
+        
+        for _ in range(n_samples):
+            config = {}
+            for param, values in search_space.items():
+                config[param] = np.random.choice(values)
+            configs.append(config)
+        
+        return configs
+    
+    def _save_results(self):
+        """Save results to JSON file."""
+        results_file = self.output_dir / "tuning_results.json"
+        
+        data = {
+            "baseline": self.baseline_results,
+            "experiments": self.results,
+            "timestamp": datetime.now().isoformat(),
+            "config": {
+                "max_samples": self.max_samples,
+                "num_steps": self.num_steps,
+                "dataset_type": self.dataset_type,
+                "model_variant": self.model_variant,
+            }
+        }
+        
+        with open(results_file, 'w') as f:
+            json.dump(data, f, indent=2)
+        
+        print(f"\nResults saved to: {results_file}")
+    
+    def analyze_results(self) -> Dict[str, Any]:
+        """Analyze results and find best configuration."""
+        if not self.results:
+            print("No results to analyze!")
+            return {}
+        
+        print("\n" + "="*80)
+        print("ANALYSIS: FINDING BEST CONFIGURATION")
+        print("="*80)
+        
+        # Filter out failed experiments
+        successful_results = [r for r in self.results if "metrics" in r]
+        
+        if not successful_results:
+            print("No successful experiments!")
+            return {}
+        
+        # Compute aggregate scores
+        for result in successful_results:
+            metrics = result["metrics"]
+            result["aggregate_score"] = self.compute_aggregate_score(metrics)
+        
+        # Sort by aggregate score
+        successful_results.sort(key=lambda x: x["aggregate_score"], reverse=True)
+        
+        # Print top 5 configurations
+        print("\nTop 5 Configurations:")
+        print("="*80)
+        
+        for i, result in enumerate(successful_results[:5], 1):
+            print(f"\n#{i} - Aggregate Score: {result['aggregate_score']:.4f}")
+            print(f"Config: {result['config']}")
+            print(f"Metrics:")
+            for metric, value in result['metrics'].items():
+                print(f"  {metric}: {value:.4f}")
+            if result.get('improvements'):
+                print(f"Improvements over baseline:")
+                for metric, value in result['improvements'].items():
+                    print(f"  {metric}: {value:+.2f}%")
+        
+        # Save best config
+        best_result = successful_results[0]
+        best_config_file = self.output_dir / "best_config.json"
+        
+        with open(best_config_file, 'w') as f:
+            json.dump({
+                "config": best_result["config"],
+                "metrics": best_result["metrics"],
+                "aggregate_score": best_result["aggregate_score"],
+                "improvements": best_result.get("improvements", {}),
+            }, f, indent=2)
+        
+        print(f"\n✓ Best configuration saved to: {best_config_file}")
+        
+        return best_result
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Hyperparameter tuning for gradient ascent")
+    parser.add_argument("--output_dir", type=str, default="tuning_results",
+                        help="Directory to save tuning results")
+    parser.add_argument("--max_samples", type=int, default=30,
+                        help="Number of samples to use for tuning")
+    parser.add_argument("--num_steps", type=int, default=20,
+                        help="Number of inference steps (fixed)")
+    parser.add_argument("--dataset_type", type=str, default="pickapic",
+                        choices=["coco", "pickapic"],
+                        help="Dataset to use")
+    parser.add_argument("--model_variant", type=str, default="lpo",
+                        choices=["origin", "spo", "diffusion_dpo", "lpo"],
+                        help="Model variant to use")
+    parser.add_argument("--cuda", type=int, default=0,
+                        help="CUDA device ID")
+    parser.add_argument("--search_type", type=str, default="grid",
+                        choices=["grid", "random"],
+                        help="Type of hyperparameter search")
+    parser.add_argument("--metrics", type=str, nargs="+",
+                        default=["clip", "aesthetic", "pickscore", "hpsv2", "imagereward"],
+                        help="Metrics to evaluate")
+    parser.add_argument("--start_idx", type=int, default=0,
+                        help="Starting index for experiments (for GPU distribution)")
+    parser.add_argument("--end_idx", type=int, default=None,
+                        help="Ending index for experiments (for GPU distribution)")
+    
+    args = parser.parse_args()
+    
+    # Create tuner
+    tuner = HyperparameterTuner(
+        output_dir=args.output_dir,
+        max_samples=args.max_samples,
+        num_steps=args.num_steps,
+        dataset_type=args.dataset_type,
+        model_variant=args.model_variant,
+        cuda_id=args.cuda,
+        metrics=args.metrics,
+    )
+    
+    # Run search
+    results = tuner.run_search(
+        search_type=args.search_type,
+        start_idx=args.start_idx,
+        end_idx=args.end_idx
+    )
+    
+    # Analyze results
+    best_result = tuner.analyze_results()
+    
+    print("\n" + "="*80)
+    print("TUNING COMPLETE!")
+    print("="*80)
+    print(f"Total experiments: {len(results)}")
+    print(f"Results directory: {args.output_dir}")
+    
+    if best_result:
+        print(f"\nBest configuration:")
+        print(json.dumps(best_result["config"], indent=2))
+        print(f"\nAggregate score: {best_result['aggregate_score']:.4f}")
+
+
+if __name__ == "__main__":
+    main()

evaluation/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

evaluation/aesthetic_score.py

@@ -0,0 +1,71 @@
+"""
+Adapted from https://github.com/christophschuhmann/improved-aesthetic-predictor. Originally Apache License, Version 2.0, January 2004.
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+from transformers import CLIPModel, CLIPProcessor
+from PIL import Image
+from io import BytesIO
+
+def open_image(image):
+    if isinstance(image, bytes):
+        image = Image.open(BytesIO(image))
+    elif isinstance(image, str):
+        image = Image.open(image)
+    image = image.convert("RGB")
+    return image
+
+
+class MLP(nn.Module):
+    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):
+    def __init__(self, dtype, clip_name_or_path, aesthetic_name_or_path):
+        super().__init__()
+        self.clip = CLIPModel.from_pretrained(clip_name_or_path)
+        self.processor = CLIPProcessor.from_pretrained(clip_name_or_path)
+        self.mlp = MLP()
+        state_dict = torch.load(aesthetic_name_or_path, map_location='cpu')
+        self.mlp.load_state_dict(state_dict)
+        self.dtype = dtype
+        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)
+ 
+    
+if __name__ == "__main__":
+    clip_name_or_path = "openai/clip-vit-large-patch14"
+    aesthetic_name_or_path = "./sac+logos+ava1-l14-linearMSE.pth"
+    aesthetic_scorer = AestheticScorer(torch.float32, clip_name_or_path, aesthetic_name_or_path)
+    image0 = open_image('./image0.png')
+    image1 = open_image('./image1.png')
+    print(aesthetic_scorer(image0))
+    print(aesthetic_scorer(image1))
+    print(aesthetic_scorer([image0, image1]))
+    

evaluation/blip/__init__.py

@@ -0,0 +1 @@
+from .blip_pretrain import *

evaluation/general_score_pap_val.py

@@ -0,0 +1,265 @@
+from datasets import load_dataset
+from pick_score import PickScorer
+from aesthetic_score import AestheticScorer
+from hpsv2_score import HPSv2Scorer
+from imagereward_score import load_imagereward
+from diffusers import AutoencoderKL, StableDiffusionPipeline, \
+                        StableDiffusionXLPipeline, DDIMScheduler, \
+                        UNet2DConditionModel
+import torch
+import os
+import json
+from tqdm import tqdm
+from huggingface_hub import hf_hub_download
+from argparse import ArgumentParser
+
+
+def load_origin_sd_v1_5(scheduler, inference_dtype):
+    pipe = StableDiffusionPipeline.from_pretrained(
+        'stable-diffusion-v1-5/stable-diffusion-v1-5',
+        torch_dtype=inference_dtype,
+        scheduler=scheduler,
+        safety_checker=None,
+    )
+    guidance_scale = 7.5
+    return pipe, guidance_scale
+
+
+def load_spo_sd_v1_5(scheduler, inference_dtype):
+    pipe = StableDiffusionPipeline.from_pretrained(
+        'SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep',
+        torch_dtype=inference_dtype,
+        scheduler=scheduler,
+        safety_checker=None,
+    )
+    guidance_scale = 5.0
+    return pipe, guidance_scale
+
+
+def load_diffusion_dpo_sd_v1_5(scheduler, inference_dtype):
+    unet = UNet2DConditionModel.from_pretrained('mhdang/dpo-sd1.5-text2image-v1', subfolder="unet", torch_dtype=inference_dtype)
+    pipe = StableDiffusionPipeline.from_pretrained(
+        'stable-diffusion-v1-5/stable-diffusion-v1-5',
+        torch_dtype=inference_dtype,
+        scheduler=scheduler,
+        safety_checker=None,
+        unet=unet,
+    )
+    guidance_scale = 7.5
+    return pipe, guidance_scale
+
+
+def load_lpo_sd_v1_5(scheduler, inference_dtype):
+    unet = UNet2DConditionModel.from_pretrained(
+        'casiatao/LPO',
+        subfolder="lpo_sd15_merge/unet",
+        torch_dtype=inference_dtype
+    )
+    pipe = StableDiffusionPipeline.from_pretrained(
+        'stable-diffusion-v1-5/stable-diffusion-v1-5',
+        torch_dtype=inference_dtype,
+        scheduler=scheduler,
+        safety_checker=None,
+        unet=unet
+    )
+    guidance_scale = 5.0
+    return pipe, guidance_scale
+
+
+def load_origin_sdxl(scheduler, inference_dtype):
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        'stabilityai/stable-diffusion-xl-base-1.0',
+        torch_dtype=inference_dtype,
+        scheduler=scheduler,
+    )
+    vae = AutoencoderKL.from_pretrained(
+        'madebyollin/sdxl-vae-fp16-fix',
+        torch_dtype=torch.float16,
+    )
+    pipe.vae = vae
+    guidance_scale = 5.0
+    return pipe, guidance_scale
+
+
+def load_spo_sdxl(scheduler, inference_dtype):
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        'SPO-Diffusion-Models/SPO-SDXL_4k-p_10ep',
+        torch_dtype=inference_dtype,
+        scheduler=scheduler,
+    )
+    vae = AutoencoderKL.from_pretrained(
+        'madebyollin/sdxl-vae-fp16-fix',
+        torch_dtype=torch.float16,
+    )
+    pipe.vae = vae
+    guidance_scale = 5.0
+    return pipe, guidance_scale
+
+
+def load_diffusion_dpo_sdxl(scheduler, inference_dtype):
+    unet = UNet2DConditionModel.from_pretrained('mhdang/dpo-sdxl-text2image-v1', subfolder="unet", torch_dtype=inference_dtype)
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        'stabilityai/stable-diffusion-xl-base-1.0',
+        torch_dtype=inference_dtype,
+        scheduler=scheduler,
+        unet=unet,
+    )
+    vae = AutoencoderKL.from_pretrained(
+        'madebyollin/sdxl-vae-fp16-fix',
+        torch_dtype=torch.float16,
+    )
+    pipe.vae = vae
+    guidance_scale = 5.0
+    return pipe, guidance_scale
+
+
+def load_lpo_sdxl(scheduler, inference_dtype):
+    unet = UNet2DConditionModel.from_pretrained(
+        'casiatao/LPO',
+        subfolder="lpo_sdxl_merge/unet",
+        torch_dtype=inference_dtype
+    )
+    vae = AutoencoderKL.from_pretrained(
+        'madebyollin/sdxl-vae-fp16-fix',
+        torch_dtype=torch.float16,
+    )
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        'stabilityai/stable-diffusion-xl-base-1.0',
+        torch_dtype=inference_dtype,
+        scheduler=scheduler,
+        unet=unet,
+        vae=vae
+    )
+    guidance_scale = 5.0
+    return pipe, guidance_scale
+
+
+
+model_dict = {
+    'origin_sd15': load_origin_sd_v1_5,
+    'spo_sd15': load_spo_sd_v1_5,
+    'diffusion_dpo_sd15': load_diffusion_dpo_sd_v1_5,
+    'lpo_sd15': load_lpo_sd_v1_5,
+    'origin_sdxl': load_origin_sdxl,
+    'spo_sdxl': load_spo_sdxl,
+    'diffusion_dpo_sdxl': load_diffusion_dpo_sdxl,
+    'lpo_sdxl': load_lpo_sdxl,
+}
+
+
+if __name__ == "__main__":
+    # hyperparameter
+    parser = ArgumentParser()
+    parser.add_argument("--model_name", type=str, default="origin_sdxl")
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--num_image_per_prompt", type=int, default=4)
+    parser.add_argument("--sample_steps", type=int, default=20)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--device", type=str, default="cuda")
+    args = parser.parse_args()
+
+    model_name = args.model_name
+    batch_size = args.batch_size
+    num_image_per_prompt = args.num_image_per_prompt
+    sample_steps = args.sample_steps
+    seed = args.seed
+    device = args.device
+
+    # load preference model
+    pickscorer = PickScorer(processor_name_or_path="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", model_pretrained_name_or_path="yuvalkirstain/PickScore_v1", device=device)
+    
+    aesthetic_scorer = AestheticScorer(torch.float32, "openai/clip-vit-large-patch14", "./sac+logos+ava1-l14-linearMSE.pth")
+    aesthetic_scorer = aesthetic_scorer.to(device)
+    
+
+    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"),
+        model_pretrained_name_or_path=hf_hub_download(repo_id="xswu/HPSv2", filename="HPS_v2_compressed.pt"),
+        device=device
+    )
+
+    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"),
+        model_pretrained_name_or_path=hf_hub_download(repo_id="xswu/HPSv2", filename="HPS_v2.1_compressed.pt"),
+        device=device
+    )
+
+    imagereward_scorer = load_imagereward(
+        model_path=hf_hub_download(repo_id="THUDM/ImageReward", filename="ImageReward.pt"), 
+        med_config=hf_hub_download(repo_id="THUDM/ImageReward", filename="med_config.json"), 
+        device=device
+    )
+    
+    load_model_func = model_dict[model_name]
+
+    # load diffusion model
+    inference_dtype = torch.float16
+    scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
+    pipe, guidance_scale = load_model_func(scheduler, inference_dtype)
+    pipe.to(device)
+    pipe.set_progress_bar_config(disable=True)
+    
+    # load dataset
+    val_dataset = load_dataset("pickapic-anonymous/pickapic_v1", split="validation_unique", streaming=True)
+
+    # calculate preference score
+    caption_list = []
+    for i, sample in enumerate(val_dataset):
+        caption_list.append(sample['caption'])
+    
+    batch_num = len(caption_list) // batch_size if len(caption_list) % batch_size == 0 else len(caption_list) // batch_size + 1
+    batched_caption_list = [caption_list[i*batch_size:(i+1)*batch_size] for i in range(batch_num)]
+    
+    pickscore_list = []
+    aesthetic_score_list = []
+    hpsv2score_list = []
+    hpsv21score_list = []
+    imagereward_list = []
+    
+    for batch_prompt in tqdm(batched_caption_list):
+        generator=torch.Generator(device=device).manual_seed(seed)
+        images = pipe(
+            batch_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=sample_steps,
+            generator=generator,
+            output_type='pil',
+            num_images_per_prompt=num_image_per_prompt,
+        ).images
+        
+        for prompt, image in zip(batch_prompt, images):
+            pickscore = pickscorer(prompt, [image])[0]
+            pickscore_list.append(pickscore)
+            
+            aesthetic_score = aesthetic_scorer(image)[0].item()
+            aesthetic_score_list.append(aesthetic_score)
+            
+            hpsv2_score = hpsv2_scorer.score(image, prompt)[0]
+            hpsv2score_list.append(hpsv2_score)
+            
+            hpsv21_score = hpsv21_scorer.score(image, prompt)[0]
+            hpsv21score_list.append(hpsv21_score)
+            
+            imagereward_score = imagereward_scorer.score(prompt, image)
+            imagereward_list.append(imagereward_score)
+    
+    
+    res_save_dir = './eval_results/pick_a_pic_val_score'
+    os.makedirs(res_save_dir, exist_ok=True)
+
+    file_name = f"{model_name}_ddim_cfg{guidance_scale}_step{sample_steps}_seed{seed}_{num_image_per_prompt}image_batch{batch_size}.json"
+    with open(os.path.join(res_save_dir, file_name), 'w', encoding='utf-8') as f:
+        json.dump({
+            'pickscore': torch.mean(torch.tensor(pickscore_list)).item(),
+            'aestheticscore': torch.mean(torch.tensor(aesthetic_score_list)).item(),
+            'hpsv2score': torch.mean(torch.tensor(hpsv2score_list)).item(),
+            'hpsv21score': torch.mean(torch.tensor(hpsv21score_list)).item(),
+            'imagerewardscore': torch.mean(torch.tensor(imagereward_list)).item(),
+        }, f, indent=4)
+    
+    print(f"Pickscore: {torch.mean(torch.tensor(pickscore_list))}")
+    print(f"Aestheticscore: {torch.mean(torch.tensor(aesthetic_score_list))}")
+    print(f"HPSv2score: {torch.mean(torch.tensor(hpsv2score_list))}")
+    print(f"HPSv21score: {torch.mean(torch.tensor(hpsv21score_list))}")
+    print(f"Imagerewardscore: {torch.mean(torch.tensor(imagereward_list))}")
+        

upload.py

@@ -0,0 +1,270 @@
+#!/usr/bin/env python3
+"""Upload this codebase to Hugging Face Hub while excluding the .back folder.
+
+Examples:
+  /g/data/rr81/aev/bin/python upload.py
+  /g/data/rr81/aev/bin/python upload.py --repo-id aryadomain/all_code_base
+  /g/data/rr81/aev/bin/python upload.py --method large
+  /g/data/rr81/aev/bin/python upload.py --repo-id my-user/all_code_base
+"""
+
+from __future__ import annotations
+
+import argparse
+import inspect
+from pathlib import Path
+from typing import List
+
+from huggingface_hub import HfApi
+
+
+def build_ignore_patterns(extra_ignore: List[str]) -> List[str]:
+    # Always exclude local metadata and the requested backup folder.
+    patterns = [
+        ".git",
+        ".git/**",
+        ".back",
+        ".back/**",
+        "**/.back/**",
+    ]
+    patterns.extend(extra_ignore)
+    return patterns
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Upload a folder to Hugging Face Hub, excluding .back."
+    )
+    parser.add_argument(
+        "--source-dir",
+        type=str,
+        default=".",
+        help="Local folder to upload (default: current directory).",
+    )
+    parser.add_argument(
+        "--repo-name",
+        type=str,
+        default="all_code_base",
+        help="Repo name used when --repo-id is not provided.",
+    )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default="aryadomain/all_code_base",
+        help="Full Hugging Face repo id like user_or_org/repo_name (default: aryadomain/all_code_base).",
+    )
+    parser.add_argument(
+        "--namespace",
+        type=str,
+        default=None,
+        help="Optional user/org namespace override when using --repo-name.",
+    )
+    parser.add_argument(
+        "--repo-type",
+        type=str,
+        default="model",
+        choices=["model", "dataset", "space"],
+        help="Hub repo type.",
+    )
+    parser.add_argument(
+        "--private",
+        action="store_true",
+        help="Create the repo as private.",
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default="main",
+        help="Target branch/revision (default: main).",
+    )
+    parser.add_argument(
+        "--commit-message",
+        type=str,
+        default="Upload codebase excluding .back",
+        help="Commit message for upload.",
+    )
+    parser.add_argument(
+        "--extra-ignore",
+        nargs="*",
+        default=[],
+        help="Additional ignore patterns for upload_folder.",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print resolved settings and exit without uploading.",
+    )
+    parser.add_argument(
+        "--method",
+        type=str,
+        default="large",
+        choices=["auto", "folder", "large"],
+        help=(
+            "Upload method: auto chooses large upload for big folders, "
+            "folder forces upload_folder, large forces upload_large_folder compatibility path."
+        ),
+    )
+    parser.add_argument(
+        "--large-threshold-gb",
+        type=float,
+        default=10.0,
+        help="In auto mode, switch to large upload when folder size exceeds this many GB.",
+    )
+    return parser.parse_args()
+
+
+def resolve_repo_id(api: HfApi, args: argparse.Namespace) -> str:
+    if args.repo_id:
+        return args.repo_id
+
+    if args.namespace:
+        return f"{args.namespace}/{args.repo_name}"
+
+    who = api.whoami()
+    username = who.get("name")
+    if not username:
+        raise RuntimeError("Could not resolve username from Hugging Face login.")
+    return f"{username}/{args.repo_name}"
+
+
+def folder_size_bytes(folder: Path) -> int:
+    total = 0
+    for p in folder.rglob("*"):
+        if p.is_file() and ".back" not in p.parts:
+            try:
+                total += p.stat().st_size
+            except OSError:
+                pass
+    return total
+
+
+def upload_folder_compat(
+    api: HfApi,
+    repo_id: str,
+    repo_type: str,
+    source_dir: Path,
+    revision: str,
+    commit_message: str,
+    ignore_patterns: List[str],
+):
+    kwargs = {
+        "repo_id": repo_id,
+        "repo_type": repo_type,
+        "folder_path": str(source_dir),
+        "path_in_repo": ".",
+        "revision": revision,
+        "commit_message": commit_message,
+        "ignore_patterns": ignore_patterns,
+    }
+    sig = inspect.signature(api.upload_folder)
+    if "multi_commits" in sig.parameters:
+        kwargs["multi_commits"] = True
+    if "multi_commits_verbose" in sig.parameters:
+        kwargs["multi_commits_verbose"] = True
+    return api.upload_folder(**kwargs)
+
+
+def upload_large_compat(
+    api: HfApi,
+    repo_id: str,
+    repo_type: str,
+    source_dir: Path,
+    revision: str,
+    ignore_patterns: List[str],
+):
+    if not hasattr(api, "upload_large_folder"):
+        raise RuntimeError("Installed huggingface_hub does not provide upload_large_folder")
+
+    sig = inspect.signature(api.upload_large_folder)
+    kwargs = {
+        "repo_id": repo_id,
+        "repo_type": repo_type,
+        "folder_path": str(source_dir),
+    }
+    if "revision" in sig.parameters:
+        kwargs["revision"] = revision
+    if "path_in_repo" in sig.parameters:
+        kwargs["path_in_repo"] = "."
+    if "ignore_patterns" in sig.parameters:
+        kwargs["ignore_patterns"] = ignore_patterns
+
+    return api.upload_large_folder(**kwargs)
+
+
+def main() -> None:
+    args = parse_args()
+    source_dir = Path(args.source_dir).resolve()
+
+    if not source_dir.exists() or not source_dir.is_dir():
+        raise FileNotFoundError(f"Source directory not found or not a directory: {source_dir}")
+
+    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)
+
+    if args.method == "auto":
+        use_large = total_size_gb >= args.large_threshold_gb
+    else:
+        use_large = args.method == "large"
+
+    print("Source directory:", source_dir)
+    print("Repo id:", repo_id)
+    print("Repo type:", args.repo_type)
+    print("Private:", args.private)
+    print("Revision:", args.revision)
+    print("Ignore patterns:", ignore_patterns)
+    print(f"Folder size (excluding .back): {total_size_gb:.2f} GB")
+    print("Upload method:", "large" if use_large else "folder")
+
+    if args.dry_run:
+        print("Dry run requested. Exiting before create/upload.")
+        return
+
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type=args.repo_type,
+        private=args.private,
+        exist_ok=True,
+    )
+
+    if use_large:
+        try:
+            commit_info = upload_large_compat(
+                api=api,
+                repo_id=repo_id,
+                repo_type=args.repo_type,
+                source_dir=source_dir,
+                revision=args.revision,
+                ignore_patterns=ignore_patterns,
+            )
+        except Exception as e:
+            print(f"Large upload path failed ({e}). Falling back to upload_folder with multi-commit mode.")
+            commit_info = upload_folder_compat(
+                api=api,
+                repo_id=repo_id,
+                repo_type=args.repo_type,
+                source_dir=source_dir,
+                revision=args.revision,
+                commit_message=args.commit_message,
+                ignore_patterns=ignore_patterns,
+            )
+    else:
+        commit_info = upload_folder_compat(
+            api=api,
+            repo_id=repo_id,
+            repo_type=args.repo_type,
+            source_dir=source_dir,
+            revision=args.revision,
+            commit_message=args.commit_message,
+            ignore_patterns=ignore_patterns,
+        )
+
+    print("Upload completed.")
+    print("Commit:", commit_info)
+    print("Repo URL: https://huggingface.co/" + repo_id)
+
+
+if __name__ == "__main__":
+    main()