File size: 2,389 Bytes
ef8f3ad | 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 | # SANA LRM Architecture
## System Overview
This implementation reuses the modular trainer architecture from `flux`:
- `scripts/train.py` orchestrates load/prepare/train/eval/save.
- `accelerators/*` manage distributed runtime, logging, and checkpoints.
- `datasets/*` build pairwise preference batches with timestep controls.
- `models/*` provide latent-space text-image feature extraction.
- `criterions/*` compute pairwise preference loss.
- `tasks/*` compute inference probabilities and accuracy metrics.
## Data Flow
1. Dataset loads pairwise examples from `pickapic-anonymous/pickapic_v1`.
2. Each sample yields prompt tokens, image pair tensors, labels, and timestep pair.
3. Criterion concatenates image pairs and passes text/image/timestep to the model.
4. Model returns normalized text/image embeddings in shared projection space.
5. Criterion computes pairwise logits and weighted preference loss.
6. Task evaluation converts scores to probabilities and computes accuracy.
## SANA Model Path
1. Text branch:
- Load SANA tokenizer/text encoder(s) by checkpoint profile.
- Build prompt conditioning embeddings for transformer cross-attention.
- Build pooled text representation for reward embedding projection.
2. Image branch:
- Preprocess RGB image to configured resolution.
- Encode image with SANA VAE to latents.
- Apply scheduler-based noise at selected timesteps.
- Run SANA transformer on noisy latents with text conditioning.
- Pool transformer latent outputs to image representation.
3. Reward projection:
- Project text and image representations to common dimension.
- L2-normalize features.
- Compute scaled similarity using learnable `logit_scale`.
## Training Controls
- Mixed precision and gradient accumulation from accelerator config.
- Distributed feature gather for pairwise loss consistency.
- Validation/test evaluation at configured intervals.
- Periodic checkpointing and best-metric tracking.
## Logging and Outputs
- Output root: `logs/lrm/{project_name}/{run_name}`.
- Saved artifacts:
- config snapshot
- train log
- periodic checkpoints
- final model state
## Variant-Aware Profile Design
A model profile selects:
- pretrained checkpoint id
- default image size (512 or 1024)
- tokenizer/text path behavior
- optional special overrides for memory or precision
This keeps one codebase while enabling all four requested SANA checkpoints.
|