qwen image 21 control, update flex forcing and fix bug in qwen image 21

2026-09-23 13:53:03 +08:00
parent bce42b301f
commit 261e1d85d3
17 changed files with 6828 additions and 84 deletions
@@ -1112,7 +1112,7 @@ def main():
            aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()}
            
            if args.fix_sample_size is not None:
-                fix_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size]
+                fix_sample_size = [int(x / 32) * 32 for x in args.fix_sample_size]  # 32 = vae_scale_factor(16)*2: keeps the latent grid even so h*w % 4 == 0 (joint stream expands image slots 4x)
            elif args.random_ratio_crop:
                if rng is None:
                    random_sample_size = aspect_ratio_random_crop_sample_size[
@@ -1122,10 +1122,10 @@ def main():
                    random_sample_size = aspect_ratio_random_crop_sample_size[
                        rng.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB)
                    ]
-                random_sample_size = [int(x / 16) * 16 for x in random_sample_size]
+                random_sample_size = [int(x / 32) * 32 for x in random_sample_size]  # 32 = vae_scale_factor(16)*2: keep latent dims even
            else:
                closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size)
-                closest_size = [int(x / 16) * 16 for x in closest_size]
+                closest_size = [int(x / 32) * 32 for x in closest_size]  # 32 = vae_scale_factor(16)*2: keep latent dims even

            for example in examples:
                if args.fix_sample_size is not None:
