Reward SANA Idealized
This folder is a SANA-only reward-guided inference package.
What is inside
models/reward_model.py- Local SANA reward wrapper (no trainer import from other directories).
- Loads base SANA diffusers modules and local reward checkpoint weights.
pipelines/sana_reward_pipeline.py- SANA pipeline with per-step reward tracking.
pipelines/sana_gradient_ascent_pipeline.py- SANA pipeline with gradient-ascent latent updates.
eval.py- End-to-end evaluation script.
examples.sh- Cluster entrypoint for prefetch and evaluation.
Default checkpoint
examples.sh defaults to:
/g/data/rr81/LPO/lrm/lrm_sana/logs/v8/reward_model/step_sana_sana_600m_512_variable-t_lr1e-5_step-8000_filter2_time951/checkpoint-gstep76000
Override with:
LRM_MODEL_PATH=/path/to/checkpoint-dir-or-model.safetensors
Run (10-sample smoke test)
cd /g/data/rr81/LPO/Reward_sana_idealized
OFFLINE_MODE=1 MAX_SAMPLES=10 MODE=gradient_ascent MODEL_PROFILE=sana_600m_512 ./examples.sh
Notes
- Uses existing Python env:
/g/data/rr81/aev/bin/python. - GPU nodes should run with offline HF cache.