Update one-forcing and causal forcing params

2026-10-30 16:33:15 +08:00
parent a15b293128
commit f00a7d480e
17 changed files with 4866 additions and 86 deletions
@@ -0,0 +1,303 @@
+import os
+import sys
+
+import numpy as np
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
+from PIL import Image
+
+current_file_path = os.path.abspath(__file__)
+project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
+for project_root in project_roots:
+    sys.path.insert(0, project_root) if project_root not in sys.path else None
+
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
+                               WanT5EncoderModel,
+                               WanTransformer3DModel_SelfForcing)
+from videox_fun.pipeline import WanSelfForcingPipeline
+from videox_fun.utils import (FlowDPMSolverMultistepScheduler,
+                              FlowUniPCMultistepScheduler,
+                              apply_gpu_memory_mode, filter_kwargs, merge_lora,
+                              save_videos_grid, unmerge_lora)
+
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# model_full_load means that the entire model will be moved to the GPU.
+#
+# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
+# and the transformer model has been quantized to float8, which can save more GPU memory.
+#
+# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
+#
+# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
+# and the transformer model has been quantized to float8, which can save more GPU memory.
+#
+# model_group_offload transfers internal layer groups between CPU/CUDA,
+# balancing memory efficiency and speed between full-module and leaf-level offloading methods.
+#
+# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
+# resulting in slower speeds but saving a large amount of GPU memory.
+GPU_memory_mode     = "model_full_load"
+# Multi GPUs config
+# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used.
+# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4.
+# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1.
+ulysses_degree      = 1
+ring_degree         = 1
+# Use FSDP to save more GPU memory in multi gpus.
+fsdp_dit            = False
+fsdp_text_encoder   = True
+# Compile will give a speedup in fixed resolution and need a little GPU memory.
+# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload.
+compile_dit         = False
+
+# Config and model path
+config_path         = "config/wan2.1/wan_civitai.yaml"
+# One-Forcing distills a Wan2.1-T2V-1.3B student (the 14B model is only the DMD teacher at training time).
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+
+# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++".
+sampler_name        = "Flow"
+# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics.
+# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
+shift               = 5.0
+# One-Forcing is a DMD-distilled student, so it uses stochastic sampling.
+stochastic_sampling = True
+
+# One-Forcing checkpoint to overlay on top of the Wan2.1 base model. Point at the trainer
+# output checkpoint dir (its live/EMA weights are auto-resolved below) or a plain `.pt`/`.safetensors` file.
+# Default: the official One-Forcing checkpoint, downloaded via
+#   hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+# The GAN discriminator head is a training-only submodule of fake_score, so it never appears in
+# the generator checkpoint and inference needs no discriminator wiring.
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+# Load the EMA generator (transformer_ema/) when True, otherwise the live generator weights.
+# Only applies to trainer-output checkpoint dirs; a plain `.pt`/`.safetensors` file is loaded as-is.
+use_ema             = False
+vae_path            = None
+lora_path           = None
+
+# Other params
+sample_size         = [480, 832]
+video_length        = 81
+fps                 = 16
+
+# One-Forcing causal inference config
+# `num_frame_per_block`: 1 = frame-wise (One-Forcing default), 3 = chunk-wise.
+num_frame_per_block     = 1
+# Local attention window size (-1 for global attention)
+local_attn_size         = -1
+# Others
+independent_first_frame = False
+context_noise           = 0.0
+
+# Use torch.float16 if GPU does not support torch.bfloat16
+# Some graphics cards, such as v100, 2080ti, do not support torch.bfloat16
+weight_dtype        = torch.bfloat16
+prompt              = "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."
+negative_prompt     = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+# -------- One-Forcing step selector --------
+# One-Forcing bakes CFG into the student weights, so inference MUST use guidance_scale=1.0.
+# The official recipe generates each frame in a single step (denoising_step_list=[1000]):
+#   1-step (One-Forcing default): num_inference_steps = 1
+#   4-step (more stable, slower): num_inference_steps = 4 -> [1000, 750, 500, 250]
+# With stochastic_sampling=True, num_inference_steps=1 yields the warped single timestep [1000],
+# matching the official framewise 1-step generation.
+guidance_scale      = 1.0
+num_inference_steps = 1
+seed                = 43
+lora_weight         = 0.55
+save_path           = "samples/wan-videos-one-forcing"
+
+device = set_multi_gpus_devices(ulysses_degree, ring_degree)
+config = OmegaConf.load(config_path)
+
+# Load transformer with causal inference support if enabled
+transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs'])
+transformer_additional_kwargs['local_attn_size'] = local_attn_size
+
+transformer = WanTransformer3DModel_SelfForcing.from_pretrained(
+    os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+    transformer_additional_kwargs=transformer_additional_kwargs,
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+if transformer_path is not None:
+    def _resolve_transformer_path(raw_path: str, prefer_ema: bool) -> str:
+        """Resolve a checkpoint path to the actual weights file to load.
+
+        File paths are returned unchanged so external `.pt` ckpts keep working. For trainer-output
+        dirs, prefer the EMA generator when asked and available, otherwise fall back to the live
+        generator. The FSDP trainer writes the EMA under `transformer_ema/`; `ema_transformer/` is
+        kept as a compatibility alias.
+        """
+        if os.path.isfile(raw_path):
+            return raw_path
+        if os.path.isdir(raw_path):
+            candidates = []
+            if prefer_ema:
+                candidates.append(os.path.join(raw_path, "transformer_ema", "diffusion_pytorch_model.safetensors"))
+                candidates.append(os.path.join(raw_path, "ema_transformer", "diffusion_pytorch_model.safetensors"))
+            candidates.append(os.path.join(raw_path, "transformer", "diffusion_pytorch_model.safetensors"))
+            candidates.append(os.path.join(raw_path, "diffusion_pytorch_model.safetensors"))
+            for c in candidates:
+                if os.path.isfile(c):
+                    return c
+        raise FileNotFoundError(
+            f"transformer_path={raw_path!r} is neither a file nor a checkpoint dir "
+            f"with a known safetensors layout (transformer_ema/, transformer/ or root)."
+        )
+
+    _raw_transformer_path = transformer_path
+    transformer_path = _resolve_transformer_path(transformer_path, prefer_ema=use_ema)
+    if transformer_path != _raw_transformer_path:
+        print(f"use_ema={use_ema}: resolved {_raw_transformer_path} -> {transformer_path}")
+    print(f"From checkpoint: {transformer_path}")
+    if transformer_path.endswith("safetensors"):
+        from safetensors.torch import load_file
+        state_dict = load_file(transformer_path)
+    else:
+        state_dict = torch.load(transformer_path, map_location="cpu")
+
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+    state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict
+    state_dict = state_dict["generator"] if "generator" in state_dict else state_dict
+    # FSDP-saved ckpts keep the `model._fsdp_wrapped_module.` prefix; strip it before the generic
+    # `model.` strip so both FSDP and plain ckpts land at bare parameter names.
+    if any("._fsdp_wrapped_module." in k for k in state_dict.keys()):
+        state_dict = {k.replace("model._fsdp_wrapped_module.", "model.", 1) if k.startswith("model._fsdp_wrapped_module.") else k: v for k, v in state_dict.items()}
+    if any(k.startswith("model.") for k in state_dict.keys()):
+        state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()}
+
+    m, u = transformer.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# Get Vae
+vae = AutoencoderKLWan.from_pretrained(
+    os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),
+    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+).to(weight_dtype)
+
+if vae_path is not None:
+    print(f"From checkpoint: {vae_path}")
+    if vae_path.endswith("safetensors"):
+        from safetensors.torch import load_file
+        state_dict = load_file(vae_path)
+    else:
+        state_dict = torch.load(vae_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = vae.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# Get Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
+)
+
+# Get Text encoder
+text_encoder = WanT5EncoderModel.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+    additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+# Get Scheduler
+Chosen_Scheduler = {
+    "Flow": FlowMatchEulerDiscreteScheduler,
+    "Flow_Unipc": FlowUniPCMultistepScheduler,
+    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
+}[sampler_name]
+if sampler_name == "Flow_Unipc" or sampler_name == "Flow_DPM++":
+    config['scheduler_kwargs']['shift'] = 1
+scheduler = Chosen_Scheduler(
+    **filter_kwargs(Chosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
+)
+
+# Get Pipeline
+pipeline = WanSelfForcingPipeline(
+    transformer=transformer,
+    vae=vae,
+    tokenizer=tokenizer,
+    text_encoder=text_encoder,
+    scheduler=scheduler,
+)
+
+if ulysses_degree > 1 or ring_degree > 1:
+    from functools import partial
+    transformer.enable_multi_gpus_inference()
+    if fsdp_dit:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.transformer = shard_fn(pipeline.transformer)
+        print("Add FSDP DIT")
+    if fsdp_text_encoder:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.text_encoder = shard_fn(pipeline.text_encoder)
+        print("Add FSDP TEXT ENCODER")
+
+if compile_dit:
+    for i in range(len(pipeline.transformer.blocks)):
+        pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
+    print("Add Compile")
+
+# Quantize (when the mode carries an "_and_<quant>" suffix) and then place the pipeline.
+# The order lives inside the helper: quantization has to happen before the offload hooks
+# are installed.
+apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype)
+
+generator = torch.Generator(device=device).manual_seed(seed)
+
+if lora_path is not None:
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
+
+with torch.no_grad():
+    video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
+    latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
+
+    sample = pipeline(
+        prompt,
+        num_frames = video_length,
+        negative_prompt = negative_prompt,
+        height      = sample_size[0],
+        width       = sample_size[1],
+        generator   = generator,
+        guidance_scale          = guidance_scale,
+        num_inference_steps     = num_inference_steps,
+        shift                   = shift,
+        num_frame_per_block     = num_frame_per_block,
+        independent_first_frame = independent_first_frame,
+        context_noise           = context_noise,
+        stochastic_sampling     = stochastic_sampling,
+    ).videos
+
+if lora_path is not None:
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
+
+def save_results():
+    if not os.path.exists(save_path):
+        os.makedirs(save_path, exist_ok=True)
+
+    index = len([path for path in os.listdir(save_path)]) + 1
+    prefix = str(index).zfill(8)
+    if video_length == 1:
+        video_path = os.path.join(save_path, prefix + ".png")
+
+        image = sample[0, :, 0]
+        image = image.transpose(0, 1).transpose(1, 2)
+        image = (image * 255).numpy().astype(np.uint8)
+        image = Image.fromarray(image)
+        image.save(video_path)
+        print(f"Saved image to: {video_path}")
+    else:
+        video_path = os.path.join(save_path, prefix + ".mp4")
+        save_videos_grid(sample, video_path, fps=fps)
+
+if ulysses_degree * ring_degree > 1:
+    import torch.distributed as dist
+    if dist.get_rank() == 0:
+        save_results()
+else:
+    save_results()
@@ -173,7 +173,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --num_train_epochs=100 \
  --checkpointing_steps=200 \
  --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
