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.