File size: 2,105 Bytes
b4efe93 | 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 | # Flux LRM Implementation Plan
## Goal
Build a working latent-space reward model for FLUX.1-schnell using the same pairwise preference dataset protocol used by the SD1.5 and SDXL variants.
## Scope
- Reuse the existing trainer architecture (accelerator/task/criterion/dataset/model split).
- Use FLUX.1-schnell latent + transformer path for reward feature extraction.
- Train on the same Pick-a-Pic style pairwise data format.
- Keep docs for this variant inside flux/docs.
## Implementation Phases
1. Scaffold and rename
- Create a dedicated flux package with trainer modules and run script.
- Ensure all config groups are registered with Flux names.
2. Flux model wrapper
- Load FLUX components: VAE, scheduler, transformer, CLIP tokenizer+encoder, T5 tokenizer+encoder.
- Encode prompts with dual encoders.
- Encode images to latents, apply flow-style noising, and pack latents.
- Run Flux transformer and pool token outputs to image features.
- Project text/image features into shared reward embedding space.
3. Dataset and criterion
- Keep pairwise data contract compatible with existing task/criterion.
- Use CLIP tokenizer for input_ids and T5 tokenizer for input_ids_2.
- Keep timestep sampling support (constant/variable and comparison mode).
- Reuse pairwise loss logic from SD variants.
4. Config and training wiring
- Provide step_flux_base Hydra config with Flux defaults.
- Keep optimizer/scheduler/accelerator knobs aligned with existing variants.
5. Validation and smoke tests
- Verify imports and Python syntax.
- Compose Hydra config.
- Run a minimal initialization smoke test.
## Current Status
- Scaffold and naming migration: in progress/completed for main files.
- Flux model implementation: in progress.
- Dataset and criterion adaptation: in progress.
- Config wiring: in progress.
- Smoke validation: pending.
## Risks
- Flux model memory footprint is high; batch size may require reduction for first run.
- Timestep indexing must stay consistent with scheduler timesteps/sigmas.
- External model download/auth may block runtime tests if network credentials are missing.
|