[niehen6174]

Summary

Adds qwen_image21 as a separate train/rollout family from Qwen-Image 1.0, plus one official 1-GPU T2I recipe: Flow-GRPO + PickScore.

2.1 is Qwen3-VL conditioning, single-stream block-causal DiT, RGBA VAE, and CFG-distilled (default guidance 1.0). This PR only covers the T2I path. I2I is served by the engine but has no recipe yet.

Qwen-Image 2.1 support is already on sglang main and diffusers main (landed via sglang #39983 and diffusers #14804). This PR does not pin those packages to a released version yet; until miles-diffusion's default deps include them, point SGLANG_21_PYTHON / DIFFUSERS_21_SRC at current main checkouts.

What's included

Train config (qwen_image21)

• Timesteps ÷ 1000 (matches diffusers t / 1000)
• CFG combine without cond-norm rescale
• Rebuild img_mask / img_shapes from --diffusion-height/width
• Slice noise pred to the image tail ([:, -target_tokens:])
• Rebuild RoPE caches on CUDA after materialize
• LoRA targets: to_q/k/v, to_out.0, img_mlp.{proj,out,gate_layer}

Rollout patch group qwen_image21

• Disable 2.1 fused kernels
• Wipe prefix caches every step so the engine reruns the full joint sequence (same as train)
• Reuse the 1.0 PEFT-ordered LoRA forwards for IPC-synced adapters

Recipe

• scripts/run_diffusion_grpo_qwenimage21_pickscore_1gpu_v2.py
• 512² / 10 steps / CFG 1.0 / SDE 3–5 / LoRA 64/128
• 16 prompts × 16 samples, lr 1e-4, 200 rollouts
• --rollout-return-full-trajectory so SDE-window conversion can index a full 0..T latent array

Docs / tests

• Family page: docs/models/qwen-image/qwen-image-21.md
• Listed in README / docs index / CLI help
• Fast tests for config, registry, FSDP plan, and patch-group registration

How to run

Recipes must pass --diffusion-model-family qwen_image21. A bare qwen-image substring (including Qwen/Qwen-Image-2.1) still resolves to the 1.0 family.

python3 scripts/run_diffusion_grpo_qwenimage21_pickscore_1gpu_v2.py \
  --hf-checkpoint /path/to/Qwen-Image-2.1 \
  --prompt-data /path/to/flowgrpo_pickscore

If the installed sglang / diffusers are not yet new enough, set SGLANG_21_PYTHON and DIFFUSERS_21_SRC at current main checkouts. Real-weight GRPO needs ≥80GB; 32GB only finished generate.

Results

PickScore reward_mean

log_prob_mean_abs_diff

Out of scope / follow-ups

• Pin sglang / diffusers (and any related extras) to released versions that already include 2.1, so recipes do not need SGLANG_21_PYTHON / DIFFUSERS_21_SRC
• Engine echo of latent_step_indices — this PR uses --rollout-return-full-trajectory instead
• Bitwise train/rollout (needs a joint engine forward + more op patches; dummy alignment is log_prob_mean_abs_diff ≈ 5e-5)
• I2I recipe

Test plan

• Fast tests: qwen_image21 config / registry / FSDP plan / patch group
• Dummy alignment: log_prob_mean_abs_diff ~5e-5
• Real ckpt generate
• 1-GPU GRPO on ≥80GB with the v2 recipe