Reward_sd15_idealized/examples.sh

#!/usr/bin/env bash
set -euo pipefail

clear
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Shared HF cache used on this cluster.
HF_HUB_CACHE_DIR="${HF_HUB_CACHE_DIR:-/scratch/rr81/ma5430/.cache/huggingface/hub}"
export HF_HUB_CACHE="$HF_HUB_CACHE_DIR"
export HUGGINGFACE_HUB_CACHE="$HF_HUB_CACHE_DIR"
export HF_HOME="$(dirname "$HF_HUB_CACHE_DIR")"

# GPU nodes have no internet, while login nodes do.
# Auto default: offline on GPU nodes, online on login nodes.
DEFAULT_OFFLINE_MODE="1"
if ! (command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1); then
    DEFAULT_OFFLINE_MODE="0"
fi
OFFLINE_MODE="${OFFLINE_MODE:-$DEFAULT_OFFLINE_MODE}"

if [[ "$OFFLINE_MODE" == "1" ]]; then
    export HF_DATASETS_OFFLINE="1"
    export HF_METRICS_OFFLINE="1"
    export HF_MODULES_OFFLINE="1"
    export TRANSFORMERS_OFFLINE="1"
    export DIFFUSERS_OFFLINE="1"
    export HF_HUB_OFFLINE="1"
else
    export HF_DATASETS_OFFLINE="0"
    export HF_METRICS_OFFLINE="0"
    export HF_MODULES_OFFLINE="0"
    export TRANSFORMERS_OFFLINE="0"
    export DIFFUSERS_OFFLINE="0"
    export HF_HUB_OFFLINE="0"
fi

# Existing environment requested by user.
PYTHON_BIN="${PYTHON_BIN:-/g/data/rr81/aev/bin/python}"
if [[ ! -x "$PYTHON_BIN" ]]; then
    echo "[examples.sh] Missing Python executable: $PYTHON_BIN" >&2
    exit 1
fi

DATASET_NAME="${DATASET_NAME:-pickapic}"      # coco | pickapic
GRAD_CONFIG="${GRAD_CONFIG:-one_step_rectification_config}"
MODEL_VARIANT="${MODEL_VARIANT:-spo}"      # origin | spo | diffusion_dpo | lpo
MODE="${MODE:-gradient_ascent}"               # gradient_ascent | baseline | both
# Empty MAX_SAMPLES means evaluate all available samples.
MAX_SAMPLES="${MAX_SAMPLES:-}"
NUM_STEPS="${NUM_STEPS:-20}"
CFG_SCALE="${CFG_SCALE:-3}"
METRICS="${METRICS:-clip aesthetic pickscore hpsv2 hpsv21 imagereward}"
PREFETCH_ONLY="${PREFETCH_ONLY:-0}"

# Override this path whenever you want to swap reward weights.
LRM_MODEL_PATH="${LRM_MODEL_PATH:-/g/data/rr81/LPO/lrm/lrm_15/LRM}"

if [[ -z "${GPU_ID:-}" ]]; then
    if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
        GPU_ID="$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits | sort -k2 -n | head -n1 | cut -d',' -f1 | tr -d ' ')"
        GPU_ID="${GPU_ID:-0}"
    else
        GPU_ID="0"
        echo "[examples.sh] No visible NVIDIA GPU on this node. Defaulting GPU_ID=0."
        echo "[examples.sh] eval.py will run on CPU if CUDA is unavailable."
    fi
fi

echo "Using GPU ID: $GPU_ID"
echo "Using LRM weights: $LRM_MODEL_PATH"
echo "HF offline mode: $OFFLINE_MODE"

if [[ "$PREFETCH_ONLY" == "1" ]]; then
    echo "[examples.sh] PREFETCH_ONLY=1 -> downloading required model files to shared cache and exiting."
    export MODEL_VARIANT
    export METRICS
    "$PYTHON_BIN" - <<'PY'
import os
from huggingface_hub import hf_hub_download, snapshot_download

cache_dir = os.environ["HF_HUB_CACHE"]
model_variant = os.environ.get("MODEL_VARIANT", "origin")
metrics = set(os.environ.get("METRICS", "clip aesthetic").split())

def snap(repo_id):
    print(f"[prefetch] snapshot_download: {repo_id}")
    snapshot_download(repo_id=repo_id, cache_dir=cache_dir, local_files_only=False)

def one(repo_id, filename):
    print(f"[prefetch] hf_hub_download: {repo_id}/{filename}")
    hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir, local_files_only=False)

# Base model for generation + reward backbone
snap("stable-diffusion-v1-5/stable-diffusion-v1-5")

# Variant-specific generation checkpoints (if used)
if model_variant == "spo":
    snap("SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep")
elif model_variant == "diffusion_dpo":
    snap("mhdang/dpo-sd1.5-text2image-v1")
elif model_variant == "lpo":
    snap("casiatao/LPO")

# Required for CLIP-based metrics and LRM text projection init fallback
if "clip" in metrics or "aesthetic" in metrics:
    snap("openai/clip-vit-large-patch14")

if "pickscore" in metrics:
    snap("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
    snap("yuvalkirstain/PickScore_v1")

if "hpsv2" in metrics or "hpsv21" in metrics:
    one("laion/CLIP-ViT-H-14-laion2B-s32B-b79K", "open_clip_pytorch_model.bin")
if "hpsv2" in metrics:
    one("xswu/HPSv2", "HPS_v2_compressed.pt")
if "hpsv21" in metrics:
    one("xswu/HPSv2", "HPS_v2.1_compressed.pt")

if "imagereward" in metrics:
    one("THUDM/ImageReward", "ImageReward.pt")
    one("THUDM/ImageReward", "med_config.json")

print("[prefetch] done")
PY
    exit 0
fi

read -r -a METRICS_ARR <<< "$METRICS"

CMD=(
    "$PYTHON_BIN" eval.py
    --model_variant "$MODEL_VARIANT"
    --dataset_type "$DATASET_NAME"
    --lrm_model "$LRM_MODEL_PATH"
    --grad_config "$GRAD_CONFIG"
    --metrics "${METRICS_ARR[@]}"
    --num_steps "$NUM_STEPS"
    --cfg_scale "$CFG_SCALE"
    --hf_cache_dir "$HF_HUB_CACHE_DIR"
    --output_dir "RESULTS/$DATASET_NAME/${GRAD_CONFIG}_${MODEL_VARIANT}"
    --cuda "$GPU_ID"
    --mode "$MODE"
)

if [[ -n "$MAX_SAMPLES" ]]; then
    CMD+=(--max_samples "$MAX_SAMPLES")
fi

if [[ "$OFFLINE_MODE" == "1" ]]; then
    CMD+=(--offline)
fi

"${CMD[@]}"