@@ -0,0 +1,548 @@
+# Qwen-Image 2.1 Control (ControlNet-Union) Training Guide
+
+This document provides a complete workflow for training a **ControlNet-Union** adapter on top of the frozen
+**Qwen-Image 2.1** base transformer: environment setup, data preparation, distributed training, CFG
+distillation, and inference testing.
+
+A parallel chain of zero-initialized `control_blocks` produces a per-layer skip (`hints`) that is added back into
+the frozen base blocks, so the adapter starts as an identity skip and learns control gradually. Only the control
+modules are trained (`--trainable_modules "control"`).
+
+The adapter is a **union** of control + inpaint: the conditioning tensor `control_context` packs
+`[control_latents (64) | mask (1) | masked-image latents (64)] = 129` channels, so one adapter handles both
+spatial control (depth / canny / pose / …) and inpainting.
+
+---
+
+## Table of Contents
+- [1. Environment Setup](#1-environment-setup)
+- [2. Data Preparation](#2-data-preparation)
+  - [2.1 Quick Test Dataset](#21-quick-test-dataset)
+  - [2.2 Dataset Structure](#22-dataset-structure)
+  - [2.3 metadata.json Format](#23-metadatajson-format)
+  - [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage)
+- [3. Control Training](#3-control-training)
+  - [3.1 Download Pre-trained Model](#31-download-pre-trained-model)
+  - [3.2 Quick Start (DeepSpeed-Zero-2)](#32-quick-start-deepspeed-zero-2)
+  - [3.3 Common Training Parameters](#33-common-training-parameters)
+  - [3.4 Training Validation](#34-training-validation)
+  - [3.5 Training with FSDP](#35-training-with-fsdp)
+  - [3.6 Other Backends](#36-other-backends)
+  - [3.7 Multi-machine Distributed Training](#37-multi-machine-distributed-training)
+  - [3.8 CFG Distillation](#38-cfg-distillation)
+- [4. Inference Testing](#4-inference-testing)
+  - [4.1 Inference Parameters](#41-inference-parameters)
+  - [4.2 Single GPU Inference](#42-single-gpu-inference)
+  - [4.3 Multi-GPU Parallel Inference](#43-multi-gpu-parallel-inference)
+- [5. Additional Resources](#5-additional-resources)
+
+---
+
+## 1. Environment Setup
+
+**Option 1: Using requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**Option 2: Manual Installation**
+
+```bash
+pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
+pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
+pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
+pip install yunchang xfuser modelscope openpyxl
+pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
+pip install opencv-python-headless
+pip install deepspeed==0.17.0 numpy==1.26.4
+```
+
+**Option 3: Using Docker**
+
+When using Docker, please ensure that the GPU drivers and CUDA environment are correctly installed, then execute the following commands:
+
+```
+# pull image
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# enter image
+docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+```
+
+> **Qwen-Image 2.1 specific**: the text encoder is a **Qwen3-VL** model, so the environment needs a `transformers`
+> build that ships the `qwen3_vl` architecture (newer than the base pin in `requirements.txt`). If
+> `Qwen3VLForConditionalGeneration` / `Qwen3VLProcessor` import as `None`, your `transformers` is too old.
+>
+> The YOLO object-mask feature (see [2.3](#23-metadatajson-format)) needs `ultralytics`; `yolov8x-seg.pt`
+> downloads automatically on first use.
+
+---
+
+## 2. Data Preparation
+
+Control training uses `ImageVideoControlDataset`.
+
+### 2.1 Quick Test Dataset
+
+We provide a test dataset containing several training samples with corresponding control files.
+
+```bash
+# Download official example dataset
+modelscope download --dataset PAI/X-Fun-Images-Controls-Demo --local_dir ./datasets/X-Fun-Images-Controls-Demo
+```
+
+### 2.2 Dataset Structure
+
+```
+📦 datasets/
+├── 📂 my_dataset/
+│   ├── 📂 train/                 # target images (what the model should generate)
+│   │   ├── 📄 image001.jpg
+│   │   └── 📄 ...
+│   ├── 📂 control/               # paired control / condition images (pose, canny, depth, ...)
+│   │   ├── 📄 image001.png
+│   │   └── 📄 ...
+│   └── 📄 metadata.json
+```
+
+### 2.3 metadata.json Format
+
+The manifest is the standard image metadata JSON plus one extra `control_file_path` field that pairs each
+**target** image with its **control** image.
+
+```json
+[
+    {
+      "file_path": "train/image001.jpg",
+      "control_file_path": "control/image001.png",
+      "text": "A young woman, studio lighting, high quality.",
+      "width": 1024,
+      "height": 1024,
+      "type": "image"
+    }
+]
+```
+
+**Key field descriptions**:
+- `file_path`: the **target** image (relative or absolute).
+- `control_file_path`: the **control / condition** image (pose map, edge map, depth, gray sketch, …). It is loaded
+  as RGB and resized/cropped with the **same** transform as the target, so they stay pixel-aligned.
+- `text`: caption.
+- `width` / `height`: recommended for bucket training; use `scripts/process_json_add_width_and_height.py` to add
+  them to a JSON that lacks them.
+- `type`: `"image"`.
+
+> **You only supply the target + control images. You do NOT supply masks.** The inpaint mask is generated on the
+> fly:
+> - A random rectangular hole via `get_random_mask` in the collate.
+> - Then, on a random ~70% subset of frames, an **irregular object-shaped mask** produced by a **YOLO-seg**
+>   detector (`ObjectInstanceDetector`), with its edges randomly dilated / eroded / Gaussian-blurred. This mirrors
+>   `scripts/qwenimage_fun/train_control.py` and gives the model realistic, object-shaped inpaint holes instead of
+>   only rectangles.
+> - The masked image fed to the union branch is always `target * (1 - mask)`.
+
+> **RGBA note**: the 2.1 VAE reads RGBA. Training images are loaded as RGB and automatically composited over an
+> opaque alpha channel before encoding, so you do not need to provide RGBA data.
+
+### 2.4 Relative vs Absolute Path Usage
+
+**Relative paths** (small, local dataset):
+```bash
+export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json"
+```
+
+**Absolute paths** (NAS / OSS / multi-machine shared data):
+```bash
+export DATASET_NAME=""
+export DATASET_META_NAME="/mnt/data/metadata.json"
+```
+
+> If the dataset is stored on external storage or shared across machines, prefer absolute paths.
+
+---
+
+## 3. Control Training
+
+### 3.1 Download Pre-trained Model
+
+Point `MODEL_NAME` at a local **Qwen-Image 2.1** checkpoint directory. Its `transformer/` subfolder supplies the
+frozen base weights; the control modules are zero-initialized on load.
+
+```bash
+mkdir -p models/Diffusion_Transformer
+# Place your Qwen-Image 2.1 weights here, e.g.
+# models/Diffusion_Transformer/Qwen-Image-2.1/{transformer,vae,text_encoder,...}
+```
+
+> **No released 2.1 ControlNet-Union checkpoint.** Unlike Qwen-Image 2512, there is currently no published
+> `...-Fun-Controlnet-Union.safetensors` for 2.1, so training starts from scratch with the **zero-initialized**
+> control branch. Consequently the launcher leaves `--transformer_path` out; only add it to resume or fine-tune a
+> control checkpoint you have already trained (or produced with `scripts/*/extract_control_weights.py`).
+
+### 3.2 Quick Start (DeepSpeed-Zero-2)
+
+It is recommended to use DeepSpeed-Zero-2 or FSDP for training, which can save a significant amount of GPU memory.
+
+After following **2.1 Quick Test Dataset** and **3.1 Download Pre-trained Model**, you can directly copy and run the following command:
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image-2.1"
+export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/qwenimage21_fun/train_control.py \
+  --config_path="config/qwenimage21/qwenimage21_control.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --train_batch_size=1 \
+  --image_sample_size=1024 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-05 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_qwen_image_21_control" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --enable_bucket \
+  --uniform_sampling \
+  --trainable_modules "control"
+```
+
+### 3.3 Common Training Parameters
+
+| Parameter | Description | Example Value |
+|-----------|-------------|---------------|
+| `--config_path` | **Required.** Builds the control transformer with `control_layers` / `control_in_dim` | `config/qwenimage21/qwenimage21_control.yaml` |
+| `--pretrained_model_name_or_path` | Base Qwen-Image 2.1 model (frozen weights) | `models/Diffusion_Transformer/Qwen-Image-2.1` |
+| `--train_data_dir` / `--train_data_meta` | Dataset root / manifest JSON | `""` / `/path/metadata.json` |
+| `--trainable_modules` | `"control"` trains only `control_blocks.*` + `control_img_in.*`; base stays frozen | `"control"` |
+| `--transformer_path` | **Omit for from-scratch** training; add only to resume/finetune a trained control branch | *(none)* |
+| `--image_sample_size` | Max training resolution, auto bucketing | `1024` |
+| `--train_batch_size` / `--gradient_accumulation_steps` | Per-device batch / accumulation | `1` / `1` |
+| `--learning_rate` | Initial learning rate | `2e-05` |
+| `--lr_scheduler` / `--lr_warmup_steps` | Scheduler / warmup | `constant_with_warmup` / `100` |
+| `--checkpointing_steps` | Save a checkpoint every N steps | `50` |
+| `--gradient_checkpointing` | Activation recomputation | flag |
+| `--vae_mini_batch` | Mini-batch size for VAE encoding (control encodes 3 streams) | `1` |
+| `--max_grad_norm` | Gradient clipping | `0.05` |
+| `--enable_bucket` | Bucket training by resolution without cropping | flag |
+| `--uniform_sampling` | Uniform timestep sampling | flag |
+| `--low_vram` | Offload VAE / text encoder when idle to save memory | flag (optional) |
+
+> **Memory**: each control step encodes **three** latent streams (target, control, masked image), so it is heavier
+> than base training. Keep `--vae_mini_batch=1`; add `--low_vram` if you are tight on memory.
+
+### 3.4 Training Validation
+
+Configure validation during training to periodically render control previews:
+
+```bash
+  --validation_paths "asset/pose.jpg" \
+  --validation_steps=50 \
+  --validation_epochs=500 \
+  --validation_prompts="1girl, black_hair, brown_eyes, ... solo, upper_body"
+```
+
+- `--validation_prompts` and `--validation_paths` must have **matching counts**; entry `i` of each pair is used
+  together. The output resolution is derived from each control image's aspect ratio via
+  `calculate_dimensions(image_sample_size^2, w/h)`.
+- Validation triggers on either `--validation_steps` or `--validation_epochs`.
+- Previews are written to `{output_dir}/sample/`. Because the 2.1 VAE decodes to **RGBA**, previews are saved as
+  **`.png`** (JPEG cannot store alpha).
+- `log_validation` is wrapped in `try/except`: a bad control path only logs `Eval error on rank N` and never
+  crashes training. To confirm validation actually produced images, check `output_dir/sample/` **and** grep the log
+  for `Eval error`.
+
+### 3.5 Training with FSDP
+
+If DeepSpeed-Zero-2 runs out of GPU memory, you can switch to FSDP for training. The launcher `scripts/qwenimage21_fun/train_control.sh` runs
+exactly the command below; edit the paths at the top (`MODEL_NAME`, `DATASET_META_NAME`, …) and run it with `bash` if you prefer
+(the wrap classes must match the control model's blocks):
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image-2.1"
+export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+  --fsdp_transformer_layer_cls_to_wrap=BaseQwenImage21TransformerBlock,QwenImage21ControlTransformerBlock \
+  --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+  --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+  scripts/qwenimage21_fun/train_control.py \
+  --config_path="config/qwenimage21/qwenimage21_control.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --train_batch_size=1 \
+  --image_sample_size=1024 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-05 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_qwen_image_21_control" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --enable_bucket \
+  --uniform_sampling \
+  --trainable_modules "control"
+```
+
+### 3.6 Other Backends
+
+#### 3.6.1 Training without DeepSpeed and FSDP
+
+Using neither DeepSpeed nor FSDP may result in insufficient GPU memory; only recommended when GPU memory is
+sufficient. Plain DDP also has to replicate the 2.1 base transformer on every GPU, so it is generally not
+recommended:
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image-2.1"
+export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" scripts/qwenimage21_fun/train_control.py \
+  --config_path="config/qwenimage21/qwenimage21_control.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --train_batch_size=1 \
+  --image_sample_size=1024 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-05 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_qwen_image_21_control" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --enable_bucket \
+  --uniform_sampling \
+  --trainable_modules "control"
+```
+
+### 3.7 Multi-machine Distributed Training
+
+**Suitable for**: Ultra-large-scale datasets, faster training speed
+
+#### 3.7.1 Environment Configuration
+
+When using multi-machine training, please set the following environment variables:
+
+```bash
+export MASTER_ADDR="your master address"
+export MASTER_PORT=10086
+export WORLD_SIZE=1 # The number of machines
+export NUM_PROCESS=8 # The number of processes, such as WORLD_SIZE * 8
+export RANK=0 # The rank of this machine
+
+accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK scripts/qwenimage21_fun/train_control.py \
+  [other training parameters...]
+```
+
+#### 3.7.2 Multi-machine Training Considerations
+
+- **Network Requirements**:
+   - Recommended: RDMA/InfiniBand (high performance)
+   - Without RDMA, add environment variables:
+     ```bash
+     export NCCL_IB_DISABLE=1
+     export NCCL_P2P_DISABLE=1
+     ```
+
+- **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage)
+
+### 3.8 CFG Distillation
+
+`train_control_distill.py` / `train_control_distill.sh` are an **optional second stage**. They distill
+classifier-free guidance (CFG) into the trained control branch, so that **inference needs no guidance scale**.
+
+Algorithm (identical in spirit to `scripts/minimax_h3_fun` / `scripts/flux2_fun` control distillation):
+- A **frozen teacher** (a second copy of the same control model, loaded with the **same** `--transformer_path` —
+  your stage-1 trained control branch) runs two forward passes per step, on the prompt and on an **empty**
+  negative prompt, both **with the control condition**. The two velocities combine into the CFG target:
+  `target = uncond + (cond - uncond) * real_guidance_scale`.
+- The trainable **student** (control branch) runs a single conditional forward and regresses onto that target
+  (MSE in velocity space). Only `--trainable_modules "control"` is trained.
+
+Run it after you have a trained control checkpoint:
+
+```bash
+# Set CONTROL_TRANSFORMER_PATH to the stage-1 checkpoint's
+# output_dir_qwen_image_21_control/<ts>/checkpoint-<step>/diffusion_pytorch_model.safetensors
+bash scripts/qwenimage21_fun/train_control_distill.sh
+```
+
+Distillation-specific parameters:
+
+| Parameter | Description | Example Value |
+|-----------|-------------|---------------|
+| `--transformer_path` | **Required**: the trained control branch both student and teacher load | `/root/diffusion_pytorch_model.safetensors` |
+| `--real_guidance_scale` | CFG scale applied to the teacher to build the target | `3.5` |
+| `--learning_rate` | Lower LR for distillation | `2e-06` |
+| `--output_dir` | Separate output dir for the distilled adapter | `output_dir_qwen_image_21_control_distill` |
+
+> The teacher is a separate, unsharded bf16 copy on each GPU (only the student is FSDP-sharded), which is
+> memory-heavy. Add `--low_vram` to stream the teacher only for its two forward passes. A distilled student is run
+> at inference with `guidance_scale = 1.0` (CFG is already baked into the weights).
+
+---
+
+## 4. Inference Testing
+
+### 4.1 Inference Parameters
+
+| Parameter | Description | Example Value |
+|-----------|-------------|---------------|
+| `config_path` | Must match the trained adapter's config | `config/qwenimage21/qwenimage21_control.yaml` |
+| `model_name` | Base Qwen-Image 2.1 path | `models/Diffusion_Transformer/Qwen-Image-2.1` |
+| `transformer_path` | Trained control weights (`control_*` keys load with `strict=False`), or `None` for the base | `output_dir_qwen_image_21_control/.../diffusion_pytorch_model.safetensors` |
+| `sampler_name` | Flow-matching sampler | `Flow` |
+| `sample_size` | Output canvas `[height, width]` | `[1728, 992]` |
+| `control_image` | Control condition image (`predict_t2i_control.py`) | `asset/pose.jpg` |
+| `control_image_path` | Optional control condition image (`predict_i2i_inpaint.py`, defaults to `None`) | `asset/pose.jpg` |
+| `control_context_scale` | Control-branch strength (the value the adapter was trained to consume) | `1.0` |
+| `image_path` / `mask_path` | Inpaint source image / mask (only `predict_i2i_inpaint.py`, see 4.2) | `asset/pose.jpg` / `asset/mask.png` |
+| `guidance_scale` | CFG strength. `1.0` for a CFG-distilled checkpoint | `1.0` |
+| `weight_dtype` | Use `torch.float16` on GPUs without bf16 (v100, 2080Ti, …) | `torch.bfloat16` |
+| `GPU_memory_mode` | GPU memory management mode, see table below | `model_group_offload` |
+| `ulysses_degree` / `ring_degree` | Multi-GPU parallelism (see 4.3). `ring_degree` must stay `1` | `1` / `1` |
+| `num_inference_steps` / `seed` | Sampling steps / seed | `40` / `43` |
+| `save_path` | Output directory | `samples/qwenimage21-control-images` |
+
+**GPU Memory Management Modes**:
+
+| Mode | Description | Memory Usage |
+|------|------|---------|
+| `model_full_load` | Load entire model to GPU | Highest |
+| `model_full_load_and_qfloat8` | Full load + FP8 quantization | High |
+| `model_cpu_offload` | Offload model to CPU after use | Medium |
+| `model_cpu_offload_and_qfloat8` | CPU offload + FP8 quantization | Medium-Low |
+| `model_group_offload` | Switch layer groups between CPU/CUDA | Low |
+| `sequential_cpu_offload` | Layer-by-layer offload (slowest) | Lowest |
+
+### 4.2 Single GPU Inference
+
+#### Quick Start
+
+```bash
+python examples/qwenimage21_fun/predict_t2i_control.py
+```
+
+Edit the top-of-file constants to match your setup. The pipeline preprocesses and VAE-encodes `control_image`
+(accepts a PIL image / a path), builds the 129-channel `control_context`, and injects the control skips:
+
+```python
+GPU_memory_mode     = "model_group_offload"
+model_name          = "models/Diffusion_Transformer/Qwen-Image-2.1"
+transformer_path    = "models/Personalized_Model/Qwen-Image-2.1-Fun-Controlnet-Union.safetensors"  # or your trained checkpoint's diffusion_pytorch_model.safetensors
+control_image       = "asset/pose.jpg"
+control_context_scale = 1.0
+prompt              = "A young woman with long straight black hair ..."
+sample_size         = [1728, 992]
+num_inference_steps = 40
+```
+
+Results are saved to `samples/qwenimage21-control-images/*.png`.
+
+> The KV cache is **disabled automatically** when `control_context` is present: the control skip depends on the
+> per-step base joint stream, so a prefix cache would be invalid.
+
+**Image Inpainting Inference**:
+
+The union adapter also does inpainting. `predict_i2i_inpaint.py` feeds `image_path` + `mask_path` (and leaves
+`control_image_path` as `None`), so only the inpaint half of the 129-channel context is used:
+
+```bash
+python examples/qwenimage21_fun/predict_i2i_inpaint.py
+```
+
+`mask_path` semantics: **white (`>= 0.5`) = repaint**, black = keep. You can supply a control image *and* the
+inpaint pair together to use the full union.
+
+### 4.3 Multi-GPU Parallel Inference
+
+**Suitable for**: High-resolution generation, faster inference
+
+Qwen-Image 2.1 supports **Ulysses (head-parallel) sequence parallelism only**.
+
+#### Install Parallel Inference Dependencies
+
+```bash
+pip install xfuser==0.4.2 yunchang==0.6.2
+```
+
+#### Configure Parallel Strategy
+
+Edit `examples/qwenimage21_fun/predict_t2i_control.py`:
+
+```python
+# Ensure ulysses_degree × ring_degree = number of GPUs
+# For example, using 4 GPUs:
+ulysses_degree = 4  # Head dimension parallelism
+ring_degree = 1     # Sequence dimension parallelism, must stay 1
+```
+
+**Configuration Principles**:
+- `ulysses_degree` must divide `num_attention_heads` (32): one of `1/2/4/8/16/32`.
+- `ring_degree` **must stay 1** — ring attention rotates KV chunks and cannot express 2.1's block-causal mask or
+  its prefix KV cache.
+
+**Example Configurations**:
+
+| GPU count | ulysses_degree | ring_degree |
+|-----------|----------------|-------------|
+| 1 | 1 | 1 |
+| 4 | 4 | 1 |
+| 8 | 8 | 1 |
+
+#### Run Multi-GPU Inference
+
+```bash
+# Set ulysses_degree > 1, keep ring_degree = 1, and GPU count = ulysses_degree * ring_degree.
+torchrun --nproc-per-node=4 examples/qwenimage21_fun/predict_t2i_control.py
+```
+
+---
+
+## 5. Additional Resources
+
+- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
+- **Qwen-Image Official Repository**: https://github.com/QwenLM/Qwen-Image
@@ -0,0 +1,531 @@
+# Qwen-Image 2.1 Control(ControlNet-Union)训练指南
+
+本文档提供在冻结的 **Qwen-Image 2.1** 基座 transformer 之上训练 **ControlNet-Union** 适配器的完整流程,包括环境配置、
+数据准备、分布式训练、CFG 蒸馏与推理测试。
+
+一整套零初始化的 `control_blocks` 并行链会产出逐层残差(`hints`),再加回冻结的基座 block,因此适配器初始时等价于恒等
+跳连,并逐步学习控制信号。训练时只更新控制模块(`--trainable_modules "control"`)。
+
+该适配器是 control + inpaint 的 **union**:条件张量 `control_context` 打包了
+`[control_latents(64) | mask(1) | masked-image latents(64)] = 129` 个通道,因此单个适配器同时处理空间控制
+(depth / canny / pose / …)与图像修补(inpainting)。
+
+---
+
+## 目录
+- [一、环境配置](#一环境配置)
+- [二、数据准备](#二数据准备)
+  - [2.1 快速测试数据集](#21-快速测试数据集)
+  - [2.2 数据集结构](#22-数据集结构)
+  - [2.3 metadata.json 格式](#23-metadatajson-格式)
+  - [2.4 相对路径与绝对路径使用方案](#24-相对路径与绝对路径使用方案)
+- [三、Control 训练](#三control-训练)
+  - [3.1 下载预训练模型](#31-下载预训练模型)
+  - [3.2 快速开始(DeepSpeed-Zero-2)](#32-快速开始deepspeed-zero-2)
+  - [3.3 训练常用参数解析](#33-训练常用参数解析)
+  - [3.4 训练验证](#34-训练验证)
+  - [3.5 使用 FSDP 训练](#35-使用-fsdp-训练)
+  - [3.6 其他后端](#36-其他后端)
+  - [3.7 多机分布式训练](#37-多机分布式训练)
+  - [3.8 CFG 蒸馏](#38-cfg-蒸馏)
+- [四、推理测试](#四推理测试)
+  - [4.1 推理参数解析](#41-推理参数解析)
+  - [4.2 单卡推理](#42-单卡推理)
+  - [4.3 多卡并行推理](#43-多卡并行推理)
+- [五、更多资源](#五更多资源)
+
+---
+
+## 一、环境配置
+
+**方式 1：使用requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**方式 2：手动安装依赖**
+
+```bash
+pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
+pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
+pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
+pip install yunchang xfuser modelscope openpyxl
+pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
+pip install opencv-python-headless
+pip install deepspeed==0.17.0 numpy==1.26.4
+```
+
+**方式 3：使用docker**
+
+使用docker的情况下,请保证机器中已经正确安装显卡驱动与CUDA环境,然后依次执行以下命令:
+
+```
+# pull image
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# enter image
+docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+```
+
+> **Qwen-Image 2.1 特有**:文本编码器是 **Qwen3-VL** 模型,因此环境需要一个包含 `qwen3_vl` 结构的 `transformers`
+> 版本(比 `requirements.txt` 里的基线更新)。如果 `Qwen3VLForConditionalGeneration` / `Qwen3VLProcessor` 导入为
+> `None`,说明你的 `transformers` 太旧。
+>
+> YOLO 目标掩膜功能(见 [2.3](#23-metadatajson-格式))需要 `ultralytics`;`yolov8x-seg.pt` 会在首次使用时自动下载。
+
+---
+
+## 二、数据准备
+
+Control 训练使用 `ImageVideoControlDataset`。
+
+### 2.1 快速测试数据集
+
+我们提供了一个测试的数据集，其中包含若干训练数据以及对应的控制文件。
+
+```bash
+# 下载官方示例数据集
+modelscope download --dataset PAI/X-Fun-Images-Controls-Demo --local_dir ./datasets/X-Fun-Images-Controls-Demo
+```
+
+### 2.2 数据集结构
+
+```
+📦 datasets/
+├── 📂 my_dataset/
+│   ├── 📂 train/                 # 目标图(模型应生成的内容)
+│   │   ├── 📄 image001.jpg
+│   │   └── 📄 ...
+│   ├── 📂 control/               # 配对的 control / 条件图(pose、canny、depth…)
+│   │   ├── 📄 image001.png
+│   │   └── 📄 ...
+│   └── 📄 metadata.json
+```
+
+### 2.3 metadata.json 格式
+
+清单文件是标准图像 metadata JSON,外加一个 `control_file_path` 字段,把每张**目标图**与其 **control 图**配对。
+
+```json
+[
+    {
+      "file_path": "train/image001.jpg",
+      "control_file_path": "control/image001.png",
+      "text": "A young woman, studio lighting, high quality.",
+      "width": 1024,
+      "height": 1024,
+      "type": "image"
+    }
+]
+```
+
+**关键字段说明**:
+- `file_path`:**目标图**(相对或绝对路径)。
+- `control_file_path`:**control / 条件图**(姿态图、边缘图、深度图、线稿…)。它以 RGB 载入,并与目标图使用**完全相同**的
+  变换做 resize / crop,从而保证像素对齐。
+- `text`:描述(caption)。
+- `width` / `height`:推荐提供,用于 bucket 训练;可用 `scripts/process_json_add_width_and_height.py` 为缺失字段的
+  JSON 补上。
+- `type`:图像数据为 `"image"`。
+
+> **你只需提供目标图 + control 图,不需要提供掩膜。** inpaint 掩膜是即时生成的:
+> - 先在 collate 中用 `get_random_mask` 生成随机矩形遮挡。
+> - 然后在随机约 70% 的帧上,再由 **YOLO-seg** 检测器(`ObjectInstanceDetector`)生成**不规则的目标形状掩膜**,
+>   其边缘随机做膨胀 / 腐蚀 / 高斯模糊。这与 `scripts/qwenimage_fun/train_control.py` 一致,能让模型见到真实、
+>   目标形状的修补空洞,而不只是矩形。
+> - 送进 union 分支的被遮罩图始终是 `target * (1 - mask)`。
+
+> **RGBA 说明**:2.1 VAE 读取 RGBA。训练图以 RGB 载入,在编码前会自动合成到不透明 alpha 通道,因此你不需要提供 RGBA 数据。
+
+### 2.4 相对路径与绝对路径使用方案
+
+**相对路径**(小型本地数据集):
+```bash
+export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json"
+```
+
+**绝对路径**(NAS / OSS / 多机共享数据):
+```bash
+export DATASET_NAME=""
+export DATASET_META_NAME="/mnt/data/metadata.json"
+```
+
+> 如果数据集存放在外部存储或被多台机器共享,推荐使用绝对路径。
+
+---
+
+## 三、Control 训练
+
+### 3.1 下载预训练模型
+
+将 `MODEL_NAME` 指向本地的 **Qwen-Image 2.1** checkpoint 目录。其 `transformer/` 子目录提供冻结的基座权重;控制模块在
+载入时零初始化。
+
+```bash
+mkdir -p models/Diffusion_Transformer
+# 将你的 Qwen-Image 2.1 权重放到这里,例如
+# models/Diffusion_Transformer/Qwen-Image-2.1/{transformer,vae,text_encoder,...}
+```
+
+> **没有公开的 2.1 ControlNet-Union checkpoint。** 与 Qwen-Image 2512 不同,2.1 目前并没有发布的
+> `...-Fun-Controlnet-Union.safetensors`,因此训练从零开始,control 分支**零初始化**。所以启动脚本不写 `--transformer_path`;
+> 只有在你需要 resume / fine-tune 一个已训练好的 control checkpoint(或你自己用 `scripts/*/extract_control_weights.py` 得到的)
+> 时才加上它。
+
+### 3.2 快速开始（DeepSpeed-Zero-2）
+
+推荐使用 DeepSpeed-Zero-2 或 FSDP 方案进行训练,可以节省大量显存。
+
+如果按照 **2.1 快速测试数据集**下载数据与 **3.1 下载预训练模型**放置权重后,直接复制以下启动指令进行启动。
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image-2.1"
+export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/qwenimage21_fun/train_control.py \
+  --config_path="config/qwenimage21/qwenimage21_control.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --train_batch_size=1 \
+  --image_sample_size=1024 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-05 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_qwen_image_21_control" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --enable_bucket \
+  --uniform_sampling \
+  --trainable_modules "control"
+```
+
+### 3.3 训练常用参数解析
+
+| 参数 | 说明 | 示例值 |
+|------|------|--------|
+| `--config_path` | **必填。** 用 `control_layers` / `control_in_dim` 构建 control transformer | `config/qwenimage21/qwenimage21_control.yaml` |
+| `--pretrained_model_name_or_path` | 基座 Qwen-Image 2.1 模型(冻结权重) | `models/Diffusion_Transformer/Qwen-Image-2.1` |
+| `--train_data_dir` / `--train_data_meta` | 数据根目录 / 清单 JSON | `""` / `/path/metadata.json` |
+| `--trainable_modules` | `"control"` 只训练 `control_blocks.*` + `control_img_in.*`,基座冻结 | `"control"` |
+| `--transformer_path` | **从零训练时省略**;仅在 resume / fine-tune 已训练的 control 分支时加上 | *(无)* |
+| `--image_sample_size` | 最大训练分辨率,自动 bucket | `1024` |
+| `--train_batch_size` / `--gradient_accumulation_steps` | 单卡 batch / 梯度累积 | `1` / `1` |
+| `--learning_rate` | 初始学习率 | `2e-05` |
+| `--lr_scheduler` / `--lr_warmup_steps` | 学习率调度 / 预热 | `constant_with_warmup` / `100` |
+| `--checkpointing_steps` | 每 N 步保存 checkpoint | `50` |
+| `--gradient_checkpointing` | 激活重计算 | flag |
+| `--vae_mini_batch` | VAE 编码 mini-batch(control 要编码 3 路 latent) | `1` |
+| `--max_grad_norm` | 梯度裁剪 | `0.05` |
+| `--enable_bucket` | 按分辨率分组、不裁剪的 bucket 训练 | flag |
+| `--uniform_sampling` | 均匀 timestep 采样 | flag |
+| `--low_vram` | 空闲时卸载 VAE / 文本编码器以省显存 | flag(可选) |
+
+> **显存**:每个 control step 要编码**三路** latent(目标图、control 图、被遮罩图),比普通训练更重。请保持
+> `--vae_mini_batch=1`;显存紧张时再加 `--low_vram`。
+
+### 3.4 训练验证
+
+在训练时配置验证参数,定期渲染 control 预览:
+
+```bash
+  --validation_paths "asset/pose.jpg" \
+  --validation_steps=50 \
+  --validation_epochs=500 \
+  --validation_prompts="1girl, black_hair, brown_eyes, ... solo, upper_body"
+```
+
+- `--validation_prompts` 与 `--validation_paths` 的数量必须**匹配**,每对的第 i 项一起使用。输出分辨率由每张 control 图的
+  宽高比经 `calculate_dimensions(image_sample_size^2, w/h)` 推出。
+- 验证在 `--validation_steps` 或 `--validation_epochs` 任一满足时触发。
+- 预览写入 `{output_dir}/sample/`。由于 2.1 VAE 解码为 **RGBA**,预览保存为 **`.png`**(JPEG 无法存 alpha)。
+- `log_validation` 包在 `try/except` 里:错误的 control 路径只会打印 `Eval error on rank N`,不会让训练崩溃。要确认验证
+  是否真的产出了图,请查看 `output_dir/sample/` **并**在日志里 grep `Eval error`。
+
+### 3.5 使用 FSDP 训练
+
+如果 DeepSpeed-Zero-2 显存不足,可以切换使用 FSDP 进行训练。配套的启动脚本 `scripts/qwenimage21_fun/train_control.sh` 运行的就是
+下面这条命令,修改其顶部路径(`MODEL_NAME`、`DATASET_META_NAME` 等)后可直接 `bash` 运行(wrap 类必须与 control 模型的 block 同名):
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image-2.1"
+export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+  --fsdp_transformer_layer_cls_to_wrap=BaseQwenImage21TransformerBlock,QwenImage21ControlTransformerBlock \
+  --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+  --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+  scripts/qwenimage21_fun/train_control.py \
+  --config_path="config/qwenimage21/qwenimage21_control.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --train_batch_size=1 \
+  --image_sample_size=1024 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-05 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_qwen_image_21_control" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --enable_bucket \
+  --uniform_sampling \
+  --trainable_modules "control"
+```
+
+### 3.6 其他后端
+
+#### 3.6.1 不使用 DeepSpeed 与 FSDP 训练
+
+不使用 DeepSpeed 或 FSDP 可能会导致显存不足,仅建议在显存充足的情况下使用;普通 DDP 还需要在每张卡上完整复制 2.1 基座
+transformer,通常并不推荐:
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image-2.1"
+export DATASET_NAME="datasets/X-Fun-Images-Controls-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Images-Controls-Demo/metadata_add_width_height.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" scripts/qwenimage21_fun/train_control.py \
+  --config_path="config/qwenimage21/qwenimage21_control.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --train_batch_size=1 \
+  --image_sample_size=1024 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-05 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_qwen_image_21_control" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --enable_bucket \
+  --uniform_sampling \
+  --trainable_modules "control"
+```
+
+### 3.7 多机分布式训练
+
+**适合场景**：超大规模数据集、需要更快的训练速度
+
+#### 3.7.1 环境配置
+
+当使用多机训练时，请设置以下环境变量：
+
+```bash
+export MASTER_ADDR="your master address"
+export MASTER_PORT=10086
+export WORLD_SIZE=1 # The number of machines
+export NUM_PROCESS=8 # The number of processes, such as WORLD_SIZE * 8
+export RANK=0 # The rank of this machine
+
+accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK scripts/qwenimage21_fun/train_control.py \
+  [其他训练参数...]
+```
+
+#### 3.7.2 多机训练注意事项
+
+- **网络要求**：
+   - 推荐 RDMA/InfiniBand（高性能）
+   - 无 RDMA 时添加环境变量：
+     ```bash
+     export NCCL_IB_DISABLE=1
+     export NCCL_P2P_DISABLE=1
+     ```
+
+- **数据同步**：所有机器必须能够访问相同的数据路径（NFS/共享存储）
+
+### 3.8 CFG 蒸馏
+
+`train_control_distill.py` / `train_control_distill.sh` 是一个**可选的第二阶段**。它把 classifier-free guidance(CFG)
+蒸馏进已训练好的 control 分支,使**推理阶段无需 guidance scale**。
+
+算法(思路与 `scripts/minimax_h3_fun` / `scripts/flux2_fun` 的 control 蒸馏完全一致):
+- **冻结的 teacher**(同一 control 模型的另一份拷贝,用**相同**的 `--transformer_path` 载入,即你第一阶段训练好的 control 分支)
+  每步做两次前向,分别在 prompt 与**空**负 prompt 上,两者都**带 control 条件**。两个速度合成 CFG 目标:
+  `target = uncond + (cond - uncond) * real_guidance_scale`。
+- 可训练的 **student**(control 分支)只跑一次带条件的前向,向该目标回归(速度空间 MSE)。同样只训练
+  `--trainable_modules "control"`。
+
+在你已有一个训练好的 control checkpoint 后运行:
+
+```bash
+# 将 CONTROL_TRANSFORMER_PATH 指向第一阶段 checkpoint 的
+# output_dir_qwen_image_21_control/<ts>/checkpoint-<step>/diffusion_pytorch_model.safetensors
+bash scripts/qwenimage21_fun/train_control_distill.sh
+```
+
+蒸馏专有参数:
+
+| 参数 | 说明 | 示例值 |
+|------|------|--------|
+| `--transformer_path` | **必填**:student 与 teacher 都载入的、已训练好的 control 分支 | `/root/diffusion_pytorch_model.safetensors` |
+| `--real_guidance_scale` | 作用于 teacher 以合成目标的 CFG scale | `3.5` |
+| `--learning_rate` | 蒸馏使用更低的学习率 | `2e-06` |
+| `--output_dir` | 蒸馏适配器单独的输出目录 | `output_dir_qwen_image_21_control_distill` |
+
+> teacher 是每卡上一份独立的、不分片的 bf16 拷贝(只有 student 被 FSDP 分片),因此很吃显存。可加 `--low_vram` 让 teacher
+> 只在其两次前向时上卡。蒸馏后的 student 推理时用 `guidance_scale = 1.0`(CFG 已烘焙进权重)。
+
+---
+
+## 四、推理测试
+
+### 4.1 推理参数解析
+
+| 参数 | 说明 | 示例值 |
+|------|------|--------|
+| `config_path` | 必须与训练好的适配器 config 一致 | `config/qwenimage21/qwenimage21_control.yaml` |
+| `model_name` | 基座 Qwen-Image 2.1 路径 | `models/Diffusion_Transformer/Qwen-Image-2.1` |
+| `transformer_path` | 训练好的 control 权重(`control_*` 键以 `strict=False` 载入),基线可为 `None` | `output_dir_qwen_image_21_control/.../diffusion_pytorch_model.safetensors` |
+| `sampler_name` | flow-matching 采样器 | `Flow` |
+| `sample_size` | 输出画布 `[height, width]` | `[1728, 992]` |
+| `control_image` | control 条件图(`predict_t2i_control.py`) | `asset/pose.jpg` |
+| `control_image_path` | 可选 control 条件图(`predict_i2i_inpaint.py`,默认 `None`) | `asset/pose.jpg` |
+| `control_context_scale` | control 分支强度(适配器训练时消费的取值) | `1.0` |
+| `image_path` / `mask_path` | inpaint 输入图 / 掩码(仅 `predict_i2i_inpaint.py`,见 4.2) | `asset/pose.jpg` / `asset/mask.png` |
+| `guidance_scale` | CFG 强度。CFG 蒸馏后的 checkpoint 用 `1.0` | `1.0` |
+| `weight_dtype` | 不支持 bf16 的卡(v100、2080Ti…)用 `torch.float16` | `torch.bfloat16` |
+| `GPU_memory_mode` | 显存管理模式,可选值见下表 | `model_group_offload` |
+| `ulysses_degree` / `ring_degree` | 多卡并行(见 4.3)。`ring_degree` 必须为 `1` | `1` / `1` |
+| `num_inference_steps` / `seed` | 采样步数 / 随机种子 | `40` / `43` |
+| `save_path` | 输出目录 | `samples/qwenimage21-control-images` |
+
+**显存管理模式说明**:
+
+| 模式 | 说明 | 显存占用 |
+|------|------|---------|
+| `model_full_load` | 整个模型加载到 GPU | 最高 |
+| `model_full_load_and_qfloat8` | 全量加载 + FP8 量化 | 高 |
+| `model_cpu_offload` | 使用后将模型卸载到 CPU | 中等 |
+| `model_cpu_offload_and_qfloat8` | CPU 卸载 + FP8 量化 | 中低 |
+| `model_group_offload` | 层组在 CPU/CUDA 间切换 | 低 |
+| `sequential_cpu_offload` | 逐层卸载(速度最慢) | 最低 |
+
+### 4.2 单卡推理
+
+#### 快速开始
+
+```bash
+python examples/qwenimage21_fun/predict_t2i_control.py
+```
+
+修改文件顶部常量以匹配你的环境。pipeline 会预处理并 VAE 编码 `control_image`(接受 PIL 图 / 路径),构建 129 通道的
+`control_context`,并注入 control 残差:
+
+```python
+GPU_memory_mode     = "model_group_offload"
+model_name          = "models/Diffusion_Transformer/Qwen-Image-2.1"
+transformer_path    = "models/Personalized_Model/Qwen-Image-2.1-Fun-Controlnet-Union.safetensors"  # 或训练输出的 diffusion_pytorch_model.safetensors
+control_image       = "asset/pose.jpg"
+control_context_scale = 1.0
+prompt              = "A young woman with long straight black hair ..."
+sample_size         = [1728, 992]
+num_inference_steps = 40
+```
+
+结果保存到 `samples/qwenimage21-control-images/*.png`。
+
+> 当存在 `control_context` 时,KV cache 会**自动禁用**:control 残差依赖每步的基座 joint 流,前缀缓存会失效。
+
+**图像修补推理**:
+
+union 适配器也能做 inpainting。`predict_i2i_inpaint.py` 传入 `image_path` + `mask_path`(并把 `control_image_path` 留为 `None`),
+于是只用 129 通道条件的 inpaint 半边:
+
+```bash
+python examples/qwenimage21_fun/predict_i2i_inpaint.py
+```
+
+`mask_path` 语义:**白色(`>= 0.5`)= 重绘**,黑色 = 保留。你可以同时提供 control 图与 inpaint 对,以使用完整 union。
+
+### 4.3 多卡并行推理
+
+**适合场景**:高分辨率生成、加速推理
+
+Qwen-Image 2.1 **仅支持 Ulysses(head 并行)序列并行**。
+
+#### 安装并行推理依赖
+
+```bash
+pip install xfuser==0.4.2 yunchang==0.6.2
+```
+
+#### 配置并行策略
+
+编辑 `examples/qwenimage21_fun/predict_t2i_control.py`:
+
+```python
+# 确保 ulysses_degree × ring_degree = GPU 数量
+# 例如使用 4 张 GPU:
+ulysses_degree = 4  # Head 维度并行
+ring_degree = 1     # Sequence 维度并行,必须保持为 1
+```
+
+**配置原则**:
+- `ulysses_degree` 必须整除 `num_attention_heads`(32):取 `1/2/4/8/16/32`。
+- `ring_degree` **必须为 1** —— ring attention 会轮转 KV chunk,无法表达 2.1 的 block-causal mask 或其前缀 KV cache。
+
+**示例配置**:
+
+| GPU 数 | ulysses_degree | ring_degree |
+|--------|----------------|-------------|
+| 1 | 1 | 1 |
+| 4 | 4 | 1 |
+| 8 | 8 | 1 |
+
+#### 运行多卡推理
+
+```bash
+# 设 ulysses_degree > 1,保持 ring_degree = 1,且 GPU 数 = ulysses_degree * ring_degree。
+torchrun --nproc-per-node=4 examples/qwenimage21_fun/predict_t2i_control.py
+```
+
+---
+
+## 五、更多资源
+
+- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
+- **Qwen-Image 官方仓库**：https://github.com/QwenLM/Qwen-Image
@@ -0,0 +1,119 @@
+# Extract the control branch of a trained QwenImage21ControlTransformer2DModel checkpoint.
+#
+# `train_control.py` / `train_control_distill.py` save the whole transformer (frozen base branch + trainable control
+# branch) in the diffusers layout; with FSDP the gathered state dict is written as
+# `<checkpoint>/diffusion_pytorch_model.safetensors` (no `transformer/` subdir, no `_control` suffix). The control
+# branch is everything that `QwenImage21ControlTransformer2DModel` adds on top of the base model: the
+# `control_blocks.*` list (one block per `control_layers` entry) and the `control_img_in.*` input projection. This
+# script writes just those tensors to a standalone safetensors file, which can be re-applied onto a fresh base model
+# built from `config/qwenimage21/qwenimage21_control.yaml` with `transformer.load_state_dict(..., strict=False)`
+# (only the `control_*` keys are consumed; the base keys already match the freshly-loaded base weights).
+#
+# Usage:
+#   python scripts/qwenimage21_fun/extract_control_weights.py \
+#       --model_path /path/to/train_control/checkpoint-xxx/diffusion_pytorch_model.safetensors \
+#       --output_path /path/to/control_weights.safetensors
+import argparse
+import json
+import os
+
+import torch
+from safetensors.torch import load_file, save_file
+
+CONTROL_PREFIXES = ("control_blocks.", "control_img_in.")
+# FSDP / DeepSpeed unwrap may leave wrapper prefixes on the keys; strip them to the bare model namespace.
+WRAPPER_PREFIXES = ("_fsdp_wrapped_module.", "_fsdp_wrapped_module_", "module.", "_orig_mod.")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Extract the control-branch weights (control_blocks / control_img_in) of a trained "
+        "Qwen-Image 2.1 control transformer into a standalone safetensors file."
+    )
+    parser.add_argument(
+        "--model_path", type=str, default="output_dir_qwen_image_21_control/checkpoint-10000/diffusion_pytorch_model.safetensors",
+        help="Path to the saved transformer: a directory containing diffusion_pytorch_model*.safetensors "
+        "(e.g. the FSDP `<checkpoint>` dir), or a single .safetensors file.",
+    )
+    parser.add_argument(
+        "--output_path", type=str, default="output_dir_qwen_image_21_control/checkpoint-10000/diffusion_pytorch_model_control.safetensors",
+        help="Where to write the extracted control weights.",
+    )
+    return parser.parse_args()
+
+
+def resolve_safetensor_files(model_path):
+    if os.path.isdir(model_path):
+        shards = sorted(
+            os.path.join(model_path, name)
+            for name in os.listdir(model_path)
+            if name.endswith(".safetensors")
+        )
+        if not shards:
+            raise FileNotFoundError(f"No .safetensors files found under {model_path}.")
+        return shards
+    if os.path.isfile(model_path) and model_path.endswith(".safetensors"):
+        return [model_path]
+    raise FileNotFoundError(f"--model_path must be a safetensors file or a directory of them, got {model_path}.")
+
+
+def unwrap_key(key):
+    changed = True
+    while changed:
+        changed = False
+        for prefix in WRAPPER_PREFIXES:
+            if key.startswith(prefix):
+                key = key[len(prefix):]
+                changed = True
+    return key
+
+
+def main():
+    args = parse_args()
+
+    state_dict = {}
+    for shard in resolve_safetensor_files(args.model_path):
+        state_dict.update(load_file(shard, device="cpu"))
+
+    control_state_dict = {}
+    for key, value in state_dict.items():
+        bare_key = unwrap_key(key)
+        if bare_key.startswith(CONTROL_PREFIXES):
+            control_state_dict[bare_key] = value.contiguous()
+
+    if not control_state_dict:
+        raise ValueError(
+            f"No control-branch keys (control_blocks.* / control_img_in.*) found in {args.model_path}; "
+            "this checkpoint does not look like a Qwen-Image 2.1 control training output."
+        )
+
+    # Carry the branch layout next to the weights so a loader can rebuild the same control model without
+    # inspecting the full training config. `config.json` of the saved transformer records both fields; keep
+    # the safetensors metadata strings-only.
+    metadata = {"format": "pt"}
+    config_path = os.path.join(args.model_path, "config.json") if os.path.isdir(args.model_path) else None
+    if config_path is not None and os.path.isfile(config_path):
+        with open(config_path, "r") as file:
+            config = json.load(file)
+        for field in ("control_layers", "control_in_dim"):
+            if field in config:
+                metadata[field] = json.dumps(config[field])
+
+    os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True)
+    save_file(control_state_dict, args.output_path, metadata=metadata)
+
+    num_params = sum(value.numel() for value in control_state_dict.values())
+    block_ids = sorted({
+        int(key.split(".")[1]) for key in control_state_dict if key.startswith("control_blocks.")
+    })
+    print(f"Extracted {len(control_state_dict)} control tensors ({num_params / 1e9:.3f}B params) -> {args.output_path}")
+    print(f"  control_blocks indices: {block_ids}")
+    if "control_in_dim" in metadata:
+        print(f"  control_in_dim: {metadata['control_in_dim']}, control_layers: {metadata['control_layers']}")
+    for key in sorted(control_state_dict):
+        if key.endswith(".weight") and control_state_dict[key].dim() >= 2:
+            print(f"  {key}: {list(control_state_dict[key].shape)} {control_state_dict[key].dtype}")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,39 @@
+export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image-2.1"
+export DATASET_NAME="datasets/internal_datasets/"
+export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=BaseQwenImage21TransformerBlock,QwenImage21ControlTransformerBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/qwenimage21_fun/train_control.py \
+  --config_path="config/qwenimage21/qwenimage21_control.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --train_batch_size=1 \
+  --image_sample_size=1328 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-05 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_qwen_image_21_control" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --enable_bucket \
+  --random_hw_adapt \
+  --low_vram \
+  --uniform_sampling \
+  --trainable_modules "control" \
+  --resume_from_checkpoint="latest"
@@ -0,0 +1,40 @@
+export MODEL_NAME="models/Diffusion_Transformer/Qwen-Image-2.1"
+export DATASET_NAME="datasets/internal_datasets/"
+export DATASET_META_NAME="datasets/internal_datasets/metadata.json"
+# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
+# export NCCL_IB_DISABLE=1
+# export NCCL_P2P_DISABLE=1
+NCCL_DEBUG=INFO
+
+accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=BaseQwenImage21TransformerBlock,QwenImage21ControlTransformerBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/qwenimage21_fun/train_control_distill.py \
+  --config_path="config/qwenimage21/qwenimage21_control.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --train_batch_size=1 \
+  --image_sample_size=1328 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --checkpointing_steps=50 \
+  --learning_rate=2e-06 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
+  --seed=42 \
+  --output_dir="output_dir_qwen_image_21_control_distill" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --adam_weight_decay=3e-2 \
+  --adam_epsilon=1e-10 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=0.05 \
+  --enable_bucket \
+  --random_hw_adapt \
+  --low_vram \
+  --uniform_sampling \
+  --transformer_path="models/Personalized_Model/Qwen-Image-2.1-Fun-Controlnet-Union.safetensors" \
+  --trainable_modules "control" \
+  --resume_from_checkpoint="latest
@@ -81,8 +81,8 @@ from videox_fun.pipeline import (WanI2VPipeline, WanPipeline,
                                 WanFlexForcingPipeline,
                                 WanSelfForcingPipeline)
 from videox_fun.utils.discrete_sampler import DiscreteSampling
-from videox_fun.utils.flex_chunking import (UNIFORM_BLOCK_PROB,
-                                            broadcast_chunk_sizes,
+from videox_fun.utils.flex_chunking import (broadcast_chunk_sizes,
+                                            build_full_then_blocks_partitions,
                                            build_pyramid_partitions,
                                            chunk_boundaries,
                                            sample_flexible_chunks,
@@ -317,7 +317,7 @@ def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer
                            # would validate out of distribution too. The uniform
                            # block layout is the one band that is both fixed
                            # across checkpoints and actually trained:
-                            # UNIFORM_BLOCK_PROB of iterations use it, against
+                            # FLEX_ARM_PROB of iterations use it, against
                            # ~1.3% for the most common random partition.
                            flex_kwargs = dict(
                                chunk_spec=args.num_frame_per_block,
@@ -379,14 +379,16 @@ def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=Non
    return torch.clip(t.to(torch.int32), low, high - 1)


-# Fraction of pyramid iterations whose level 0 is the whole clip as a single
-# chunk, i.e. fully bidirectional. 3.2's coarse planning step is exactly that at
-# inference (`denoise_mode="pyramid"` with no `chunk_spec` builds `[[F], ...]`),
-# so it has to appear in training; the remaining iterations keep the random
-# partition of 3.1 so the causal end and every layout in between stay covered.
-# Deliberately not a CLI flag - it is a property of the paper's schedule, not a
-# knob the launcher should have to keep in sync.
-COARSE_GLOBAL_PROB = 0.5
+# Each training iteration draws its level-0 layout from four equally likely arms
+# (FLEX_ARM_PROB = 25% of iterations each): the whole-clip planning chunk (3.2's
+# coarse step, `[[F], ...]` at inference), the launcher's own uniform
+# `num_frame_per_block`, the fixed "first step full, then block-major" ladder
+# (`denoise_mode="full_then_blocks"`), and a random 3.1 partition. Equal shares
+# keep the coarse, causal and every in-between layout covered while giving the
+# full_then_blocks ladder enough iterations to train the block-major-at-high-noise
+# trajectory the binary pyramid never reaches. Deliberately not a CLI flag - it is
+# a property of the schedule mixture, not a knob the launcher keeps in sync.
+FLEX_ARM_PROB = 0.25

 # How many steps of a launch report the partition they drew. Counted from this
 # process rather than from `global_step`, so a resumed run still gets its own
@@ -406,59 +408,110 @@ def sample_flex_partitions(args, num_frames, num_denoising_steps, torch_rng,
    per denoising step. Returns ``None`` when Flex-Forcing is off, which leaves
    the inherited uniform ``num_frame_per_block`` masks untouched.

-    The level-0 layout is a three-way mixture decided by a single uniform draw,
-    so the constants below are the actual iteration shares: ``COARSE_GLOBAL_PROB``
-    of iterations use one chunk over the whole clip (what inference's coarse
-    planning step uses), ``UNIFORM_BLOCK_PROB`` pin the launcher's own uniform
-    ``num_frame_per_block``, and the rest draw a random 2..10 partition. All of
-    them are then refined into the same ladder, so a single set of weights covers
-    `denoise_mode="pyramid"` whether or not the caller pins `chunk_spec`. With no
-    pyramid the coarse band is empty and the split is 10% uniform / 90% random.
+    The level-0 layout is a four-way mixture decided by a single uniform draw,
+    each arm taking ``FLEX_ARM_PROB`` (25%) of iterations: one chunk over the
+    whole clip (inference's coarse planning step), the launcher's own uniform
+    ``num_frame_per_block``, the fixed "first step full, then block-major" ladder
+    (`denoise_mode="full_then_blocks"`), and a random 2..10 partition. The first,
+    second and last are then refined into the same binary pyramid, so a single set
+    of weights covers `denoise_mode="pyramid"` whether or not the caller pins
+    `chunk_spec`; the full_then_blocks arm bypasses that refinement and emits its
+    two-level ladder whole, drawing the block-major width between fully causal (1
+    frame / block) and `num_frame_per_block`. With no pyramid the coarse arm is
+    empty (a `[F]` level 0 would be a degenerate single-chunk rollout), so the
+    split is 25% uniform / 25% full_then_blocks / 50% random.
    """
    if not args.flex_forcing:
        return None
-    u = torch.rand((), generator=torch_rng, device=device).item()
-    coarse_prob = COARSE_GLOBAL_PROB if args.flex_pyramid_levels > 1 else 0.0
-    if u < coarse_prob:
+    # The arm selector must be IDENTICAL on every rank, not just its broadcast
+    # count. Each arm now post-processes the received partition its own way (the
+    # coarse/uniform/random trio funnel through the shared refine tail below, but
+    # full_then_blocks sets the ladder directly from a different tensor), so a
+    # per-rank `u` would leave ranks in different arms building different ladders
+    # -> different walk forward counts -> FSDP all-gather desync (NCCL hang).
+    # Broadcast rank 0's draw so all ranks take the same arm; the arm's existing
+    # single broadcast then reconciles its within-arm randomness (the random base,
+    # the block width) to rank 0, making every ladder byte-identical.
+    u = torch.rand((), generator=torch_rng, device=device)
+    if dist.is_available() and dist.is_initialized():
+        dist.broadcast(u.reshape(1), src=0)
+    u = u.item()
+    pyramid = args.flex_pyramid_levels > 1
+    # Four equally likely arms; the coarse (whole-clip planning) one only when a
+    # pyramid is on, else its 25% folds into the random arm.
+    coarse_hi = FLEX_ARM_PROB if pyramid else 0.0
+    uniform_hi = coarse_hi + FLEX_ARM_PROB
+    ftb_hi = uniform_hi + FLEX_ARM_PROB
+    ladder = None
+    if u < coarse_hi:
        # Coarse end of 3.1: reuse the ladder builder's own level-0 rule so the
        # `independent_first_frame` handling cannot drift from inference's.
        base = build_pyramid_partitions(
            num_frames, num_levels=1, base_chunks=None,
            independent_first_frame=args.independent_first_frame)[0]
        arm = "whole clip, the 3.2 coarse planning layout"
-    elif u < coarse_prob + UNIFORM_BLOCK_PROB:
+    elif u < uniform_hi:
        # `uniform_chunks`, not `normalize_chunk_spec`: the latter takes no
        # `independent_first_frame` argument (it encodes that as a leading 1), so
-        # it would silently disagree with the two bands around it. This is the
-        # band `log_validation` renders when there is no pyramid.
+        # it would silently disagree with the bands around it. This is the band
+        # `log_validation` renders when there is no pyramid.
        base = uniform_chunks(
            num_frames, args.num_frame_per_block,
            independent_first_frame=args.independent_first_frame)
        arm = f"uniform {args.num_frame_per_block}-frame blocks, the launcher's own layout"
+    elif u < ftb_hi:
+        # "First step full, every later step block-major": a fixed two-level
+        # ladder that does NOT go through the binary pyramid refinement, so the
+        # block-major level is reached at high noise (the 2nd step) instead of as
+        # a deep refinement - the one trajectory the pyramid arms never sample.
+        # The block-major width is drawn between fully causal (1 frame / block)
+        # and the launcher's `num_frame_per_block`, so both the tight-AR and the
+        # coarse-block refinement of the whole-clip plan get trained. The block
+        # width is a per-rank draw, so it has to be reconciled with EXACTLY ONE
+        # broadcast - the same collective count every other arm issues (they each
+        # broadcast their single base) - or the ranks desync and NCCL deadlocks.
+        # Level 0 is just [F], identical on every rank and needing no sync, so we
+        # broadcast only the block-major level; rank 0's draw wins and every rank
+        # rebuilds the same two-level ladder locally.
+        block_choices = sorted({1, max(1, int(args.num_frame_per_block))})
+        pick = int(torch.rand((), generator=torch_rng, device=device).item()
+                   * len(block_choices))
+        block = block_choices[min(pick, len(block_choices) - 1)]
+        block = min(block, int(num_frames))
+        ftb = build_full_then_blocks_partitions(
+            num_frames, block,
+            independent_first_frame=args.independent_first_frame)
+        ftb[1] = broadcast_chunk_sizes(ftb[1], device=device)
+        ladder = ftb[:num_denoising_steps]
+        arm = (f"full clip -> uniform {max(ftb[1])}-frame blocks, "
+               "the full_then_blocks ladder")
    else:
        base = sample_flexible_chunks(
            num_frames, min_chunk=args.flex_chunk_min, max_chunk=args.flex_chunk_max,
            generator=torch_rng, device=device,
            independent_first_frame=args.independent_first_frame)
        arm = f"random {args.flex_chunk_min}..{args.flex_chunk_max}-frame blocks, the 3.1 spectrum"
-    # Every rank has to train the same layout: the FlexAttention mask, and the
-    # `num_frame_per_block` derived from it, must agree across the SP/FSDP group.
-    base = broadcast_chunk_sizes(base, device=device)
-    ladder = [base]
-    if args.flex_pyramid_levels > 1:
-        ladder = build_pyramid_partitions(
-            num_frames, num_levels=args.flex_pyramid_levels,
-            min_num_frame_per_block=args.flex_min_num_frame_per_block, base_chunks=base,
-            independent_first_frame=args.independent_first_frame)
-        if len(ladder) > num_denoising_steps:
-            # Same short-circuit as at inference, where the rollout stops refining at
-            # the last step: deeper levels would never be reached, so drop them
-            # instead of reporting a pyramid that was not actually trained.
-            if verbose:
-                print(f"--flex_pyramid_levels={args.flex_pyramid_levels} builds "
-                      f"{len(ladder)} levels but only {num_denoising_steps} denoising "
-                      f"steps are trained; keeping the first {num_denoising_steps}.")
-            ladder = ladder[:num_denoising_steps]
+    if ladder is None:
+        # Every rank has to train the same layout: the FlexAttention mask, and the
+        # `num_frame_per_block` derived from it, must agree across the SP/FSDP
+        # group. (The full_then_blocks arm above is deterministic given `args` and
+        # already broadcast each level, so it skips this path.)
+        base = broadcast_chunk_sizes(base, device=device)
+        ladder = [base]
+        if pyramid:
+            ladder = build_pyramid_partitions(
+                num_frames, num_levels=args.flex_pyramid_levels,
+                min_num_frame_per_block=args.flex_min_num_frame_per_block, base_chunks=base,
+                independent_first_frame=args.independent_first_frame)
+            if len(ladder) > num_denoising_steps:
+                # Same short-circuit as at inference, where the rollout stops
+                # refining at the last step: deeper levels would never be reached,
+                # so drop them instead of reporting a pyramid not actually trained.
+                if verbose:
+                    print(f"--flex_pyramid_levels={args.flex_pyramid_levels} builds "
+                          f"{len(ladder)} levels but only {num_denoising_steps} denoising "
+                          f"steps are trained; keeping the first {num_denoising_steps}.")
+                ladder = ladder[:num_denoising_steps]
    if verbose:
        # Every level, not just the drawn one: level 0 is what the mixture above
        # picked, the rest are derived from it, and they are the sub-spans the