@@ -192,8 +192,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --enable_bucket \
  --num_frame_per_block=3 \
  --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --randomize_step_indices \
  --fake_guidance_scale=0.0 \
  --gen_update_interval=5 \
@@ -222,7 +224,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--num_train_epochs` | Number of training epochs | 100 |
 | `--checkpointing_steps` | Save checkpoint every N steps | 200 |
 | `--learning_rate` | Generator learning rate | 2e-06 |
-| `--learning_rate_critic` | Critic learning rate | 2e-06 |
+| `--learning_rate_critic` | Critic learning rate | 4.0e-07 |
 | `--lr_scheduler` | LR scheduler type | `constant_with_warmup` |
 | `--lr_warmup_steps` | LR warmup steps | 100 |
 | `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_dmd` |
@@ -235,6 +237,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--max_grad_norm` | Gradient clipping threshold | 10.0 |
 | `--trainable_modules` | Trainable modules (`"."` = all) | `"."` |
 | `--low_vram` | Enable low VRAM mode (offload VAE/text encoder) | - |
+| `--use_ema` | Maintain an EMA copy of the generator weights | - |

 **Video Sampling Parameters**:

@@ -255,8 +258,8 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh

 | Parameter | Description | Example Value |
 |-----------|-------------|---------------|
-| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Example runs 4-step; parser default `[1000, 500]` = 2-step | `1000 667 334 1` |
-| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 6.0 |
+| `--denoising_step_indices_list` | Denoising step indices (DMD core param). Example runs 4-step; parser default `[1000, 500]` = 2-step | `1000 750 500 250` |
+| `--real_guidance_scale` | CFG scale for the real-score (14B teacher) | 4.0 |
 | `--randomize_step_indices` | Randomize the denoising step indices during training | - |
 | `--fake_guidance_scale` | CFG scale for the fake-score (generator). 0.0 = no CFG | 0.0 |
 | `--gen_update_interval` | Generator update interval (generator updates every N critic steps) | 5 |
@@ -273,7 +276,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh

 ### 5.4 Optional One-Forcing GAN and Precision Controls

