# SANA LRM Implementation Checklist

## A. Documentation First
- [x] Create `docs/plan.md`.
- [x] Create `docs/architecture.md`.
- [x] Create `docs/checklist.md`.
- [x] Create `docs/migration_notes.md`.

## B. Code Scaffold
- [x] Ensure `lrm_sana/setup.py` exists.
- [x] Ensure `lrm_sana/train_lrm_sana.sh` exists.
- [x] Ensure full `lrm_sana/trainer/*` structure exists.

## C. Config Registration
- [x] Add SANA config registration in `trainer/models/__init__.py`.
- [x] Add SANA config registration in `trainer/datasets/__init__.py`.
- [x] Add SANA config registration in `trainer/tasks/__init__.py`.
- [x] Add SANA config registration in `trainer/criterions/__init__.py`.
- [x] Add SANA trainer config dataclass in `trainer/configs/step_sana_configs.py`.
- [x] Add `trainer/conf/step_sana_base.yaml`.

## D. Model Integration
- [x] Implement `trainer/models/sana_preference_model.py`.
- [x] Implement variant-aware checkpoint selection support.
- [x] Confirm latent encode/noise/transformer/pooling path.
- [x] Confirm text/image projection + `logit_scale` logic.

## E. Dataset/Task/Criterion
- [x] Implement `trainer/datasets/step_sana_hf_dataset.py`.
- [x] Implement `trainer/tasks/step_sana_task.py`.
- [x] Implement `trainer/criterions/step_clip_criterion_sana.py`.
- [x] Preserve pseudo-preference filtering and timestep behavior.

## F. Runtime Launcher
- [x] Adapt `train_lrm_sana.sh` from flux launcher.
- [x] Keep `RUN_PROFILE=main|quick` flow.
- [x] Keep offline cache and DeepSpeed launch behavior.
- [x] Add model profile switch for 4 SANA checkpoints.

## G. Validation
- [x] Run static import checks.
- [ ] Run hydra config composition check.
- [ ] Run quick smoke training (`max_steps=1`).
- [ ] Verify logs/checkpoints/config snapshot output.
- [ ] Verify each SANA profile initializes cleanly.