Hugging Face's logo

aryadomain
/
all_code_base

Model card Files
xet
 Community
all_code_base
File size: 5,134 Bytes
6d6dbbc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#!/usr/bin/env bash
set -euo pipefail

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

"${CMD[@]}"