-These options reuse the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms in `scripts/wan2.1_self_forcing/train_distill.py`. All GAN/R1/R2 weights default to **0**, and both precision flags default to **off** in argparse. The AR diffusion, CCD, and DMD launchers enable `--fsdp_reduce_fp32`; GAN/R1/R2 and `--score_timestep_clamp` remain off. AR diffusion and CCD do not gain GAN losses.
+These options reuse the [One-Forcing](https://github.com/Aurora-edu/One-Forcing) mechanisms in `scripts/wan2.1_self_forcing/train_distill.py`. All GAN/R1/R2 weights default to **0**, and both precision flags default to **off** in argparse. The AR diffusion, CCD, and DMD launchers enable `--fsdp_reduce_fp32`; GAN/R1/R2 remain off, and the DMD launcher additionally enables `--score_timestep_clamp`. AR diffusion and CCD do not gain GAN losses.

 The discriminator shares the **fake-score backbone**, with one learnable register token per selected feature layer. Its head is optimized by the critic optimizer, not the generator optimizer. Generator loss is `DMD + weighted GAN`; critic loss is `denoising + weighted GAN + R1 + R2`.

@@ -290,7 +293,7 @@ The discriminator shares the **fake-score backbone**, with one learnable registe
 | `--r2_weight` | Normalized finite-difference penalty on noisy generated latents | `0.0` |
 | `--r2_sigma` | R2 perturbation standard deviation | `0.01` |
 | `--fsdp_reduce_fp32` | FSDP FP32 gradient reduction; parameter gathering stays low precision | off (enabled in the launcher) |
-| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off |
+| `--score_timestep_clamp` | Clamp DMD, generator GAN, and critic timesteps to `[0.02, 0.98] * num_train_timesteps` | off (enabled in the DMD launcher) |

 For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` for the generator and `softplus(-D_real) + softplus(D_fake)` for the discriminator. Relativistic losses instead use `softplus(-(D_fake - D_real))` and `softplus(-(D_real - D_fake))`.

@@ -173,7 +173,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --num_train_epochs=100 \
  --checkpointing_steps=200 \
  --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
@@ -192,8 +192,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --enable_bucket \
  --num_frame_per_block=3 \
  --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --randomize_step_indices \
  --fake_guidance_scale=0.0 \
  --gen_update_interval=5 \
@@ -222,7 +224,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--num_train_epochs` | 训练 epoch 数 | 100 |
 | `--checkpointing_steps` | 每 N 步保存 checkpoint | 200 |
 | `--learning_rate` | 生成器学习率 | 2e-06 |
-| `--learning_rate_critic` | 判别器学习率 | 2e-06 |
+| `--learning_rate_critic` | 判别器学习率 | 4.0e-07 |
 | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_causal_forcing_dmd` |
@@ -235,6 +237,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--max_grad_norm` | 梯度裁剪阈值 | 10.0 |
 | `--trainable_modules` | 可训练模块（`"."` = 全部） | `"."` |
 | `--low_vram` | 低显存模式（卸载 VAE/文本编码器） | - |
+| `--use_ema` | 维护生成器权重的 EMA 副本 | - |

 **视频采样参数**：

@@ -255,8 +258,8 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh

 | 参数 | 说明 | 示例值 |
 |------|------|--------|
-| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。示例为 4 步；解析器默认 `[1000, 500]` = 2 步 | `1000 667 334 1` |
-| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 6.0 |
+| `--denoising_step_indices_list` | 去噪步骤索引（DMD 核心参数）。示例为 4 步；解析器默认 `[1000, 500]` = 2 步 | `1000 750 500 250` |
+| `--real_guidance_scale` | real-score（14B teacher）的 CFG scale | 4.0 |
 | `--randomize_step_indices` | 训练时随机化去噪步骤索引 | - |
 | `--fake_guidance_scale` | fake-score（生成器）的 CFG scale。0.0 = 无 CFG | 0.0 |
 | `--gen_update_interval` | 生成器更新间隔（每 N 步判别器更新后更新 1 次生成器） | 5 |
@@ -272,7 +275,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh

 ### 5.4 可选 One-Forcing GAN 与精度控制

-这些选项复用 `scripts/wan2.1_self_forcing/train_distill.py` 中的 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 机制。GAN/R1/R2 权重默认均为 **0**，两个精度开关在 argparse 中默认均**关闭**。AR diffusion、CCD 和 DMD 的 launcher 均启用 `--fsdp_reduce_fp32`；GAN/R1/R2 和 `--score_timestep_clamp` 仍关闭。AR diffusion 和 CCD 不添加 GAN 损失。
+这些选项复用 `scripts/wan2.1_self_forcing/train_distill.py` 中的 [One-Forcing](https://github.com/Aurora-edu/One-Forcing) 机制。GAN/R1/R2 权重默认均为 **0**，两个精度开关在 argparse 中默认均**关闭**。AR diffusion、CCD 和 DMD 的 launcher 均启用 `--fsdp_reduce_fp32`；GAN/R1/R2 仍关闭，DMD 的 launcher 额外启用 `--score_timestep_clamp`。AR diffusion 和 CCD 不添加 GAN 损失。

 判别器共享 **fake-score backbone**，每个选定特征层对应一个可学习 register token。判别器头由 critic optimizer 更新，不加入 generator optimizer。生成器损失为 `DMD + 加权 GAN`，critic 损失为 `denoising + 加权 GAN + R1 + R2`。

@@ -289,7 +292,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh
 | `--r2_weight` | 噪声生成 latent 上的归一化有限差分正则权重 | `0.0` |
 | `--r2_sigma` | R2 扰动标准差 | `0.01` |
 | `--fsdp_reduce_fp32` | FSDP 使用 FP32 梯度归约，参数仍以低精度聚合 | 关闭（launcher 已启用） |
-| `--score_timestep_clamp` | 将 DMD、生成器 GAN 和 critic 时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关闭 |
+| `--score_timestep_clamp` | 将 DMD、生成器 GAN 和 critic 时间步限制到 `[0.02, 0.98] * num_train_timesteps` | 关闭（DMD launcher 已启用） |

 设 logits 为 `D_fake` 和 `D_real`，普通 GAN 的生成器损失为 `softplus(-D_fake)`，判别器损失为 `softplus(-D_real) + softplus(D_fake)`。相对论模式分别使用 `softplus(-(D_fake - D_real))` 和 `softplus(-(D_real - D_fake))`。

@@ -37,6 +37,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_ar_diffusion" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=3e-2 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
@@ -41,6 +41,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --output_dir="output_dir_wan2.1_causal_forcing_ccd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -38,13 +38,14 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --num_train_epochs=100 \
  --checkpointing_steps=200 \
  --learning_rate=2.0e-06 \
-  --learning_rate_critic=2.0e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_causal_forcing_dmd" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
  --adam_weight_decay=0.0 \
  --adam_beta1=0.0 \
  --adam_beta2=0.999 \
@@ -56,8 +57,10 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \
  --enable_bucket \
  --num_frame_per_block=3 \
  --use_kv_cache_training \
-  --denoising_step_indices_list 1000 667 334 1 \
-  --real_guidance_scale=6.0 \
+  --use_ema \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --randomize_step_indices \
  --fake_guidance_scale=0.0 \
  --gen_update_interval=5 \
@@ -0,0 +1,356 @@
+# Wan2.1 One-Forcing: Stable One-Step Autoregressive Video Generation
+
+This document provides the complete workflow for **One-Forcing** on Wan2.1-T2V-1.3B, a recipe ported from [Aurora-edu/One-Forcing](https://github.com/Aurora-edu/One-Forcing).
+
+> **What is One-Forcing?**
+>
+> One-Forcing (*Towards Stable One-Step Autoregressive Video Generation*) is **Causal-Forcing Stage 3 — Distribution Matching Distillation (DMD)** *plus* a **shared noised-latent adversarial critic**. The warm-start generator is produced by the first two Causal-Forcing stages, and One-Forcing then polishes it into a **stable one-step, frame-wise autoregressive** video model:
+>
+> 1. **Stage 1 — AR Diffusion** (`scripts/wan2.1_causal_forcing/train_ar_diffusion.py`): Train the model with **teacher forcing** on clean video latents, learning to denoise chunk-by-chunk (or frame-by-frame) in a causal autoregressive manner. This produces a strong AR backbone. See [README_TRAIN_AR_DIFFUSION.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md).
+> 2. **Stage 2 — Causal Consistency Distillation (CCD)** (`scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py`): Distill the multi-step AR model into a **one-step-per-block** consistency model using an EMA teacher with CFG guidance. Its output checkpoint is the **warm-start generator** for One-Forcing. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md).
+> 3. **Stage 3 — One-Forcing** (`train_one_forcing.py`, this document), which adds two mechanisms on top of DMD:
+>    - **DMD core**: the generator matches the real distribution through a 14B real-score teacher and a 1.3B fake-score critic (`train_causal_dmd.py`).
+>    - **Adversarial critic**: a discriminator head rides on the **fake-score backbone** (one register token per selected feature layer). Both the DMD denoising loss and the adversarial loss operate on the **same noised generated latent**, and the critic uses **one shared noise/timestep** for fake and real samples.
+>    - **One step**: `denoising_step_list = [1000]` — each frame is generated in a single denoising step (the "One" in One-Forcing).
+>
+> Instead of training Stage 1/2 yourself, you can also warm-start from the **official ODE-initialized checkpoint** released by the One-Forcing authors (see Section 3).
+>
+> **The trainer** `scripts/wan2.1_one_forcing/train_one_forcing.py` implements the One-Forcing generator/critic GAN losses, the relativistic option, and finite-difference R1/R2 penalties. For the underlying DMD stage, see [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md).
+
+---
+
+## Table of Contents
+- [1. Prerequisites](#1-prerequisites)
+- [2. Environment Setup](#2-environment-setup)
+- [3. Download Pretrained Models](#3-download-pretrained-models)
+- [4. Prepare Training Data](#4-prepare-training-data)
+  - [4.1 Quick Demo Dataset](#41-quick-demo-dataset)
+  - [4.2 Dataset Structure](#42-dataset-structure)
+  - [4.3 metadata.json Format](#43-metadatajson-format)
+- [5. Training](#5-training)
+  - [5.1 Quick Start](#51-quick-start)
+  - [5.2 Recipe Parameters](#52-recipe-parameters)
+  - [5.3 GAN Prerequisites](#53-gan-prerequisites)
+  - [5.4 Optional Regularizers (Relativistic / R1 / R2)](#54-optional-regularizers-relativistic--r1--r2)
+- [6. Inference](#6-inference)
+  - [6.1 Quick Start](#61-quick-start)
+  - [6.2 1-Step vs 4-Step Selector](#62-1-step-vs-4-step-selector)
+  - [6.3 First-Frame Schedule](#63-first-frame-schedule)
+- [7. Checkpoint Round-Trip](#7-checkpoint-round-trip)
+- [8. Additional Resources](#8-additional-resources)
+
+---
+
+## 1. Prerequisites
+
+One-Forcing is the final polish on top of a distilled generator, so it needs:
+
+1. A **warm-start generator checkpoint** — a distilled/ODE generator that the DMD + adversarial critic then polishes (e.g. a **Causal-Forcing Stage-2 CCD** checkpoint, or the official ODE-initialized checkpoint in Section 3).
+2. A **Wan2.1-T2V-14B** model as the DMD real-score teacher.
+3. **Captioned real videos** — the adversarial critic needs clean real latents, so a prompt-only JSON is **not** sufficient (unlike plain DMD).
+
+---
+
+## 2. Environment Setup
+
+**Method 1: Using requirements.txt**
+
+```bash
+pip install -r requirements.txt
+```
+
+**Method 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
+```
+
+**Method 3: Using Docker**
+
+```bash
+# 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
+```
+
+---
+
+## 3. Download Pretrained Models
+
+One-Forcing requires **two** pretrained models:
+
+- **Wan2.1-T2V-1.3B**: base model for the generator/critic.
+- **Wan2.1-T2V-14B**: the non-causal real-score teacher used by DMD to compute the real distribution score.
+
+```bash
+# Create model directory
+mkdir -p models/Diffusion_Transformer
+
+# Download Wan2.1 T2V 1.3B (student base model)
+modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
+
+# Download Wan2.1 T2V 14B (DMD real-score teacher)
+modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
+```
+
+The official One-Forcing checkpoints are hosted on HuggingFace (`JiaqiFeng/OneForcing`):
+
+```bash
+# Warm-start generator for training (the official ODE-initialized checkpoint),
+# loaded via --ode_transformer_path.
+hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+
+# Trained One-Forcing generator for inference (see Section 6).
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+If you trained Causal-Forcing Stage 1/2 yourself, you can instead point `--ode_transformer_path` at your own Stage-2 CCD checkpoint.
+
+---
+
+## 4. Prepare Training Data
+
+One-Forcing enables the GAN branch by default (`--gan_g_weight=0.03 --gan_d_weight=0.03`), so the data loader always uses an **ImageVideoDataset**: the adversarial critic needs **captioned real videos** to build its clean latents. A prompt-only JSON is **not** sufficient (unlike plain DMD).
+
+### 4.1 Quick Demo Dataset
+
+```bash
+# Download official demo dataset
+modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
+```
+
+### 4.2 Dataset Structure
+
+Same as Causal-Forcing Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md#42-dataset-structure) for details.
+
+### 4.3 metadata.json Format
+
+Each entry needs `file_path`, `type="video"`, and `text`; the launcher uses the `metadata_add_width_height.json` variant, which additionally carries `width`/`height` so bucketing does not need to decode the media files:
+
+```json
+[
+  {
+    "file_path": "train/00000000.mp4",
+    "text": "A young woman gently turns her head to the right, tilting it slightly as she offers a soft greeting...",
+    "type": "video",
+    "width": 832,
+    "height": 480
+  }
+]
+```
+
+> **Note**: Keep enough frames per video for `--video_sample_stride=1` and `--video_sample_n_frames=81`: the GAN branch aborts when the source yields fewer latent frames than `score_num_frames=21` (see Section 5.3).
+
+---
+
+## 5. Training
+
+### 5.1 Quick Start
+
+The ready-to-use launcher is [train_one_forcing.sh](./train_one_forcing.sh). It calls the One-Forcing trainer (`train_one_forcing.py`) with the GAN recipe enabled:
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=1.0e-05 \
+  --learning_rate_critic=1.0e-05 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."
+```
+
+Or simply:
+
+```bash
+bash scripts/wan2.1_one_forcing/train_one_forcing.sh
+```
+
+### 5.2 Recipe Parameters
+
+The One-Forcing recipe sets these launcher flags:
+
+| Parameter | Launcher flag | Value |
+|------------------------|---------------|-------|
+| `denoising_step_list: [1000]` | `--denoising_step_indices_list` | `1000` (one step) |
+| `num_frame_per_block: 1` | `--num_frame_per_block` | `1` (frame-wise) |
+| `num_training_frames: 21` / `score_num_frames: 21` | `--video_sample_n_frames=81` @ 480×832 | 21 latent frames |
+| `guidance_scale: 5.0` | `--real_guidance_scale` | `6.0` (see CFG note below) |
+| `lr: 1e-5` / `lr_critic: 1e-5` | `--learning_rate` / `--learning_rate_critic` | `1.0e-05` |
+| `beta1: 0.0`, `beta2: 0.999`, `weight_decay: 0.01` | `--adam_beta1`, `--adam_beta2`, `--adam_weight_decay` | `0.0`, `0.999`, `0.01` |
+| `ema_weight: 0.99` | `--use_ema`, `--ema_decay` | on, `0.99` |
+| `max_steps: 200` | `--max_train_steps` | `200` |
+| `dfake_gen_update_ratio: 5` | `--gen_update_interval` | `5` |
+| `gan_feature_layers: [21, 29]` | `--gan_feature_layers` | `21 29` |
+| `gan_block_ffn_dim: 2048` | `--gan_block_ffn_dim` | `2048` |
+| `gan_block_num_heads: 12` | `--gan_block_num_heads` | `12` |
+| `gan_g_weight: 0.03`, `gan_d_weight: 0.03` | `--gan_g_weight`, `--gan_d_weight` | `0.03`, `0.03` |
+| `relativistic_discriminator: false` | `--relativistic_discriminator` | omitted (off) |
+| `r1_weight: 0.0`, `r2_weight: 0.0` | `--r1_weight`, `--r2_weight` | omitted (off) |
+| `context_noise: 0` | `--context_noise` | `0` (default) |
+
+> **CFG convention**: this repo computes guidance as `cond + scale * (cond - uncond)`, and this recipe uses `--real_guidance_scale=6.0`.
+
+For the general training/video-sampling parameters shared with the DMD stage, see [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md#52-key-parameters) Section 5.2.
+
+### 5.3 GAN Prerequisites
+
+The GAN branch is validated at startup in `train_one_forcing.py`. It requires:
+
+- **Captioned real videos** (`type="video"` + `file_path` + `text`), not a prompt-only JSON. Enabling any positive GAN/R1/R2 weight switches the loader from `TextDataset` to `ImageVideoDataset`.
+- **`--gradient_accumulation_steps=1`** and **no DeepSpeed** — the trainer raises `ValueError` otherwise.
+- **`--video_sample_stride=1`**: with `stride=2`, an 81-frame source yields ~10 latent frames, which is below `score_num_frames=21` and aborts.
+- **Valid feature-layer indices** for the fake-score backbone. The defaults `21 29` (layers) and `12` (heads) fit Wan2.1-T2V-1.3B; the head count must divide the fake-score hidden dimension.
+
+Trainable fake-score and head master weights stay **FP32**; the forward pass uses the configured autocast/FSDP precision, and logits/softplus/finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient-communication precision only.
+
+### 5.4 Optional Regularizers (Relativistic / R1 / R2)
+
+These regularizers are **off** by default. To experiment, append to the launcher:
+
+```bash
+# Relativistic discriminator: use paired real/fake logit differences
+--relativistic_discriminator
+
+# Finite-difference gradient penalties (APT form, not double-backward)
+--r1_weight=1.0 --r1_sigma=0.1     # on noisy real latents:   r1_weight * mean((D_real' - D_real)^2)
+--r2_weight=0.001 --r2_sigma=0.01  # on noisy gen latents: 0.5 * r2_weight * mean(((D_fake' - D_fake)/sigma)^2)
+```
+
+For logits `D_fake` and `D_real`, ordinary GAN losses are `softplus(-D_fake)` (generator) and `softplus(-D_real) + softplus(D_fake)` (discriminator); relativistic losses use the logit differences instead. Real/fake samples share timestep and noise; perturbations use independent noise.
+
+---
+
+## 6. Inference
+
+### 6.1 Quick Start
+
+The inference entry is [examples/wan2.1_one_forcing/predict_t2v.py](../../examples/wan2.1_one_forcing/predict_t2v.py). It loads the **1.3B student** (the 14B model is only the DMD teacher at training time) and runs frame-wise one-step generation. By default it uses the **official One-Forcing checkpoint**:
+
+```bash
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+```python
+# examples/wan2.1_one_forcing/predict_t2v.py
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+use_ema             = False   # only applies to trainer-output checkpoint dirs
+stochastic_sampling = True    # DMD-distilled student
+guidance_scale      = 1.0     # CFG is baked into the distilled weights
+num_inference_steps = 1       # One-Forcing one-step generation
+num_frame_per_block = 1       # frame-wise
+sample_size         = [480, 832]
+video_length        = 81
+```
+
+Run it:
+
+```bash
+python examples/wan2.1_one_forcing/predict_t2v.py
+```
+
+`transformer_path` may be the official checkpoint file, a plain `.pt`/`.safetensors` file, or a trainer-output **checkpoint dir** (its live/EMA weights are auto-resolved). The GAN discriminator head is a training-only submodule of fake_score, so inference needs **no discriminator wiring**.
+
+### 6.2 1-Step vs 4-Step Selector
+
+With `stochastic_sampling=True`, `num_inference_steps` selects the warped timestep schedule:
+
+| `num_inference_steps` | Denoise timesteps | Use case |
+|-----------------------|-------------------|----------|
+| `1` (One-Forcing default) | `[1000]` | Fastest; frame-wise one-step generation |
+| `4` | `[1000, 750, 500, 250]` | More stable, slower |
+
+### 6.3 First-Frame Schedule
+
+The shared `WanSelfForcingPipeline` applies **one denoising schedule to all blocks**. A separate first-frame schedule (e.g. `[1000, 750, 500, 250]` for the first frame and `[1000]` for later frames) is therefore not applied; the pipeline is shared by multiple examples and is intentionally left unmodified.
+
+---
+
+## 7. Checkpoint Round-Trip
+
+The FSDP trainer writes each `checkpoint-{step}/` as follows, so training output feeds inference with no key surgery:
+
+| Component | Saved to | Contains `gan_head`? |
+|-----------|----------|----------------------|
+| **Generator** (live) | `checkpoint-{step}/diffusion_pytorch_model.safetensors` (root, via FSDP hook) | No |
+| **Generator** (EMA) | `checkpoint-{step}/transformer_ema/` | No |
+| **Fake-score + discriminator head** | `checkpoint-{step}/fake_score/` | Yes (FP32 head tensors) |
+
+- The **generator never contains `gan_head`** keys (the head lives on fake-score), so `predict_t2v.py` loads it cleanly with `strict=False`.
+- `_resolve_transformer_path` in the inference script resolves, in order: `transformer_ema/`, `ema_transformer/` (compat alias), `transformer/`, then the root safetensors.
+- **Resume** GAN training with the same GAN/head configuration: the head is recreated from CLI flags before loading and restored from `fake_score/`.
+
+---
+
+## 8. Additional Resources
+
+- **One-Forcing (source recipe)**: https://github.com/Aurora-edu/One-Forcing
+- **Causal-Forcing (DMD stage)**: https://github.com/thu-ml/Causal-Forcing
+- **Stage 1 (AR Diffusion) training guide**: [README_TRAIN_AR_DIFFUSION.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md)
+- **Stage 2 (CCD) training guide**: [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md)
+- **Underlying DMD trainer guide**: [README_TRAIN_CAUSAL_DMD_DISTILL.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md)
+- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
@@ -0,0 +1,355 @@
+# Wan2.1 One-Forcing：稳定的单步自回归视频生成
+
+本文档提供在 Wan2.1-T2V-1.3B 上进行 **One-Forcing** 的完整工作流，该配方移植自 [Aurora-edu/One-Forcing](https://github.com/Aurora-edu/One-Forcing)。
+
+> **什么是 One-Forcing？**
+>
+> One-Forcing（*Towards Stable One-Step Autoregressive Video Generation*）= **Causal-Forcing Stage 3 — 分布匹配蒸馏 (DMD)** *加上*一个**共享噪声 latent 的对抗判别器（critic）**。热启动生成器由 Causal-Forcing 的前两个阶段产出，One-Forcing 再将其打磨为**稳定的单步、逐帧自回归**视频模型：
+>
+> 1. **Stage 1 — AR Diffusion**（`scripts/wan2.1_causal_forcing/train_ar_diffusion.py`）：在干净视频 latent 上以 **teacher forcing** 训练模型，学习按块（或按帧）因果自回归地去噪，产出强 AR 基座。参见 [README_TRAIN_AR_DIFFUSION_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md)。
+> 2. **Stage 2 — 因果一致性蒸馏 (CCD)**（`scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py`）：使用 EMA teacher + CFG 将多步 AR 模型蒸馏为**每块一步**的一致性模型，其产出 checkpoint 即 One-Forcing 的**热启动生成器**。参见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)。
+> 3. **Stage 3 — One-Forcing**（`train_one_forcing.py`，即本文档），在 DMD 之上增加两个机制：
+>    - **DMD 核心**：生成器通过 14B real-score teacher 与 1.3B fake-score critic 匹配真实分布（`train_causal_dmd.py`）。
+>    - **对抗 critic**：判别器头挂在 **fake-score backbone** 上（每个选定特征层对应一个 register token）。DMD 去噪损失与对抗损失作用在**同一个噪声生成 latent** 上，critic 对真实/生成样本使用**同一份共享噪声与时间步**。
+>    - **单步**：`denoising_step_list = [1000]` —— 每帧在单个去噪步内生成（即 One-Forcing 中的 “One”）。
+>
+> 你也可以不自己训练 Stage 1/2，而是直接使用 One-Forcing 作者发布的**官方 ODE 初始化 checkpoint** 作为热启动权重（见第三节）。
+>
+> **训练器** `scripts/wan2.1_one_forcing/train_one_forcing.py` 实现了 One-Forcing 的生成器/critic GAN 损失、相对论（relativistic）选项以及有限差分 R1/R2 正则。底层 DMD 阶段请参见 [README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md)。
+
+---
+
+## 目录
+- [一、前置条件](#一前置条件)
+- [二、环境配置](#二环境配置)
+- [三、下载预训练模型](#三下载预训练模型)
+- [四、数据准备](#四数据准备)
+  - [4.1 快速测试数据集](#41-快速测试数据集)
+  - [4.2 数据集结构](#42-数据集结构)
+  - [4.3 metadata.json 格式](#43-metadatajson-格式)
+- [五、训练](#五训练)
+  - [5.1 快速开始](#51-快速开始)
+  - [5.2 配方参数](#52-配方参数)
+  - [5.3 GAN 前置要求](#53-gan-前置要求)
+  - [5.4 可选正则（相对论 / R1 / R2）](#54-可选正则相对论--r1--r2)
+- [六、推理](#六推理)
+  - [6.1 快速开始](#61-快速开始)
+  - [6.2 单步与 4 步选择器](#62-单步与-4-步选择器)
+  - [6.3 首帧调度](#63-首帧调度)
+- [七、Checkpoint 往返（训练→推理）](#七checkpoint-往返训练推理)
+- [八、更多资源](#八更多资源)
+
+---
+
+## 一、前置条件
+
+One-Forcing 是在已蒸馏生成器之上的最后打磨，因此需要：
+
+1. **热启动生成器 checkpoint** —— 一个已蒸馏/ODE 生成器，由 DMD + 对抗 critic 进一步打磨（例如 **Causal-Forcing Stage-2 CCD** checkpoint，或第三节的官方 ODE 初始化 checkpoint）。
+2. **Wan2.1-T2V-14B** 模型，作为 DMD real-score teacher。
+3. **带描述的真实视频** —— 对抗 critic 需要干净的真实 latent，因此仅有提示词的 JSON **不够用**（这一点与普通 DMD 不同）。
+
+---
+
+## 二、环境配置
+
+**方式 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**
+
+```bash
+# Pull the image
+docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
+
+# Enter the container
+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
+```
+
+---
+
+## 三、下载预训练模型
+
+One-Forcing 需要**两个**预训练模型：
+
+- **Wan2.1-T2V-1.3B**：生成器/判别器的基础模型。
+- **Wan2.1-T2V-14B**：非因果的 real-score teacher，用于计算真实分布得分。
+
+```bash
+# Create the model directory
+mkdir -p models/Diffusion_Transformer
+
+# Download the Wan2.1 T2V 1.3B student base model
+modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B
+
+# Download the Wan2.1 T2V 14B real-score teacher
+modelscope download --model Wan-AI/Wan2.1-T2V-14B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-14B
+```
+
+One-Forcing 官方 checkpoint 托管在 HuggingFace（`JiaqiFeng/OneForcing`）：
+
+```bash
+# 训练用热启动生成器（官方 ODE 初始化 checkpoint），通过 --ode_transformer_path 加载。
+hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+
+# 推理用已训练 One-Forcing 生成器（见第六节）。
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+如果你自己训练了 Causal-Forcing Stage 1/2，也可以将 `--ode_transformer_path` 指向你产出的 Stage-2 CCD checkpoint。
+
+---
+
+## 四、数据准备
+
+One-Forcing 默认开启 GAN 分支（`--gan_g_weight=0.03 --gan_d_weight=0.03`），因此数据加载始终使用 **ImageVideoDataset**：对抗 critic 需要**带描述的真实视频**来构建干净 latent，仅有提示词的 JSON **不够用**（这一点与普通 DMD 不同）。
+
+### 4.1 快速测试数据集
+
+```bash
+# Download the official demo dataset
+modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo
+```
+
+### 4.2 数据集结构
+
+与 Causal-Forcing Stage 1/2 相同。详见 [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md#42-数据集结构)。
+
+### 4.3 metadata.json 格式
+
+每条记录需要 `file_path`、`type="video"` 与 `text`；启动脚本使用 `metadata_add_width_height.json` 变体，额外携带 `width`/`height`，分桶训练无需现场解码媒体文件：
+
+```json
+[
+  {
+    "file_path": "train/00000000.mp4",
+    "text": "A young woman gently turns her head to the right, tilting it slightly as she offers a soft greeting...",
+    "type": "video",
+    "width": 832,
+    "height": 480
+  }
+]
+```
+
+> **说明**：每条视频须有足够的帧数以满足 `--video_sample_stride=1` 与 `--video_sample_n_frames=81`：当源视频产出的 latent 帧数少于 `score_num_frames=21` 时，GAN 分支会中止（见 5.3）。
+
+---
+
+## 五、训练
+
+### 5.1 快速开始
+
+可直接使用的启动脚本为 [train_one_forcing.sh](./train_one_forcing.sh)。它调用 One-Forcing 训练器（`train_one_forcing.py`）并开启 GAN 配方：
+
+```bash
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=1.0e-05 \
+  --learning_rate_critic=1.0e-05 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."
+```
+
+或直接运行：
+
+```bash
+bash scripts/wan2.1_one_forcing/train_one_forcing.sh
+```
+
+### 5.2 配方参数
+
+One-Forcing 配方设置以下 launcher 参数：
+
+| 参数 | Launcher 参数 | 取值 |
+|--------------------|---------------|------|
+| `denoising_step_list: [1000]` | `--denoising_step_indices_list` | `1000`（单步） |
+| `num_frame_per_block: 1` | `--num_frame_per_block` | `1`（逐帧） |
+| `num_training_frames: 21` / `score_num_frames: 21` | `--video_sample_n_frames=81` @ 480×832 | 21 个 latent 帧 |
+| `guidance_scale: 5.0` | `--real_guidance_scale` | `6.0`（见下方 CFG 说明） |
+| `lr: 1e-5` / `lr_critic: 1e-5` | `--learning_rate` / `--learning_rate_critic` | `1.0e-05` |
+| `beta1: 0.0`、`beta2: 0.999`、`weight_decay: 0.01` | `--adam_beta1`、`--adam_beta2`、`--adam_weight_decay` | `0.0`、`0.999`、`0.01` |
+| `ema_weight: 0.99` | `--use_ema`、`--ema_decay` | 开启、`0.99` |
+| `max_steps: 200` | `--max_train_steps` | `200` |
+| `dfake_gen_update_ratio: 5` | `--gen_update_interval` | `5` |
+| `gan_feature_layers: [21, 29]` | `--gan_feature_layers` | `21 29` |
+| `gan_block_ffn_dim: 2048` | `--gan_block_ffn_dim` | `2048` |
+| `gan_block_num_heads: 12` | `--gan_block_num_heads` | `12` |
+| `gan_g_weight: 0.03`、`gan_d_weight: 0.03` | `--gan_g_weight`、`--gan_d_weight` | `0.03`、`0.03` |
+| `relativistic_discriminator: false` | `--relativistic_discriminator` | 不传（关闭） |
+| `r1_weight: 0.0`、`r2_weight: 0.0` | `--r1_weight`、`--r2_weight` | 不传（关闭） |
+| `context_noise: 0` | `--context_noise` | `0`（默认） |
+
+> **CFG 约定**：本仓库的引导计算为 `cond + scale * (cond - uncond)`，本配方取 `--real_guidance_scale=6.0`。
+
+与 DMD 阶段共用的通用训练/视频采样参数，参见 [README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md#52-关键参数) 5.2 节。
+
+### 5.3 GAN 前置要求
+
+GAN 分支在 `train_one_forcing.py` 启动时会被校验，要求：
+
+- **带描述的真实视频**（`type="video"` + `file_path` + `text`），不能只提供提示词 JSON。任一 GAN/R1/R2 权重为正时，数据加载会从 `TextDataset` 切换为 `ImageVideoDataset`。
+- **`--gradient_accumulation_steps=1`** 且**不使用 DeepSpeed** —— 否则训练器抛出 `ValueError`。
+- **`--video_sample_stride=1`**：当 `stride=2` 时，81 帧源仅得到约 10 个 latent 帧，低于 `score_num_frames=21` 而中止。
+- fake-score backbone 上**有效的特征层索引**。默认 `21 29`（层）与 `12`（头）适用于 Wan2.1-T2V-1.3B；头数须整除 fake-score 隐藏维度。
+
+可训练的 fake-score 与判别器头 master weights 保持 **FP32**；前向按配置使用 autocast/FSDP 混合精度，logits/softplus/有限差分归约在 FP32 下计算。`--fsdp_reduce_fp32` 只改变梯度通信精度。
+
+### 5.4 可选正则（相对论 / R1 / R2）
+
+这些正则默认**关闭**。若要实验，向 launcher 追加：
+
+```bash
+# Relativistic discriminator: use paired real/fake logit differences
+--relativistic_discriminator
+
+# Finite-difference gradient penalties (APT form, not double-backward)
+--r1_weight=1.0 --r1_sigma=0.1     # on noisy real latents:   r1_weight * mean((D_real' - D_real)^2)
+--r2_weight=0.001 --r2_sigma=0.01  # on noisy gen latents: 0.5 * r2_weight * mean(((D_fake' - D_fake)/sigma)^2)
+```
+
+设 logits 为 `D_fake` 与 `D_real`，普通 GAN 损失为 `softplus(-D_fake)`（生成器）与 `softplus(-D_real) + softplus(D_fake)`（判别器）；相对论模式改用 logit 差值。真实/生成样本共用时间步与噪声，扰动噪声独立抽样。
+
+---
+
+## 六、推理
+
+### 6.1 快速开始
+
+推理入口为 [examples/wan2.1_one_forcing/predict_t2v.py](../../examples/wan2.1_one_forcing/predict_t2v.py)。它加载 **1.3B 学生模型**（14B 仅在训练时作为 DMD teacher），运行逐帧单步生成。默认使用 **One-Forcing 官方 checkpoint**：
+
+```bash
+hf download JiaqiFeng/OneForcing checkpoints/one_forcing.pt --local-dir models/Diffusion_Transformer/One-Forcing
+```
+
+```python
+# examples/wan2.1_one_forcing/predict_t2v.py
+model_name          = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B"
+transformer_path    = "models/Diffusion_Transformer/One-Forcing/checkpoints/one_forcing.pt"
+use_ema             = False   # only applies to trainer-output checkpoint dirs
+stochastic_sampling = True    # DMD-distilled student
+guidance_scale      = 1.0     # CFG is baked into the distilled weights
+num_inference_steps = 1       # One-Forcing one-step generation
+num_frame_per_block = 1       # frame-wise
+sample_size         = [480, 832]
+video_length        = 81
+```
+
+运行：
+
+```bash
+python examples/wan2.1_one_forcing/predict_t2v.py
+```
+
+`transformer_path` 既可以是官方 checkpoint 文件、纯 `.pt`/`.safetensors` 文件，也可以是训练输出的 **checkpoint 目录**（其 live/EMA 权重会被自动解析）。GAN 判别器头是 fake_score 的训练专用子模块，因此推理**无需任何判别器接线**。
+
+### 6.2 单步与 4 步选择器
+
+当 `stochastic_sampling=True` 时，`num_inference_steps` 决定 warped 时间步调度：
+
+| `num_inference_steps` | 去噪时间步 | 适用场景 |
+|-----------------------|------------|----------|
+| `1`（One-Forcing 默认） | `[1000]` | 最快；逐帧单步生成 |
+| `4` | `[1000, 750, 500, 250]` | 更稳定、更慢 |
+
+### 6.3 首帧调度
+
+共享的 `WanSelfForcingPipeline` 对**所有块使用同一份去噪调度**，因此此处不单独为首帧设置调度（例如首帧 `[1000, 750, 500, 250]`、后续帧 `[1000]`）；该 pipeline 被多个 example 共用，有意保持不变。
+
+---
+
+## 七、Checkpoint 往返（训练→推理）
+
+FSDP 训练器将每个 `checkpoint-{step}/` 按如下布局写出，因此训练产物可直接喂给推理，无需改动 key：
+
+| 组件 | 保存位置 | 是否含 `gan_head`？ |
+|------|----------|---------------------|
+| **生成器**（live） | `checkpoint-{step}/diffusion_pytorch_model.safetensors`（根目录，经 FSDP hook） | 否 |
+| **生成器**（EMA） | `checkpoint-{step}/transformer_ema/` | 否 |
+| **Fake-score + 判别器头** | `checkpoint-{step}/fake_score/` | 是（FP32 判别器头张量） |
+
+- **生成器从不包含 `gan_head`** 键（判别器头挂在 fake-score 上），因此 `predict_t2v.py` 可用 `strict=False` 干净加载。
+- 推理脚本中的 `_resolve_transformer_path` 依次解析：`transformer_ema/`、`ema_transformer/`（兼容别名）、`transformer/`，最后到根目录 safetensors。
+- **恢复** GAN 训练时须使用相同的 GAN/判别器头配置：加载前会根据 CLI 参数重建判别器头，并从 `fake_score/` 还原。
+
+---
+
+## 八、更多资源
+
+- **One-Forcing（源配方）**：https://github.com/Aurora-edu/One-Forcing
+- **Causal-Forcing（DMD 阶段）**：https://github.com/thu-ml/Causal-Forcing
+- **Stage 1（AR Diffusion）训练指南**：[README_TRAIN_AR_DIFFUSION_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md)
+- **Stage 2（CCD）训练指南**：[README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md)
+- **底层 DMD 训练器指南**：[README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md](../wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md)
+- **官方 GitHub**：https://github.com/aigc-apps/VideoX-Fun
@@ -0,0 +1,85 @@
+export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
+# One-Forcing needs REAL videos to build the adversarial critic's clean_latents, so the dataset
+# must be a video meta (type=video + file_path), NOT a prompt-only json. The official demo
+# dataset X-Fun-Videos-Demo ships captioned videos and is normalized by --enable_bucket
+# --random_hw_adapt --training_with_video_token_length.
+export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
+export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
+# Warm-start generator: a distilled/ODE checkpoint that the DMD + adversarial critic then polish.
+# Default: the official One-Forcing ODE-initialized checkpoint, downloaded via
+#   hf download JiaqiFeng/OneForcing checkpoints/framewise/causal_ode.pt --local-dir models/Diffusion_Transformer/One-Forcing
+# You can also point it at your own Causal-Forcing Stage-2 CCD checkpoint.
+export STAGE2_CKPT="models/Diffusion_Transformer/One-Forcing/checkpoints/framewise/causal_ode.pt"
+# 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
+
+# One-Forcing = Causal-Forcing Stage-3 DMD + a shared noised-latent adversarial critic on fake_score.
+# - 1-step frame-wise generation: --denoising_step_indices_list 1000, --num_frame_per_block=1;
+#   21 latent frames from --video_sample_n_frames=81 @ 480x832.
+# - GAN critic: gan_g_weight/gan_d_weight=0.03, feature layers [21, 29], ffn 2048, heads 12,
+#   non-relativistic, R1/R2 off. The head rides on fake_score and is training-only.
+# - GAN requires real videos, --gradient_accumulation_steps=1, no DeepSpeed, and --video_sample_stride=1
+#   (stride=2 on 81-frame sources yields ~10 latent frames < score_num_frames=21 and aborts).
+# - guidance: local CFG is cond + scale*(cond - uncond), so --real_guidance_scale=6.0.
+# - optimizer/EMA/steps: lr=lr_critic=1e-5, betas (0.0, 0.999), weight_decay=0.01,
+#   --use_ema --ema_decay=0.99, gen_update_interval=5, max_steps=200.
+accelerate launch --mixed_precision="bf16" --use_fsdp \
+    --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \
+    --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock \
+    --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT \
+    --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False \
+    scripts/wan2.1_one_forcing/train_one_forcing.py \
+  --config_path="config/wan2.1/wan_civitai.yaml" \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_pretrained_model_name_or_path=$REAL_SCORE_MODEL_NAME \
+  --train_data_dir=$DATASET_NAME \
+  --train_data_meta=$DATASET_META_NAME \
+  --ode_transformer_path=$STAGE2_CKPT \
+  --image_sample_size=640 \
+  --video_sample_size=640 \
+  --token_sample_size=640 \
+  --fix_sample_size 480 832 \
+  --video_sample_stride=1 \
+  --video_sample_n_frames=81 \
+  --train_batch_size=1 \
+  --gradient_accumulation_steps=1 \
+  --dataloader_num_workers=8 \
+  --num_train_epochs=100 \
+  --max_train_steps=200 \
+  --checkpointing_steps=100 \
+  --learning_rate=1.0e-05 \
+  --learning_rate_critic=1.0e-05 \
+  --lr_scheduler="constant" \
+  --seed=42 \
+  --output_dir="output_dir_wan2.1_one_forcing_framewise_1step" \
+  --gradient_checkpointing \
+  --mixed_precision="bf16" \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.01 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
+  --adam_epsilon=1e-08 \
+  --vae_mini_batch=1 \
+  --max_grad_norm=10.0 \
+  --random_hw_adapt \
+  --training_with_video_token_length \
+  --enable_bucket \
+  --num_frame_per_block=1 \
+  --use_kv_cache_training \
+  --use_ema \
+  --ema_decay=0.99 \
+  --denoising_step_indices_list 1000 \
+  --real_guidance_scale=6.0 \
+  --score_timestep_clamp \
+  --fake_guidance_scale=0.0 \
+  --gen_update_interval=5 \
+  --gan_g_weight=0.03 \
+  --gan_d_weight=0.03 \
+  --gan_feature_layers 21 29 \
+  --gan_block_ffn_dim=2048 \
+  --gan_block_num_heads=12 \
+  --resume_from_checkpoint="latest" \
+  --trainable_modules "."
@@ -179,6 +179,7 @@ The difference between DeepSpeed-Zero-2 and FSDP lies in whether to shard model

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -189,6 +190,7 @@ NCCL_DEBUG=INFO
 accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -204,21 +206,26 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -247,17 +254,19 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--num_train_epochs` | Number of training epochs | 100 |
 | `--checkpointing_steps` | Save checkpoint every N steps | 50 |
 | `--learning_rate` | Initial learning rate (generator) | 2e-06 |
-| `--learning_rate_critic` | Initial learning rate (critic) | 2e-06 |
+| `--learning_rate_critic` | Initial learning rate (critic) | 4.0e-07 |
 | `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` |
 | `--lr_warmup_steps` | Learning rate warmup steps | 100 |
 | `--seed` | Random seed | 42 |
 | `--output_dir` | Output directory | `output_dir_wan2.1_self_forcing_distill` |
 | `--gradient_checkpointing` | Enable gradient checkpointing | - |
 | `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
-| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
+| `--adam_weight_decay` | AdamW weight decay | 0.0 |
+| `--adam_beta1` | AdamW beta1 | 0.0 |
+| `--adam_beta2` | AdamW beta2 | 0.999 |
 | `--adam_epsilon` | AdamW epsilon | 1e-10 |
 | `--vae_mini_batch` | VAE encoding mini-batch size | 1 |
-| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
+| `--max_grad_norm` | Gradient clipping threshold | 10.0 |
 | `--fsdp_reduce_fp32` | Optional: reduce gradients in fp32 under FSDP mixed precision (default reduces in bf16) | - |
 | `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - |
 | `--random_hw_adapt` | Auto-scale images/videos to random sizes in `[min_size, max_size]` range | - |
@@ -277,7 +286,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 |-----|------|-------|
 | `--denoising_step_indices_list` | Denoising step indices list (core distillation parameter) | `1000 750 500 250` |
 | `--randomize_step_indices` | Randomize the denoising timestep indices during training | - |
-| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 |
+| `--real_guidance_scale` | Real guidance scale for scoring | 4.0 |
 | `--score_timestep_clamp` | Optional: clamp DMD/critic timesteps to `[0.02, 0.98] * num_train_timesteps` | - |
 | `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 |
 | `--gen_update_interval` | Generator update interval | 5 |
@@ -365,6 +374,7 @@ You can configure validation parameters to periodically generate test videos dur

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -375,6 +385,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -390,21 +401,26 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -427,6 +443,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/
 Training shell command is as follows:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -437,6 +454,7 @@ NCCL_DEBUG=INFO
 accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -452,21 +470,26 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -481,6 +504,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -491,6 +515,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -506,21 +531,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -540,6 +570,7 @@ Assuming 2 machines, each with 8 GPUs:
 **Machine 0 (Master)**:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Master machine IP
@@ -555,6 +586,7 @@ NCCL_DEBUG=INFO
 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 --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -570,21 +602,26 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -596,6 +633,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
 **Machine 1 (Worker)**:
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Same as Master
@@ -180,6 +180,7 @@ hf download gdhe17/Self-Forcing --local-dir models/Diffusion_Transformer/Self-Fo

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -190,6 +191,7 @@ NCCL_DEBUG=INFO
 accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -205,21 +207,26 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -248,17 +255,19 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 | `--num_train_epochs` | 训练 epoch 数 | 100 |
 | `--checkpointing_steps` | 每 N 步保存 checkpoint | 50 |
 | `--learning_rate` | 初始学习率（生成器） | 2e-06 |
-| `--learning_rate_critic` | 初始学习率（判别器） | 2e-06 |
+| `--learning_rate_critic` | 初始学习率（判别器） | 4.0e-07 |
 | `--lr_scheduler` | 学习率调度器 | `constant_with_warmup` |
 | `--lr_warmup_steps` | 学习率预热步数 | 100 |
 | `--seed` | 随机种子 | 42 |
 | `--output_dir` | 输出目录 | `output_dir_wan2.1_self_forcing_distill` |
 | `--gradient_checkpointing` | 激活重计算 | - |
 | `--mixed_precision` | 混合精度：`fp16/bf16` | `bf16` |
-| `--adam_weight_decay` | AdamW 权重衰减 | 3e-2 |
+| `--adam_weight_decay` | AdamW 权重衰减 | 0.0 |
+| `--adam_beta1` | AdamW beta1 | 0.0 |
+| `--adam_beta2` | AdamW beta2 | 0.999 |
 | `--adam_epsilon` | AdamW epsilon 值 | 1e-10 |
 | `--vae_mini_batch` | VAE 编码时的迷你批次大小 | 1 |
-| `--max_grad_norm` | 梯度裁剪阈值 | 0.05 |
+| `--max_grad_norm` | 梯度裁剪阈值 | 10.0 |
 | `--fsdp_reduce_fp32` | 可选：FSDP 混精下梯度以 fp32 做跨卡 reduce（默认 bf16） | - |
 | `--enable_bucket` | 启用分桶训练，不裁剪图片/视频，按分辨率分组训练 | - |
 | `--random_hw_adapt` | 自动缩放图片/视频到 `[min_size, max_size]` 范围内的随机尺寸 | - |
@@ -278,7 +287,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con
 |-----|------|-------|
 | `--denoising_step_indices_list` | 去噪步骤列表（蒸馏核心参数） | `1000 750 500 250` |
 | `--randomize_step_indices` | 训练时随机化去噪 timestep 索引 | - |
-| `--real_guidance_scale` | 用于评分的真实 guidance scale | 6.0 |
+| `--real_guidance_scale` | 用于评分的真实 guidance scale | 4.0 |
 | `--score_timestep_clamp` | 可选：将 DMD/critic 的 timestep 截断到 `[0.02, 0.98] * num_train_timesteps` | - |
 | `--fake_guidance_scale` | 用于评分的虚拟 guidance scale | 0.0 |
 | `--gen_update_interval` | 生成器更新间隔 | 5 |
@@ -366,6 +375,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -376,6 +386,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=CasualWanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -391,21 +402,26 @@ accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TR
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -428,6 +444,7 @@ python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/
 训练 shell 命令如下：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -438,6 +455,7 @@ NCCL_DEBUG=INFO
 accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -453,21 +471,26 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -482,6 +505,7 @@ accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag

 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 # NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA. 
@@ -492,6 +516,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -507,21 +532,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -541,6 +571,7 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
 **机器 0（Master）**：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # Master 机器 IP
@@ -556,6 +587,7 @@ NCCL_DEBUG=INFO
 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 --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -571,21 +603,26 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -597,6 +634,7 @@ accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main
 **机器 1（Worker）**：
 ```bash
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 export DATASET_NAME="datasets/X-Fun-Videos-Demo/"
 export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json"
 export MASTER_ADDR="192.168.1.100"  # 与 Master 相同
@@ -1,4 +1,5 @@
 export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B/"
+export REAL_SCORE_MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-14B"
 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. 
@@ -9,6 +10,7 @@ NCCL_DEBUG=INFO
 accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_distill.py \
  --config_path="config/wan2.1/wan_civitai.yaml" \
  --pretrained_model_name_or_path=$MODEL_NAME \
+  --real_score_model_name_or_path=$REAL_SCORE_MODEL_NAME \
  --train_data_dir=$DATASET_NAME \
  --train_data_meta=$DATASET_META_NAME \
  --image_sample_size=640 \
@@ -24,21 +26,26 @@ accelerate launch --mixed_precision="bf16" scripts/wan2.1_self_forcing/train_dis
  --num_train_epochs=100 \
  --checkpointing_steps=50 \
  --learning_rate=2e-06 \
-  --learning_rate_critic=2e-06 \
+  --learning_rate_critic=4.0e-07 \
  --lr_scheduler="constant_with_warmup" \
  --lr_warmup_steps=100 \
  --seed=42 \
  --output_dir="output_dir_wan2.1_self_forcing_distill" \
  --gradient_checkpointing \
  --mixed_precision="bf16" \
-  --adam_weight_decay=3e-2 \
+  --fsdp_reduce_fp32 \
+  --adam_weight_decay=0.0 \
+  --adam_beta1=0.0 \
+  --adam_beta2=0.999 \
  --adam_epsilon=1e-10 \
  --vae_mini_batch=1 \
-  --max_grad_norm=0.05 \
+  --max_grad_norm=10.0 \
  --random_hw_adapt \
  --training_with_video_token_length \
  --enable_bucket \
-  --uniform_sampling \
+  --denoising_step_indices_list 1000 750 500 250 \
+  --real_guidance_scale=4.0 \
+  --score_timestep_clamp \
  --use_kv_cache_training \
  --num_frame_per_block=3 \
  --train_mode="normal" \
@@ -5,6 +5,7 @@ import warnings

 import torch
 import torch.nn as nn
+import torch.nn.functional as F

 try:
    import flash_attn_interface
@@ -26,6 +26,29 @@ DETERMINISTIC_LOGVAR = -30.0
 TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))


+def _load_raw_state_dict(path):
+    """Load a flat tensor state_dict from a .safetensors/.pth/.bin checkpoint.
+
+    Mirrors the normalization done in AutoencoderTinyWan.from_pretrained so that a
+    companion encoder checkpoint can be borrowed the same way.
+    """
+    path = os.fspath(path)
+    if path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+
+        sd = load_file(path)
+    else:
+        sd = torch.load(path, map_location="cpu", weights_only=True)
+    if isinstance(sd, dict):
+        if isinstance(sd.get("state_dict"), dict):
+            sd = sd["state_dict"]
+        elif isinstance(sd.get("decoder"), dict):
+            sd = sd["decoder"]
+    if isinstance(sd, dict) and all(k.startswith("model.") for k in sd):
+        sd = {k[len("model.") :]: v for k, v in sd.items()}
+    return sd
+
+
 def conv(n_in, n_out, **kwargs):
    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)

@@ -566,6 +589,10 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
        if spatial_compression_ratio is None:
            spatial_compression_ratio = 8 * patch_size
        self.spatial_compression_ratio = spatial_compression_ratio
+        # Persist the derived value into the frozen config as well: some pipelines
+        # read `vae.config.spatial_compression_ratio` (not the attribute), and the
+        # register_to_config decorator would otherwise keep the `None` default.
+        self.register_to_config(spatial_compression_ratio=spatial_compression_ratio)

        self.model = TAEWan(
            encoder_time_downscale=tuple(encoder_time_downscale),
@@ -696,6 +723,33 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
        if all(k.startswith("model.") for k in state_dict):
            state_dict = {k[len("model.") :]: v for k, v in state_dict.items()}

+        # Borrow an encoder when the checkpoint is decoder-only. Some taew2_2_super
+        # releases ship only the Super decoder; without an encoder, encode() would
+        # raise. We merge encoder.* weights from a companion base checkpoint so the
+        # model becomes a full encoder(base)+decoder(super) TAE.
+        # Priority: explicit `encoder_path` (via additional_kwargs / vae_kwargs or
+        # kwargs) > auto sibling found by stripping "_super" from the filename.
+        encoder_path = None
+        if isinstance(additional_kwargs, dict):
+            encoder_path = additional_kwargs.get("encoder_path")
+        if encoder_path is None:
+            encoder_path = kwargs.get("encoder_path")
+        if not any(k.startswith("encoder.") for k in state_dict):
+            if encoder_path is None:
+                root, ext = os.path.splitext(pretrained_model_path)
+                sibling = root.replace("_super", "") + ext
+                if sibling != pretrained_model_path and os.path.exists(sibling):
+                    encoder_path = sibling
+            if encoder_path is not None and os.path.exists(encoder_path):
+                enc_sd = {k: v for k, v in _load_raw_state_dict(encoder_path).items() if k.startswith("encoder.")}
+                if enc_sd:
+                    state_dict = {**enc_sd, **state_dict}
+                    print(f"[TAE] decoder-only 权重：已从 {encoder_path} 借用 {len(enc_sd)} 个 encoder 键以支持 encode。")
+            if not any(k.startswith("encoder.") for k in state_dict):
+                print("[TAE][warn] 未找到 encoder 权重，将以 decoder_only 加载，encode() 不可用。"
+                      "可通过 vae_kwargs.encoder_path=<base taew2_2 权重路径> 指定，"
+                      "或把 base 权重放在与 super 同名（去掉 _super）的路径下。")
+
        if "decoder.0.weight" in state_dict and "decoder.1.weight" not in state_dict:
            arch_variant, first_key, block_key = "super", "decoder.0.weight", "decoder.1.conv.0.weight"
            first_channels, first_kernel, block_shape = 512, (1, 1), (1024, 1, 7, 7)
@@ -23,15 +23,19 @@ from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


-def stochastic_sampling_timesteps(num_inference_steps, shift, device, num_timesteps=1000):
+def stochastic_sampling_timesteps(num_inference_steps, shift, device, num_timesteps=1000, timestep_list=None):
    """Official FlashHead timestep schedule with shift transform.

    Hardcoded schedules match the CF reference for distilled few-step inference:
      - 4-step: [1000, 750, 500, 250] — Stage 2 CCD (`causal_cd_framewise.yaml`)
      - 2-step: [1000, 500]           — Stage 3 DMD (`causal_forcing_dmd_framewise_2step.yaml`)
-    Anything else falls back to linspace.
+    Anything else falls back to linspace. An explicit `timestep_list` overrides
+    all of the above (per-block schedules such as the official One-Forcing
+    first-frame list).
    """
-    if num_inference_steps == 4:
+    if timestep_list is not None:
+        timesteps = list(timestep_list)
+    elif num_inference_steps == 4:
        timesteps = [1000, 750, 500, 250]
    elif num_inference_steps == 2:
        timesteps = [1000, 500]
@@ -589,6 +593,7 @@ class WanSelfForcingPipeline(DiffusionPipeline):
        independent_first_frame: bool = True,
        context_noise: int = 0,
        stochastic_sampling: bool = True,
+        first_frame_denoising_step_list: Optional[List[int]] = None,
        streaming: bool = False,
        decode_callback: Optional[Callable[[torch.Tensor, int], None]] = None,
        forcing_kv_enable: Optional[bool] = None,
@@ -611,6 +616,12 @@ class WanSelfForcingPipeline(DiffusionPipeline):
            num_frame_per_block: Number of frames to generate per block.
            independent_first_frame: Whether to generate the first frame independently (T2V mode).
            context_noise: Context noise level for KV cache update (matches training config).
+            first_frame_denoising_step_list: Optional explicit raw timestep list
+                (e.g. [1000, 750, 500, 250]) used only for the block that starts at
+                frame 0, mirroring the official One-Forcing
+                `first_frame_denoising_step_list`: that block is the only one with
+                an empty KV context and benefits from a multi-step schedule while
+                all other blocks keep the global one. None disables the override.
            streaming: If True, decode each causal block into pixels right after it is
                generated (instead of accumulating all latents and decoding once at the end).
                The VAE's causal temporal cache is threaded across blocks so the result is
@@ -696,6 +707,7 @@ class WanSelfForcingPipeline(DiffusionPipeline):
            in_prompt_embeds = prompt_embeds

        # 4. Prepare timesteps
+        custom_timesteps = timesteps
        if stochastic_sampling:
            timesteps = stochastic_sampling_timesteps(num_inference_steps, shift, device)
        elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
@@ -868,16 +880,40 @@ class WanSelfForcingPipeline(DiffusionPipeline):
            end_idx = start_idx + current_num_frames
            noisy_input = noise[:, :, start_idx:end_idx]

-            if not stochastic_sampling and isinstance(self.scheduler, FlowUniPCMultistepScheduler):
-                self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
-                timesteps = self.scheduler.timesteps
-            elif not stochastic_sampling and isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
-                self.scheduler.set_timesteps(device=device, sigmas=sampling_sigmas)
-                timesteps = self.scheduler.timesteps
-            else:
+            # Official One-Forcing gives the block starting at frame 0 (the only
+            # block with an empty KV context) its own multi-step schedule; all
+            # other blocks keep the global one.
+            if current_start_frame == 0 and first_frame_denoising_step_list is not None:
+                if stochastic_sampling:
+                    block_timesteps = stochastic_sampling_timesteps(
+                        len(first_frame_denoising_step_list), shift, device,
+                        timestep_list=first_frame_denoising_step_list)
+                elif isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+                    self.scheduler.set_timesteps(
+                        timesteps=list(first_frame_denoising_step_list), device=device)
+                    block_timesteps = self.scheduler.timesteps
+                else:
+                    raise ValueError(
+                        "first_frame_denoising_step_list requires stochastic sampling or "
+                        f"FlowMatchEulerDiscreteScheduler, got {type(self.scheduler)}")
                self.scheduler._step_index = None
-            
-            denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
+            else:
+                if not stochastic_sampling and isinstance(self.scheduler, FlowUniPCMultistepScheduler):
+                    self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
+                    timesteps = self.scheduler.timesteps
+                elif not stochastic_sampling and isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
+                    self.scheduler.set_timesteps(device=device, sigmas=sampling_sigmas)
+                    timesteps = self.scheduler.timesteps
+                elif (not stochastic_sampling and first_frame_denoising_step_list is not None
+                        and isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler)):
+                    # Restore the global schedule overridden by the first-frame block
+                    timesteps, num_inference_steps = retrieve_timesteps(
+                        self.scheduler, num_inference_steps, device, custom_timesteps)
+                else:
+                    self.scheduler._step_index = None
+                block_timesteps = timesteps
+
+            denoise_timesteps = block_timesteps[:-1] if stochastic_sampling else block_timesteps
            with self.progress_bar(total=len(denoise_timesteps)) as progress_bar:
                for step_idx, t in enumerate(denoise_timesteps):
                    # Per-frame timesteps for causal generation
@@ -947,8 +983,8 @@ class WanSelfForcingPipeline(DiffusionPipeline):
                    
                    # compute the previous noisy sample x_t -> x_t-1
                    if stochastic_sampling:
-                        t_i = (timesteps[step_idx] / 1000).to(weight_dtype)
-                        t_i_1 = (timesteps[step_idx + 1] / 1000).to(weight_dtype)
+                        t_i = (block_timesteps[step_idx] / 1000).to(weight_dtype)
+                        t_i_1 = (block_timesteps[step_idx + 1] / 1000).to(weight_dtype)
                        denoised_pred = noisy_input - flow_pred * t_i
                        noisy_input = (1 - t_i_1) * denoised_pred + t_i_1 * torch.randn(
                            denoised_pred.shape, dtype=denoised_pred.dtype, device=device, generator=generator
@@ -91,7 +91,13 @@ def checkpoint_fsdp_forward(model, accumulated, *, context_fn=None, **kwargs):
                    else:
                        parameter.grad.add_(previous)
                if parameter.grad is not None:
-                    accumulated[parameter] = parameter.grad.detach().clone()
+                    if previous is None:
+                        accumulated[parameter] = parameter.grad.detach().clone()
+                    else:
+                        # Keep one independent accumulator per shard across blocks.
+                        # Replacing it with a clone retains both full copies until
+                        # previous_grads leaves this backward call.
+                        previous.copy_(parameter.grad)
            result = [None] * len(keys)
            for index, tensor in zip(ctx.tensor_indices, tensors):
                result[index] = tensor.grad