Add files using upload-large-folder tool

Reward_sana_idealized/open_clip/model_configs/RN50.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            6,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

Reward_sana_idealized/open_clip/model_configs/RN50x64.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 448,
+        "layers": [
+            3,
+            15,
+            36,
+            10
+        ],
+        "width": 128,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 12
+    }
+}

Reward_sana_idealized/open_clip/model_configs/ViT-B-32-plus-256.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 256,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

Reward_sana_idealized/open_clip/model_configs/ViT-B-32-quickgelu.json

@@ -0,0 +1,17 @@
+{
+    "embed_dim": 512,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

Reward_sana_idealized/open_clip/model_configs/xlm-roberta-base-ViT-B-32.json

@@ -0,0 +1,15 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "hf_model_name": "xlm-roberta-base",
+        "hf_tokenizer_name": "xlm-roberta-base",
+        "proj": "mlp",
+        "pooler_type": "mean_pooler"
+    }
+}

Reward_sdxl_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_sdxl_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_sdxl_idealized/eval.py

@@ -0,0 +1,1143 @@
+"""
+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 pathlib import Path
+from PIL import Image
+from diffusers import StableDiffusionPipeline, DDIMScheduler, UNet2DConditionModel, StableDiffusionXLPipeline
+from models.reward_model import LRMRewardModelXL
+from pipelines.sdxl_gradient_ascent_pipeline import StableDiffusionXLGradientAscentPipeline
+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 pick_score import PickScorer
+from hpsv2_score import HPSv2Scorer
+from imagereward_score import load_imagereward
+from huggingface_hub import hf_hub_download
+
+import random
+
+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"):
+    """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 dataset (streaming)...")
+        val_dataset = load_dataset("pickapic-anonymous/pickapic_v1", split="validation_unique", streaming=True)
+        
+        prompts = []
+        for i, sample in enumerate(val_dataset):
+            prompts.append(sample['caption'])
+            if max_samples and i + 1 >= max_samples:
+                break
+        
+        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
+):
+    """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
+    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)
+        
+        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,
+            )
+        
+        # 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:
+            # 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(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
+
+
+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="stabilityai/stable-diffusion-xl-base-1.0", help="Base model path")
+    parser.add_argument("--model_variant", type=str, default="origin", 
+                        choices=["spo", "lpo"],
+                        help="SDXL model variant to use (default: origin)")
+    parser.add_argument("--lrm_model", type=str, default="casiatao/LRM", help="LRM model path")
+    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)")
+
+    # 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()
+
+    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
+    dtype = torch.bfloat16
+    
+    # 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"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)
+    
+    # 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 = LRMRewardModelXL(
+        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()
+    elif dtype == torch.bfloat16:
+        reward_model = reward_model.to(torch.bfloat16)
+    reward_model.eval()
+    print("✓ Reward model loaded")
+    
+    # Load pipeline
+    print("\n" + "="*70)
+    print("3. LOADING PIPELINE")
+    print("="*70)
+    
+    # Load model based on variant
+    if args.model_variant == "spo":
+        base_pipeline = StableDiffusionXLPipeline.from_pretrained(
+            'SPO-Diffusion-Models/SPO-SDXL_4k-p_10ep',
+            torch_dtype=dtype,
+            safety_checker=None,
+        )
+        args.cfg_scale = 5.0  # SPO uses CFG 5.0
+        print(f"✓ Loaded SPO SDXL model (cfg_scale adjusted to 5.0)")
+    elif args.model_variant == "lpo":
+        unet = UNet2DConditionModel.from_pretrained(
+            'casiatao/LPO',
+            subfolder="lpo_sdxl_merge/unet",
+            torch_dtype=dtype
+        )
+        base_pipeline = StableDiffusionXLPipeline.from_pretrained(
+            args.base_model,
+            torch_dtype=dtype,
+            variant="fp16",
+            unet=unet
+        )
+        args.cfg_scale = 5.0  # LPO uses CFG 5.0
+        print(f"✓ Loaded LPO SDXL model (cfg_scale adjusted to 5.0)")
+    
+    pipeline = StableDiffusionXLGradientAscentPipeline(**base_pipeline.components)
+    pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
+    pipeline = pipeline.to(device)
+    pipeline.set_reward_model(reward_model)
+    print("✓ Pipeline loaded")
+    
+    # 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(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=dtype, device=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:
+            pick_scorer = PickScorer(
+                processor_name_or_path="laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
+                model_pretrained_name_or_path="yuvalkirstain/PickScore_v1",
+                device=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:
+            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
+            )
+            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:
+            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
+            )
+            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:
+            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
+            )
+            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, _ = 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)
+        )
+        
+        # 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 = 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)
+        )
+
+        # 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}")
+    
+    # ========== 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_sdxl_idealized/examples.sh

@@ -0,0 +1,19 @@
+export HF_HOME=/efs/vkvermaa/2025/diffusion/latent_correct/hf_cache
+Dataset_Name="pickapic" # "coco" or "pickapic"
+Grad_Config="one_step_rectification_config"
+# constant cosine_nesterov low_to_high_momentum  high_to_low_momentum
+# low_to_high_nesterov high_to_low_nesterov
+Model_Variant="spo" #lpo, spo
+GPU_ID=$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits | sort -k2 -n | head -n1 | cut -d',' -f1)
+
+python eval.py \
+    --model_variant "$Model_Variant" \
+    --dataset_type "$Dataset_Name" \
+    --grad_config "$Grad_Config" \
+    --metrics fid clip aesthetic pickscore hpsv2 hpsv21 imagereward \
+    --max_samples 500 \
+    --num_steps 20 \
+    --cfg_scale 3 \
+    --output_dir "RESULTS/$Dataset_Name/${Grad_Config}_${Model_Variant}" \
+    --cuda $GPU_ID \
+    --mode "gradient_ascent"  # gradient_ascent baseline  # /baseline

Reward_sdxl_idealized/gradient_ascent_utils.py

@@ -0,0 +1,339 @@
+"""
+Gradient Ascent utilities for reward-guided diffusion generation.
+
+This module implements gradient ascent on the LRM reward score to guide
+the diffusion process toward higher preference scores.
+"""
+
+import torch
+import torch.nn.functional as F
+from typing import Optional, Tuple, List, Literal
+from tqdm import tqdm
+from lr_scheduler import create_lr_scheduler, LRScheduler
+
+
+class RewardGuidedDiffusion:
+    """
+    Implements reward-guided generation using gradient ascent.
+    
+    During denoising, at specified timesteps, we:
+    1. Compute the reward score for current latents
+    2. Calculate gradients of reward w.r.t. latents
+    3. Update latents in the direction that increases reward
+    
+    This guides generation toward higher preference scores.
+    """
+    
+    def __init__(
+        self,
+        reward_model,
+        grad_scale: float = 1.0,
+        grad_timestep_range: Optional[Tuple[int, int]] = None,
+        num_grad_steps: int = 5,
+        grad_step_size: float = 0.1,
+        gradient_checkpoint: bool = False,
+        # LR Scheduling
+        lr_scheduler_type: Literal["constant", "linear", "cosine", "exponential", "step"] = "constant",
+        lr_scheduler_kwargs: Optional[dict] = None,
+        # Momentum
+        use_momentum: bool = False,
+        momentum: float = 0.9,
+        use_nesterov: bool = False,
+        use_iso_projection: bool = False
+    ):
+        """
+        Initialize reward-guided diffusion.
+        
+        Args:
+            reward_model: LRM reward model for computing preference scores
+            grad_scale: Scale factor for gradient updates (default: 1.0)
+            grad_timestep_range: Tuple of (min_t, max_t) for gradient ascent.
+                                If None, applies to all timesteps.
+            num_grad_steps: Number of gradient ascent steps per timestep
+            grad_step_size: Step size for each gradient update (initial LR)
+            gradient_checkpoint: Whether to use gradient checkpointing
+            lr_scheduler_type: Type of LR scheduler ("constant", "linear", "cosine", "exponential", "step")
+            lr_scheduler_kwargs: Additional kwargs for LR scheduler (e.g., end_lr, min_lr, warmup_steps)
+            use_momentum: Whether to use momentum in gradient updates
+            momentum: Momentum coefficient (typically 0.9)
+            use_nesterov: Whether to use Nesterov momentum
+            use_iso_projection: Whether to use Iso Projection
+        """
+        self.reward_model = reward_model
+        self.grad_scale = grad_scale
+        self.grad_timestep_range = grad_timestep_range
+        self.num_grad_steps = num_grad_steps
+        self.grad_step_size = grad_step_size
+        self.gradient_checkpoint = gradient_checkpoint
+        
+        # LR Scheduler
+        self.lr_scheduler_type = lr_scheduler_type
+        self.lr_scheduler_kwargs = lr_scheduler_kwargs or {}
+        self.lr_scheduler: Optional[LRScheduler] = None
+        self.global_lr_scheduler: Optional[LRScheduler] = None  # Scheduler across denoising timesteps
+        
+        # Momentum
+        self.use_momentum = use_momentum
+        self.momentum = momentum
+        self.use_nesterov = use_nesterov
+        self.velocity = None  # Will be initialized per optimization
+
+        self.use_iso_projection = use_iso_projection
+        
+        # Statistics
+        self.grad_stats = []
+        self.timestep_counter = 0  # Track which timestep we're on
+    
+    def should_apply_gradient(self, timestep: int) -> bool:
+        """Check if gradient ascent should be applied at this timestep."""
+
+        if self.grad_timestep_range is None:
+            return False
+        
+        min_t, max_t = self.grad_timestep_range
+        return min_t <= timestep <= max_t
+    
+    @torch.enable_grad()
+    def compute_reward_gradient(
+        self,
+        latents: torch.Tensor,
+        prompt: str,
+        timestep: int,
+    ) -> Tuple[torch.Tensor, float]:
+        """
+        Compute gradient of reward score w.r.t. latents in FP32 to prevent underflow.
+        """
+        # 1. Cast to FP32 and ensure we are detached from previous iterations
+        latents_fp32 = latents.detach().to(torch.float32).clone()
+        latents_fp32.requires_grad_(True)
+        
+        # 2. Compute reward score
+        # Note: Even if the model internally uses fp16/bf16, autograd will 
+        # safely accumulate the gradient in fp32 for our leaf node.
+        reward_score = self.reward_model.get_reward_score(
+            latents_fp32,
+            prompt,
+            timestep,
+            enable_grad=True
+        )
+        
+        if reward_score.numel() > 1:
+            reward_score = reward_score.mean()
+            
+        # 3. Extract gradient
+        # CRITICAL: retain_graph=True prevents the graph from dying across multiple 
+        # gradient steps if your reward model relies on cached text embeddings.
+        grad = torch.autograd.grad(
+            outputs=reward_score,
+            inputs=latents_fp32,
+            create_graph=False,
+            retain_graph=True,  # Keeps the graph alive for the next step!
+            allow_unused=True,
+        )[0]
+        
+        # 4. Handle None gradients and cast back to the pipeline's original dtype
+        if grad is None:
+            grad = torch.zeros_like(latents)
+        else:
+            grad = grad.to(latents.dtype)
+            
+        return grad, reward_score.item()
+  
+    def apply_gradient_ascent(
+        self,
+        latents: torch.Tensor,
+        prompt: str,
+        timestep: int,
+        base_noise: Optional[torch.Tensor] = None, # Required for Iso-Marginal projection
+        verbose: bool = True,
+        total_denoising_steps: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        
+        # 1. UPCAST TO FP32 AND SETUP OPTIMIZER (Targeting Latents)
+        original_latents = latents.detach().clone().to(torch.float32)
+        current_latents = torch.nn.Parameter(original_latents.clone())
+        self.reward_model.unet.conv_in.weight.requires_grad_(True)
+        
+        # Initial reward tracking
+        with torch.no_grad():
+            initial_reward = self.reward_model.get_reward_score(
+                latents, 
+                prompt, 
+                timestep
+            )
+            initial_reward_val = initial_reward.item() if initial_reward.numel() == 1 else initial_reward.mean().item()
+            
+        # Initialize tracking lists
+        grad_norms = []
+        reward_history = [initial_reward_val]
+        lr_history = []
+        
+        # 2. FORWARD PASS (downcast to FP16 just for the model forward pass)
+        reward = self.reward_model.get_reward_score(
+            current_latents.to(latents.dtype),
+            prompt,
+            timestep,
+            enable_grad=True
+        )
+            
+        loss = -reward.mean()
+        loss.backward()
+
+        # Extract latent gradient
+        raw_grad = current_latents.grad
+        reward_history.append(reward.mean().item())
+        
+        # 3. ISO-MARGINAL PROJECTION WITH ASYMMETRIC INCLUSION
+        if raw_grad is not None and base_noise is not None and self.use_iso_projection:
+            gamma = 1e-8
+            B = raw_grad.shape[0]
+
+            grad_flat = raw_grad.view(B, -1)
+            noise_flat = base_noise.view(B, -1).to(torch.float32)
+
+            # Compute projection scalar for raw_grad (which is -?R)
+            dot_product = (grad_flat * noise_flat).sum(dim=1, keepdim=True)
+            noise_norm_sq = (noise_flat * noise_flat).sum(dim=1, keepdim=True)
+
+            proj_scalar = dot_product / (noise_norm_sq + gamma)
+            proj_scalar = proj_scalar.view(B, 1, 1, 1)
+
+            # 1. Decompose
+            grad_parallel = proj_scalar * base_noise.to(torch.float32)
+            grad_perp = raw_grad - grad_parallel
+
+            # 2. Asymmetric Inclusion
+            # proj_scalar > 0 means the applied step (+?R) points toward -epsilon (Denoising. GOOD.)
+            # proj_scalar < 0 means the applied step (+?R) points toward +epsilon (Noising. BAD.)
+            safe_proj_scalar = torch.clamp(proj_scalar, min=0.0)
+            
+            beta = 1.0 # Retention factor for the safe parallel gradient
+            safe_grad_parallel = beta * (safe_proj_scalar * base_noise.to(torch.float32))
+
+            # 3. Recombine
+            #grad_perp = grad_perp + safe_grad_parallel
+        else:
+            grad_perp = raw_grad
+            if base_noise is None:
+                print("?? WARNING: base_noise missing. Skipping Iso-Marginal projection.")
+
+        # 4. KINETIC RECTIFICATION (Applied to the projected latent gradient)
+        if grad_perp is not None:
+            max_grad = grad_perp.norm().item()
+            
+            if max_grad > 0:
+                kinetic_direction = grad_perp / (max_grad + 1e-8)
+                
+                # Because the max element is 1.0, alpha is the EXACT float32 change applied.
+                alpha = self.grad_step_size 
+                
+                with torch.no_grad():
+                    rectified_latents = original_latents - (alpha * kinetic_direction)
+            else:
+                print("?? WARNING: Gradient exists but max value is 0.0")
+                rectified_latents = original_latents.clone()
+                alpha = 0.0
+        else:
+            print("?? FATAL: PyTorch completely dropped the latent gradient!")
+            rectified_latents = original_latents.clone()
+            max_grad = 0.0
+            alpha = 0.0          
+
+        if verbose:
+            print(f"    Grad step | LR: {alpha:.6f} | Reward: {reward.mean().item():.4f} | Max Grad: {max_grad:.4f}")
+                
+        # 5. DOWNCAST AND RETURN
+        final_latents = rectified_latents.detach().to(latents.dtype)
+        
+        with torch.no_grad():
+            final_reward = self.reward_model.get_reward_score(
+                final_latents, prompt, timestep
+            )
+            final_reward_val = final_reward.item() if final_reward.numel() == 1 else final_reward.mean().item()
+            
+        stats = {
+            'timestep': timestep,
+            'initial_reward': initial_reward_val,
+            'final_reward': final_reward_val,
+            'reward_improvement': final_reward_val - initial_reward_val,
+            'grad_norms': [max_grad],
+            'reward_history': reward_history,
+            'lr_history': [alpha],       # Kept for plotting logic
+            'latent_change': (final_latents - original_latents.to(latents.dtype)).norm().item(),
+        }
+        
+        self.grad_stats.append(stats)
+            
+        return final_latents, stats
+
+    def get_statistics(self) -> dict:
+        """Get aggregated statistics across all gradient ascent applications."""
+        if not self.grad_stats:
+            return {}
+        
+        total_improvement = sum(s['reward_improvement'] for s in self.grad_stats)
+        avg_improvement = total_improvement / len(self.grad_stats)
+        
+        all_grad_norms = [n for s in self.grad_stats for n in s['grad_norms']]
+        
+        return {
+            'num_applications': len(self.grad_stats),
+            'total_reward_improvement': total_improvement,
+            'avg_reward_improvement': avg_improvement,
+            'avg_grad_norm': sum(all_grad_norms) / len(all_grad_norms) if all_grad_norms else 0,
+            'max_grad_norm': max(all_grad_norms) if all_grad_norms else 0,
+            'detailed_stats': self.grad_stats,
+        }
+    
+    def reset_statistics(self):
+        """Reset statistics and global scheduler."""
+        self.grad_stats = []
+        self.global_lr_scheduler = None
+        self.timestep_counter = 0
+
+
+def create_reward_guided_generator(
+    reward_model,
+    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
+) -> RewardGuidedDiffusion:
+    """
+    Convenience function to create a reward-guided diffusion generator.
+    
+    Args:
+        reward_model: LRM reward model
+        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
+        lr_scheduler_kwargs: Additional kwargs for LR scheduler
+        use_momentum: Whether to use momentum
+        momentum: Momentum coefficient
+        use_nesterov: Whether to use Nesterov momentum
+        use_iso_projection: Whether to use Iso Projection
+    
+    Returns:
+        RewardGuidedDiffusion instance
+    """
+    return RewardGuidedDiffusion(
+        reward_model=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,
+        use_momentum=use_momentum,
+        momentum=momentum,
+        use_nesterov=use_nesterov,
+        use_iso_projection= False
+    )

Reward_sdxl_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_sdxl_idealized/models/__init__.py

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

Reward_sdxl_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_sdxl_idealized/pipelines/sdxl_gradient_ascent_pipeline.py

@@ -0,0 +1,375 @@
+"""
+Stable Diffusion XL Pipeline with Gradient Ascent Reward Guidance
+
+Fully compatible with HuggingFace diffusers StableDiffusionXLPipeline.
+Injects the Kinetic Latent Predictor-Corrector operator split safely 
+before the ODE integration step.
+"""
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+from diffusers import StableDiffusionXLPipeline
+from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
+
+from gradient_ascent_utils import RewardGuidedDiffusion
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure.
+    (Matches diffusers implementation)
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+class StableDiffusionXLGradientAscentPipeline(StableDiffusionXLPipeline):
+    """
+    SDXL Pipeline with KLPC gradient ascent reward guidance.
+    """
+    
+    def __init__(
+        self,
+        vae,
+        text_encoder,
+        text_encoder_2,
+        tokenizer,
+        tokenizer_2,
+        unet,
+        scheduler,
+        image_encoder=None,
+        feature_extractor=None,
+        force_zeros_for_empty_prompt: bool = True,
+        add_watermarker: bool = None,
+    ):
+        super().__init__(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            unet=unet,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+            force_zeros_for_empty_prompt=force_zeros_for_empty_prompt,
+            add_watermarker=add_watermarker,
+        )
+        
+        # Initialize our KLPC custom state variables
+        self.gradient_ascent_enabled = False
+        self.grad_guidance = None
+        self.reward_model = None
+        self.reward_history = []
+
+    def set_reward_model(self, reward_model):
+        self.reward_model = reward_model
+
+    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
+    ):
+        if self.reward_model is None:
+            raise ValueError("Reward model must be set first via set_reward_model().")
+            
+        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
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        sigmas: List[float] = None,
+        denoising_end: Optional[float] = None,
+        guidance_scale: float = 5.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: 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,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+        ip_adapter_image: Optional[Any] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        original_size: Optional[Tuple[int, int]] = None,
+        crops_coords_top_left: Tuple[int, int] = (0, 0),
+        target_size: Optional[Tuple[int, int]] = None,
+        negative_original_size: Optional[Tuple[int, int]] = None,
+        negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
+        negative_target_size: Optional[Tuple[int, int]] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], Dict]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        # Custom parameters for KLPC gradient ascent
+        track_rewards: bool = True,
+        apply_gradient_ascent: bool = True,
+        verbose_grad: bool = False,
+        **kwargs,
+    ):
+        # 1. Setup sizes & Batch
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 2. Define batch size
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            lora_scale=lora_scale,
+            clip_skip=clip_skip,
+        )
+        
+        # Keep HF internal state synchronized just in case external callbacks need them
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. 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,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        if self.text_encoder_2 is None:
+            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+        else:
+            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+        add_time_ids = self._get_add_time_ids(
+            original_size,
+            crops_coords_top_left,
+            target_size,
+            dtype=prompt_embeds.dtype,
+            text_encoder_projection_dim=text_encoder_projection_dim,
+        )
+
+        if negative_original_size is not None and negative_target_size is not None:
+            negative_add_time_ids = self._get_add_time_ids(
+                negative_original_size,
+                negative_crops_coords_top_left,
+                negative_target_size,
+                dtype=prompt_embeds.dtype,
+                text_encoder_projection_dim=text_encoder_projection_dim,
+            )
+        else:
+            negative_add_time_ids = add_time_ids
+
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+        # 8. Reset Tracking
+        self.reward_history = []
+        if self.grad_guidance is not None:
+            self.grad_guidance.reset_statistics()
+
+        # Extract string prompt for the reward model forward pass
+        prompt_str = prompt[0] if isinstance(prompt, list) else prompt
+
+        # 9. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+
+                # predict the noise residual
+                added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+
+                # =========================================================================
+                # KINETIC LATENT PREDICTOR-CORRECTOR (KLPC) OPERATOR SPLIT
+                # =========================================================================
+                # This executes the spatial displacement C_i(z) right BEFORE the scheduler 
+                # integrates using the "stale" velocity (noise_pred), fulfilling Theorem 1.
+                if (
+                    self.gradient_ascent_enabled 
+                    and apply_gradient_ascent 
+                    and self.grad_guidance 
+                    and self.grad_guidance.should_apply_gradient(t.item())
+                    and prompt_str is not None
+                ):
+                    with torch.enable_grad():
+                        latents, grad_stats = self.grad_guidance.apply_gradient_ascent(
+                            latents,
+                            prompt_str,
+                            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=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+            
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                if do_classifier_free_guidance and guidance_rescale > 0.0:
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+                # Track Rewards
+                if track_rewards and self.reward_model is not None and prompt_str is not None:
+                    with torch.no_grad():
+                        score = self.reward_model.get_reward_score(latents, prompt_str, t.item())
+                        score_val = score.item() if score.numel() == 1 else score.mean().item()
+                        self.reward_history.append({'timestep': t.item(), 'reward_score': score_val})
+                # =========================================================================
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    latents = latents.to(latents_dtype)
+
+                # Execute diffusers native step end callback if provided
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
+                    add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        if not output_type == "latent":
+            # make sure the VAE is in float32 mode, as it overflows in float16
+            needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+            if needs_upcasting:
+                self.upcast_vae()
+                latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
+            elif latents.dtype != self.vae.dtype:
+                # HF Native fix: Cast the VAE to match the latents dtype!
+                self.vae = self.vae.to(latents.dtype)
+
+            # unscale/denormalize the latents
+            has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None
+            has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None
+            
+            if has_latents_mean and has_latents_std:
+                latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+                latents_std = torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+                latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean
+            else:
+                latents = latents / self.vae.config.scaling_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            
+            # cast back to fp16 if needed
+            if needs_upcasting:
+                self.vae.to(dtype=torch.float16)
+                
+            image = self.image_processor.postprocess(image, output_type=output_type)
+        else:
+            image = latents
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusionXLPipelineOutput(images=image)

Reward_sdxl_idealized/timestep_convergence_analysis.ipynb

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

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

Reward_sdxl_idealized/tune_parallel.sh

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

evaluation/open_clip/model_configs/RN50.json

@@ -0,0 +1,21 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": [
+            3,
+            4,
+            6,
+            3
+        ],
+        "width": 64,
+        "patch_size": null
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

evaluation/open_clip/model_configs/ViT-B-32-plus-256.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 640,
+    "vision_cfg": {
+        "image_size": 256,
+        "layers": 12,
+        "width": 896,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 640,
+        "heads": 10,
+        "layers": 12
+    }
+}

evaluation/open_clip/model_configs/ViT-S-32.json

@@ -0,0 +1,16 @@
+{
+    "embed_dim": 384,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 384,
+        "patch_size": 32
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 384,
+        "heads": 6,
+        "layers": 12
+    }
+}

evaluation/open_clip/model_configs/convnext_large_d_320.json

@@ -0,0 +1,19 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "timm_model_name": "convnext_large",
+        "timm_model_pretrained": false,
+        "timm_pool": "",
+        "timm_proj": "mlp",
+        "timm_drop": 0.0,
+        "timm_drop_path": 0.1,
+        "image_size": 320
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 16
+    }
+}

lrm/flux/.hydra/config.yaml

@@ -0,0 +1,124 @@
+accelerator:
+  _target_: trainer.accelerators.debug_accelerator.DebugAccelerator
+  output_dir: ${output_dir}
+  mixed_precision: BF16
+  gradient_accumulation_steps: 1
+  log_with: null
+  debug:
+    activate: false
+    port: 5900
+  seed: 42
+  resume_from_checkpoint: false
+  max_steps: 8000
+  num_epochs: 10
+  validate_steps: 100
+  generalization_validate_steps: 500
+  eval_on_start: true
+  project_name: reward_model
+  run_name: step_flux_schnell_variable-t_lr1e-5_step-8000_cfg0.0_filter2_time951
+  max_grad_norm: 1.0
+  save_steps: 100
+  metric_name: accuracy
+  metric_mode: MAX
+  limit_num_checkpoints: 1
+  save_only_if_best: true
+  dynamo_backend: 'NO'
+  keep_best_ckpts: true
+task:
+  limit_examples_to_wandb: 50
+  _target_: trainer.tasks.step_flux_task.StepFluxTask
+  pretrained_model_name_or_path: ${model.pretrained_model_name_or_path}
+  tokenizer_subfolder: tokenizer
+  label_0_column_name: ${dataset.label_0_column_name}
+  label_1_column_name: ${dataset.label_1_column_name}
+  input_ids_column_name: ${dataset.input_ids_column_name}
+  input_ids_2_column_name: ${dataset.input_ids_2_column_name}
+  pixels_0_column_name: ${dataset.pixels_0_column_name}
+  pixels_1_column_name: ${dataset.pixels_1_column_name}
+  timestep_column_name: ${dataset.timestep_column_name}
+  constant_timestep: ${dataset.constant_timestep}
+model:
+  _target_: trainer.models.flux_preference_model.FluxPreferenceModel
+  pretrained_model_name_or_path: black-forest-labs/FLUX.1-schnell
+  pretrained_vae_name_or_path: black-forest-labs/FLUX.1-schnell
+  projection_dim: 1024
+  text_embed_dim: 768
+  logit_scale_init_value: 2.6592
+  freeze_text_encoder: false
+  guidance_scale: 0.0
+  noise_offset: false
+  noise_offset_coeff: 0.05
+  max_sequence_length: 512
+  image_size: 1024
+criterion:
+  _target_: trainer.criterions.step_clip_criterion_flux.StepFluxCLIPCriterion
+  is_distributed: false
+  label_0_column_name: ${dataset.label_0_column_name}
+  label_1_column_name: ${dataset.label_1_column_name}
+  input_ids_column_name: ${dataset.input_ids_column_name}
+  input_ids_2_column_name: ${dataset.input_ids_2_column_name}
+  pixels_0_column_name: ${dataset.pixels_0_column_name}
+  pixels_1_column_name: ${dataset.pixels_1_column_name}
+  num_examples_per_prompt_column_name: ${dataset.num_examples_per_prompt_column_name}
+  timestep_column_name: ${dataset.timestep_column_name}
+  loss_type: pair
+  batch_coeff: 1.0
+  aux_loss_coeff: 1.0
+dataset:
+  train_split_name: train
+  valid_split_name: validation_unique
+  test_split_name: test_unique
+  batch_size: 4
+  num_workers: 2
+  drop_last: true
+  _target_: trainer.datasets.step_flux_hf_dataset.StepFluxHFDataset
+  dataset_name: pickapic-anonymous/pickapic_v1
+  dataset_config_name: null
+  from_disk: false
+  cache_dir: null
+  caption_column_name: caption
+  input_ids_column_name: input_ids
+  input_ids_2_column_name: input_ids_2
+  image_0_column_name: jpg_0
+  image_1_column_name: jpg_1
+  label_0_column_name: label_0
+  label_1_column_name: label_1
+  are_different_column_name: are_different
+  has_label_column_name: has_label
+  pixels_0_column_name: pixel_values_0
+  pixels_1_column_name: pixel_values_1
+  timestep_column_name: timestep
+  constant_timestep: 1
+  variable_timestep: true
+  largest_timestep: 951
+  compare_between_timestep: false
+  timestep_comparison_column_name: timestep_comparison
+  timestep_interval: 1
+  num_examples_per_prompt_column_name: num_example_per_prompt
+  keep_only_different: false
+  keep_only_with_label: false
+  keep_only_with_label_in_non_train: true
+  keep_only_with_pesudo_preference: true
+  pseudo_preference_path: /g/data/rr81/LPO/lrm/flux/vqa_aes_clip_score_mp.csv
+  filter_strategy: 2
+  processor:
+    pretrained_model_name_or_path: ${model.pretrained_model_name_or_path}
+    max_sequence_length: ${model.max_sequence_length}
+    image_size: ${model.image_size}
+    random_crop: false
+    no_hflip: true
+  limit_examples_per_prompt: -1
+  only_on_best: false
+optimizer:
+  _target_: trainer.optimizers.dummy_optimizer.BaseDummyOptim
+  lr: 1.0e-05
+  weight_decay: 0.3
+lr_scheduler:
+  _target_: trainer.lr_schedulers.dummy_lr_scheduler.instantiate_dummy_lr_scheduler
+  lr: ${optimizer.lr}
+  lr_warmup_steps: 1000
+  total_num_steps: ${accelerator.max_steps}
+debug:
+  activate: false
+  port: 5900
+output_dir: logs/lrm/${accelerator.project_name}/${accelerator.run_name}

lrm/flux/.hydra/hydra.yaml

@@ -0,0 +1,166 @@
+hydra:
+  run:
+    dir: .
+  sweep:
+    dir: multirun/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.mode=RUN
+    task:
+    - accelerator.mixed_precision=BF16
+    - accelerator.log_with=null
+    - accelerator=debug
+    - criterion.is_distributed=false
+    - dataset.pseudo_preference_path=/g/data/rr81/LPO/lrm/flux/vqa_aes_clip_score_mp.csv
+  job:
+    name: train
+    chdir: null
+    override_dirname: accelerator.log_with=null,accelerator.mixed_precision=BF16,accelerator=debug,criterion.is_distributed=false,dataset.pseudo_preference_path=/g/data/rr81/LPO/lrm/flux/vqa_aes_clip_score_mp.csv
+    id: ???
+    num: ???
+    config_name: step_flux_base
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /g/data/rr81/LPO/lrm/flux
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /g/data/rr81/LPO/lrm/flux/trainer/conf
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /g/data/rr81/LPO/lrm/flux
+    choices:
+      lr_scheduler: dummy
+      optimizer: dummy
+      dataset: step_flux
+      criterion: step_clip_flux
+      model: step_flux_base
+      task: step_flux
+      accelerator: debug
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

lrm/flux/.hydra/overrides.yaml

@@ -0,0 +1,5 @@
+- accelerator.mixed_precision=BF16
+- accelerator.log_with=null
+- accelerator=debug
+- criterion.is_distributed=false
+- dataset.pseudo_preference_path=/g/data/rr81/LPO/lrm/flux/vqa_aes_clip_score_mp.csv

lrm/flux/README.md

@@ -0,0 +1,4 @@
+Please check the train_flux.sh file 
+In line 21 there is a variable named "RUN_PROFILE" which can be set to "main" or "quick for testing. If set to "main", it will run the full training loop. If set to "quick", it will run a minimal training loop for testing purposes.
+
+I have set distributed shared for quick testing mode but distributed is set for main mode.