lrm/flux/docs/plan.md

# 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.