From f00a7d480e5df0ffa8acd2f05e5d80d175c6104d Mon Sep 17 00:00:00 2001 From: bubbliiiing <3323290568@qq.com> Date: Fri, 30 Oct 2026 16:33:15 +0800 Subject: [PATCH] Update one-forcing and causal forcing params --- examples/wan2.1_one_forcing/predict_t2v.py | 303 ++ .../README_TRAIN_CAUSAL_DMD_DISTILL.md | 19 +- .../README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md | 19 +- .../train_ar_diffusion.sh | 1 + .../train_causal_consistency_distill.sh | 1 + .../wan2.1_causal_forcing/train_causal_dmd.sh | 9 +- .../wan2.1_one_forcing/README_ONE_FORCING.md | 356 ++ .../README_ONE_FORCING_zh-CN.md | 355 ++ .../wan2.1_one_forcing/train_one_forcing.py | 3490 +++++++++++++++++ .../wan2.1_one_forcing/train_one_forcing.sh | 85 + .../README_TRAIN_DISTILL.md | 86 +- .../README_TRAIN_DISTILL_zh-CN.md | 86 +- scripts/wan2.1_self_forcing/train_distill.sh | 15 +- videox_fun/models/attention_utils.py | 1 + videox_fun/models/wan_tae.py | 54 + .../pipeline/pipeline_wan_self_forcing.py | 64 +- videox_fun/utils/fsdp_utils.py | 8 +- 17 files changed, 4866 insertions(+), 86 deletions(-) create mode 100644 examples/wan2.1_one_forcing/predict_t2v.py create mode 100644 scripts/wan2.1_one_forcing/README_ONE_FORCING.md create mode 100644 scripts/wan2.1_one_forcing/README_ONE_FORCING_zh-CN.md create mode 100644 scripts/wan2.1_one_forcing/train_one_forcing.py create mode 100644 scripts/wan2.1_one_forcing/train_one_forcing.sh diff --git a/examples/wan2.1_one_forcing/predict_t2v.py b/examples/wan2.1_one_forcing/predict_t2v.py new file mode 100644 index 0000000..c66d466 --- /dev/null +++ b/examples/wan2.1_one_forcing/predict_t2v.py @@ -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_" 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() diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md index 3f935c8..e9c55ce 100644 --- a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL.md @@ -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))`. diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md index 3e399d1..64f3307 100644 --- a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md @@ -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))`。 diff --git a/scripts/wan2.1_causal_forcing/train_ar_diffusion.sh b/scripts/wan2.1_causal_forcing/train_ar_diffusion.sh index 5914fbe..c887910 100644 --- a/scripts/wan2.1_causal_forcing/train_ar_diffusion.sh +++ b/scripts/wan2.1_causal_forcing/train_ar_diffusion.sh @@ -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 \ diff --git a/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh b/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh index db6d0e8..d70c913 100644 --- a/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh +++ b/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh @@ -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 \ diff --git a/scripts/wan2.1_causal_forcing/train_causal_dmd.sh b/scripts/wan2.1_causal_forcing/train_causal_dmd.sh index 6181bc1..8873d4a 100755 --- a/scripts/wan2.1_causal_forcing/train_causal_dmd.sh +++ b/scripts/wan2.1_causal_forcing/train_causal_dmd.sh @@ -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 \ diff --git a/scripts/wan2.1_one_forcing/README_ONE_FORCING.md b/scripts/wan2.1_one_forcing/README_ONE_FORCING.md new file mode 100644 index 0000000..d919f85 --- /dev/null +++ b/scripts/wan2.1_one_forcing/README_ONE_FORCING.md @@ -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 diff --git a/scripts/wan2.1_one_forcing/README_ONE_FORCING_zh-CN.md b/scripts/wan2.1_one_forcing/README_ONE_FORCING_zh-CN.md new file mode 100644 index 0000000..3688357 --- /dev/null +++ b/scripts/wan2.1_one_forcing/README_ONE_FORCING_zh-CN.md @@ -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 diff --git a/scripts/wan2.1_one_forcing/train_one_forcing.py b/scripts/wan2.1_one_forcing/train_one_forcing.py new file mode 100644 index 0000000..e8ffe11 --- /dev/null +++ b/scripts/wan2.1_one_forcing/train_one_forcing.py @@ -0,0 +1,3490 @@ +"""Modified from https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py +""" +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and + +import argparse +import contextlib +import dataclasses +import gc +import json +import logging +import math +import os +import pickle +import shutil +import sys + +import accelerate +import diffusers +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +import torch.utils.checkpoint +import torchvision.transforms.functional as TF +import transformers +from accelerate import Accelerator, FullyShardedDataParallelPlugin +from accelerate.logging import get_logger +from accelerate.state import AcceleratorState +from accelerate.utils import ProjectConfiguration, set_seed +from diffusers import DDIMScheduler, FlowMatchEulerDiscreteScheduler +from diffusers.optimization import get_scheduler +from diffusers.training_utils import (EMAModel, + compute_density_for_timestep_sampling, + compute_loss_weighting_for_sd3) +from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils.torch_utils import is_compiled_module +from einops import rearrange +from omegaconf import OmegaConf +from packaging import version +from PIL import Image +from torch.distributed.fsdp.fully_sharded_data_parallel import ( + FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig, + ShardedStateDictConfig) +from torch.utils.data import BatchSampler, Dataset, RandomSampler +from torch.utils.tensorboard import SummaryWriter +from torchvision import transforms +from transformers import AutoTokenizer +from transformers.utils import ContextManagers + +import datasets + +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.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, + ASPECT_RATIO_RANDOM_CROP_PROB, + AspectRatioBatchImageVideoSampler, + ImageVideoDataset, ImageVideoSampler, + RandomSampler, TextDataset, get_closest_ratio, + get_random_mask) +from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel, + WanTransformer3DModel, + WanTransformer3DModel_SelfForcing, + add_wan_gan_discriminator) +from videox_fun.pipeline import (WanI2VPipeline, WanPipeline, + WanSelfForcingPipeline) +from videox_fun.utils.discrete_sampler import DiscreteSampling +from videox_fun.utils.fsdp_ema import FSDPEMA +from videox_fun.utils.tqdm_bar import PauseAwareTqdm +from videox_fun.utils.utils import (calculate_dimensions, get_image_latent, + get_image_to_video_latent, + save_videos_grid) + +if is_wandb_available(): + import wandb + + +def initialize_kv_cache_for_training(batch_size, num_frames, frame_seq_length, num_layers, num_heads, head_dim, dtype, device): + """Initialize KV cache for block-by-block training""" + kv_cache_size = num_frames * frame_seq_length + kv_cache = [] + + for _ in range(num_layers): + kv_cache.append({ + "k": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device) + }) + + return kv_cache + + +def initialize_crossattn_cache_for_training(batch_size, text_len, num_layers, num_heads, head_dim, dtype, device): + """Initialize cross-attention cache for block-by-block training""" + crossattn_cache = [] + + for _ in range(num_layers): + crossattn_cache.append({ + "k": torch.zeros([batch_size, text_len, num_heads, head_dim], dtype=dtype, device=device), + "v": torch.zeros([batch_size, text_len, num_heads, head_dim], dtype=dtype, device=device), + "is_init": False + }) + + return crossattn_cache + + +def reencode_boundary_latent(vae, pred_latents, weight_dtype, score_num_frames=21): + """ + Re-encode the boundary frame to get a clean latent for the score window. + Follows Self-Forcing reference: decode all frames before the score window, take last pixel frame, re-encode. + Input: pred_latents [B, C, F, H, W] (all generated latent frames) + Output: boundary_latent [B, C, 1, H, W] + """ + with torch.no_grad(): + # Decode all frames except the last (score_num_frames - 1) to pixels + tail_len = score_num_frames - 1 + latent_to_decode = pred_latents[:, :, :-tail_len] + # VAE expects [B, C, F, H, W], decode returns [B, C, F, H, W] pixels + pixels = vae.decode(latent_to_decode.to(vae.dtype)).sample # [B, C, F, H, W] + # Take the last frame + frame = pixels[:, :, -1:, :, :] # [B, C, 1, H, W] + # Re-encode the last frame to get clean boundary latent + boundary_latent = vae.encode(frame)[0].sample().to(weight_dtype) # [B, C, 1, H, W] + return boundary_latent + + +def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_first_frame=False): + """ + Slice the last `score_num_frames` latent frames for score computation. + If pred has more than score_num_frames, re-encode boundary frame for clean context. + Returns: (pred_for_score, score_num_frames, need_gradient_mask) + """ + num_frames = pred.shape[2] + if num_frames <= score_num_frames: + return pred, num_frames, False + + # Re-encode boundary for cleaner score input + try: + boundary_latent = reencode_boundary_latent(vae, pred, weight_dtype, score_num_frames=score_num_frames) + pred_for_score = torch.cat([boundary_latent, pred[:, :, -(score_num_frames - 1):]], dim=2) + except Exception: + # Fallback: simple slice without boundary re-encoding + pred_for_score = pred[:, :, -score_num_frames:] + + return pred_for_score, score_num_frames, True + + +def repeat_batch(value, repeats): + """Repeat a batched tensor or per-sample list along the batch dimension.""" + if value is None or repeats == 1: + return value + if torch.is_tensor(value): + return torch.cat([value] * repeats) + return list(value) * repeats + + +def filter_kwargs(cls, kwargs): + import inspect + sig = inspect.signature(cls.__init__) + valid_params = set(sig.parameters.keys()) - {'self', 'cls'} + filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params} + return filtered_kwargs + +def get_random_downsample_ratio(sample_size, image_ratio=[], + all_choices=False, rng=None): + def _create_special_list(length): + if length == 1: + return [1.0] + if length >= 2: + first_element = 0.75 + remaining_sum = 1.0 - first_element + other_elements_value = remaining_sum / (length - 1) + special_list = [first_element] + [other_elements_value] * (length - 1) + return special_list + + if sample_size >= 1536: + number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio + elif sample_size >= 1024: + number_list = [1, 1.25, 1.5, 2] + image_ratio + elif sample_size >= 768: + number_list = [1, 1.25, 1.5] + image_ratio + elif sample_size >= 512: + number_list = [1] + image_ratio + else: + number_list = [1] + + if all_choices: + return number_list + + number_list_prob = np.array(_create_special_list(len(number_list))) + if rng is None: + return np.random.choice(number_list, p = number_list_prob) + else: + return rng.choice(number_list, p = number_list_prob) + +def resize_mask(mask, latent, process_first_frame_only=True): + latent_size = latent.size() + batch_size, channels, num_frames, height, width = mask.shape + + if process_first_frame_only: + target_size = list(latent_size[2:]) + target_size[0] = 1 + first_frame_resized = F.interpolate( + mask[:, :, 0:1, :, :], + size=target_size, + mode='trilinear', + align_corners=False + ) + + target_size = list(latent_size[2:]) + target_size[0] = target_size[0] - 1 + if target_size[0] != 0: + remaining_frames_resized = F.interpolate( + mask[:, :, 1:, :, :], + size=target_size, + mode='trilinear', + align_corners=False + ) + resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2) + else: + resized_mask = first_frame_resized + else: + target_size = list(latent_size[2:]) + resized_mask = F.interpolate( + mask, + size=target_size, + mode='trilinear', + align_corners=False + ) + return resized_mask + +# Will error if the minimal version of diffusers is not installed. Remove at your own risks. +check_min_version("0.18.0.dev0") + +logger = get_logger(__name__, log_level="INFO") + +def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step): + try: + is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine' + if is_deepspeed: + origin_config = transformer3d.config + transformer3d.config = accelerator.unwrap_model(transformer3d).config + with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + logger.info("Running validation... ") + scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) + ) + + if args.train_mode != "normal": + raise NotImplementedError(f"Validation for train_mode '{args.train_mode}' is not yet supported with WanSelfForcingPipeline. Only T2V (train_mode='normal') is currently supported.") + else: + pipeline = WanSelfForcingPipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d, + scheduler=scheduler, + ) + pipeline = pipeline.to(accelerator.device) + + if args.seed is None: + generator = None + else: + rank_seed = args.seed + accelerator.process_index + generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed) + logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}") + + for i in range(len(args.validation_prompts)): + if args.train_mode != "normal": + raise NotImplementedError(f"Validation for train_mode '{args.train_mode}' is not yet supported with WanSelfForcingPipeline. Only T2V (train_mode='normal') is currently supported.") + else: + if args.fix_sample_size is not None: + height, width = args.fix_sample_size + else: + height, width = args.video_sample_size, args.video_sample_size + sample = pipeline( + args.validation_prompts[i], + num_frames = args.video_sample_n_frames, + negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + height = height, + width = width, + generator = generator, + guidance_scale = 1.0, + num_inference_steps = len(args.denoising_step_indices_list), + num_frame_per_block = args.num_frame_per_block, + independent_first_frame = args.independent_first_frame, + context_noise = args.context_noise, + # Stream-decode each causal block instead of one full-segment VAE decode: + # numerically identical output, but peak decode memory drops ~20GB -> ~1-2GB, + # which keeps in-loop validation from OOMing under One-Forcing's high resident + # footprint (generator + fake_score critic + 14B teacher all on-device). + streaming = True, + ).videos + os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True) + save_videos_grid( + sample, + os.path.join( + args.output_dir, + f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.mp4" + ) + ) + + del pipeline + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if is_deepspeed: + transformer3d.config = origin_config + except Exception as e: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + print(f"Eval error on rank {accelerator.process_index} with info {e}") + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + +def linear_decay(initial_value, final_value, total_steps, current_step): + if current_step >= total_steps: + return final_value + current_step = max(0, current_step) + step_size = (final_value - initial_value) / total_steps + current_value = initial_value + step_size * current_step + return current_value + +def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=None): + u = torch.normal(mean=0.0, std=1.0, size=shape, device=device, generator=generator) + t = 1 / (1 + torch.exp(-u)) * (high - low) + low + return torch.clip(t.to(torch.int32), low, high - 1) + +def parse_args(): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--input_perturbation", type=float, default=0, help="The scale of input perturbation. Recommended 0.1." + ) + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + required=False, + help="Revision of pretrained model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--variant", + type=str, + default=None, + help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16", + ) + parser.add_argument( + "--train_data_dir", + type=str, + default=None, + help=( + "A folder containing the training data. " + ), + ) + parser.add_argument( + "--train_data_meta", + type=str, + default=None, + help=( + "A csv containing the training data. " + ), + ) + parser.add_argument( + "--max_train_samples", + type=int, + default=None, + help=( + "For debugging purposes or quicker training, truncate the number of training examples to this " + "value if set." + ), + ) + parser.add_argument( + "--validation_prompts", + type=str, + default=None, + nargs="+", + help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."), + ) + parser.add_argument( + "--validation_paths", + type=str, + default=None, + nargs="+", + help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."), + ) + parser.add_argument( + "--negative_prompt", + type=str, + default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + help=("The negative prompt of cfg distill"), + ) + parser.add_argument( + "--output_dir", + type=str, + default="sd-model-finetuned", + help="The output directory where the model predictions and checkpoints will be written.", + ) + parser.add_argument( + "--cache_dir", + type=str, + default=None, + help="The directory where the downloaded models and datasets will be stored.", + ) + parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") + parser.add_argument( + "--random_flip", + action="store_true", + help="whether to randomly flip images horizontally", + ) + parser.add_argument( + "--use_came", + action="store_true", + help="whether to use came", + ) + parser.add_argument( + "--multi_stream", + action="store_true", + help="whether to use cuda multi-stream", + ) + parser.add_argument( + "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader." + ) + parser.add_argument( + "--vae_mini_batch", type=int, default=32, help="mini batch size for vae." + ) + parser.add_argument("--num_train_epochs", type=int, default=100) + parser.add_argument( + "--max_train_steps", + type=int, + default=None, + help="Total number of training steps to perform. If provided, overrides num_train_epochs.", + ) + parser.add_argument( + "--gradient_accumulation_steps", + type=int, + default=1, + help="Number of updates steps to accumulate before performing a backward/update pass.", + ) + parser.add_argument( + "--gradient_checkpointing", + action="store_true", + help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--learning_rate_critic", + type=float, + default=1e-4, + help="Initial learning rate (after the potential warmup period) to use.", + ) + parser.add_argument( + "--scale_lr", + action="store_true", + default=False, + help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", + ) + parser.add_argument( + "--lr_scheduler", + type=str, + default="constant", + help=( + 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' + ' "constant", "constant_with_warmup"]' + ), + ) + parser.add_argument( + "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler." + ) + parser.add_argument( + "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." + ) + parser.add_argument( + "--allow_tf32", + action="store_true", + help=( + "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" + " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + ), + ) + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.") + parser.add_argument("--ema_start_step", type=int, default=50, help="Step at which the EMA shadow is re-seeded from the live generator, mirroring official One-Forcing's lazy EMA creation.") + parser.add_argument( + "--ema_decay", type=float, default=0.9999, + help="FSDPEMA shadow decay per step (FSDP path); lower values track the live generator faster over a short run.", + ) + parser.add_argument( + "--fsdp_reduce_fp32", action="store_true", + help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.", + ) + parser.add_argument( + "--score_timestep_clamp", action="store_true", + help="Clamp DMD, GAN and critic timesteps to [0.02, 0.98] * num_train_timesteps.", + ) + parser.add_argument( + "--gradient_num_blocks", + type=int, + default=0, + help="Restrict the DMD gradient to a window of this many rollout blocks inside the score " + "window (official One-Forcing uses 17 of 21); 0 backpropagates every scored frame.", + ) + parser.add_argument( + "--gradient_window_position", + type=str, + default="tail", + choices=["first", "tail"], + help="Position of the gradient window selected by --gradient_num_blocks.", + ) + parser.add_argument( + "--dataloader_num_workers", + type=int, + default=0, + help=( + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ), + ) + parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") + parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") + parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") + parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") + parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") + parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") + parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument( + "--prediction_type", + type=str, + default=None, + help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.", + ) + parser.add_argument( + "--hub_model_id", + type=str, + default=None, + help="The name of the repository to keep in sync with the local `output_dir`.", + ) + parser.add_argument( + "--logging_dir", + type=str, + default="logs", + help=( + "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" + " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." + ), + ) + parser.add_argument( + "--report_model_info", action="store_true", help="Whether or not to report more info about model (such as norm, grad)." + ) + parser.add_argument( + "--mixed_precision", + type=str, + default=None, + choices=["no", "fp16", "bf16"], + help=( + "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" + " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" + " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." + ), + ) + parser.add_argument( + "--report_to", + type=str, + default="tensorboard", + help=( + 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' + ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' + ), + ) + parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") + parser.add_argument( + "--checkpointing_steps", + type=int, + default=500, + help=( + "Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming" + " training using `--resume_from_checkpoint`." + ), + ) + parser.add_argument( + "--checkpoints_total_limit", + type=int, + default=None, + help=("Max number of checkpoints to store."), + ) + parser.add_argument( + "--resume_from_checkpoint", + type=str, + default=None, + help=( + "Whether training should be resumed from a previous checkpoint. Use a path saved by" + ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' + ), + ) + parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.") + parser.add_argument( + "--validation_epochs", + type=int, + default=5, + help="Run validation every X epochs.", + ) + parser.add_argument( + "--validation_steps", + type=int, + default=2000, + help="Run validation every X steps.", + ) + parser.add_argument( + "--tracker_project_name", + type=str, + default="text2image-fine-tune", + help=( + "The `project_name` argument passed to Accelerator.init_trackers for" + " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator" + ), + ) + + parser.add_argument( + "--snr_loss", action="store_true", help="Whether or not to use snr_loss." + ) + parser.add_argument( + "--uniform_sampling", action="store_true", help="Whether or not to use uniform_sampling." + ) + parser.add_argument( + "--enable_text_encoder_in_dataloader", action="store_true", help="Whether or not to use text encoder in dataloader." + ) + parser.add_argument( + "--enable_bucket", action="store_true", help="Whether enable bucket sample in datasets." + ) + parser.add_argument( + "--random_ratio_crop", action="store_true", help="Whether enable random ratio crop sample in datasets." + ) + parser.add_argument( + "--random_frame_crop", action="store_true", help="Whether enable random frame crop sample in datasets." + ) + parser.add_argument( + "--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets." + ) + parser.add_argument( + "--training_with_video_token_length", action="store_true", help="The training stage of the model in training.", + ) + parser.add_argument( + "--auto_tile_batch_size", action="store_true", help="Whether to auto tile batch size.", + ) + parser.add_argument( + "--motion_sub_loss", action="store_true", help="Whether enable motion sub loss." + ) + parser.add_argument( + "--motion_sub_loss_ratio", type=float, default=0.25, help="The ratio of motion sub loss." + ) + parser.add_argument( + "--train_sampling_steps", + type=int, + default=1000, + help="Run train_sampling_steps.", + ) + parser.add_argument( + "--keep_all_node_same_token_length", + action="store_true", + help="Reference of the length token.", + ) + parser.add_argument( + "--token_sample_size", + type=int, + default=512, + help="Sample size of the token.", + ) + parser.add_argument( + "--video_sample_size", + type=int, + default=512, + help="Sample size of the video.", + ) + parser.add_argument( + "--image_sample_size", + type=int, + default=512, + help="Sample size of the image.", + ) + parser.add_argument( + "--fix_sample_size", + nargs=2, type=int, default=None, + help="Fix Sample size [height, width] when using bucket and collate_fn." + ) + parser.add_argument( + "--video_sample_stride", + type=int, + default=4, + help="Sample stride of the video.", + ) + parser.add_argument( + "--video_sample_n_frames", + type=int, + default=17, + help="Num frame of video.", + ) + parser.add_argument( + "--video_repeat", + type=int, + default=0, + help="Num of repeat video.", + ) + parser.add_argument( + "--config_path", + type=str, + default=None, + help=( + "The config of the model in training." + ), + ) + parser.add_argument( + "--transformer_path", + type=str, + default=None, + help=("If you want to load the weight from other transformers, input its path."), + ) + parser.add_argument( + "--ode_transformer_path", + type=str, + default=None, + help=("If you want to load the ode-trained weight into generator transformer3d, input its path."), + ) + parser.add_argument( + "--real_score_pretrained_model_name_or_path", + type=str, + default=None, + help=( + "Path to a non-causal pretrained model used as the DMD real_score teacher. " + "For CF Stage 3 this should point at Wan2.1-T2V-14B. If unset, falls back to " + "`--pretrained_model_name_or_path` (1.3B) — useful only for sanity testing." + ), + ) + parser.add_argument( + "--vae_path", + type=str, + default=None, + help=("If you want to load the weight from other vaes, input its path."), + ) + + parser.add_argument( + '--trainable_modules', + nargs='+', + help='Enter a list of trainable modules' + ) + parser.add_argument( + '--trainable_modules_low_learning_rate', + nargs='+', + default=[], + help='Enter a list of trainable modules with lower learning rate' + ) + parser.add_argument( + '--tokenizer_max_length', + type=int, + default=512, + help='Max length of tokenizer' + ) + parser.add_argument( + "--use_deepspeed", action="store_true", help="Whether or not to use deepspeed." + ) + parser.add_argument( + "--use_fsdp", action="store_true", help="Whether or not to use fsdp." + ) + parser.add_argument( + "--low_vram", action="store_true", help="Whether enable low_vram mode." + ) + parser.add_argument( + "--train_mode", + type=str, + default="normal", + help=( + 'The format of training data. Support `"normal"`' + ' (default), `"i2v"`.' + ), + ) + parser.add_argument( + "--gen_update_interval", + type=int, + default=5, + help="The ratio to update transformer3d.", + ) + parser.add_argument( + "--fake_guidance_scale", + type=float, + default=0.0, + help="The cfg scale for fake iscore.", + ) + parser.add_argument( + "--real_guidance_scale", + type=float, + default=4.5, + help="The cfg scale for real score.", + ) + parser.add_argument( + '--denoising_step_indices_list', + nargs='+', + type=int, + default=[1000, 500], + help="Denoising step indices (in train_sampling_steps space). CF Stage 3 DMD default = 2-step [1000, 500].", + ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) + parser.add_argument( + "--num_frame_per_block", + type=int, + default=3, + help="Number of frames per block for Self-Forcing causal training" + ) + parser.add_argument( + "--independent_first_frame", + action="store_true", + help="Whether first frame is independent ([1, N, N, ...] pattern)" + ) + parser.add_argument( + "--use_kv_cache_training", + action="store_true", + help="Use KV cache block-by-block training (matches original Self-Forcing)" + ) + parser.add_argument( + "--score_num_frames", + type=int, + default=21, + help="Number of latent frames for score computation window (default: 21, matching base model). " + "fake_score/real_score always receive this many frames." + ) + parser.add_argument( + "--min_length_prob_bias", + type=float, + default=0.0, + help="Probability bias for sampling the minimum length (score_num_frames). " + "0.0 = uniform sampling (default), 0.5 = 50%% prob for min length, " + "remaining prob distributed equally among longer lengths. " + "Use this to increase 21-frame training ratio." + ) + parser.add_argument( + "--context_noise", + type=int, + default=0, + help="Context noise level for KV cache update (matches training config)" + ) + parser.add_argument( + "--use_teacher_forcing", + action="store_true", + help="Enable teacher forcing training (pass clean_x to transformer)" + ) + parser.add_argument( + "--teacher_forcing_prob", + type=float, + default=1.0, + help="Probability of applying teacher forcing per step (1.0 = always)" + ) + + parser.add_argument( + "--gan_g_weight", type=float, default=0.0, + help="One-Forcing: non-saturating generator GAN loss weight (0 disables). Positive GAN/R1/R2 weights require real videos.", + ) + parser.add_argument( + "--gan_d_weight", type=float, default=0.0, + help="One-Forcing: discriminator GAN loss weight added to critic denoising loss (0 disables).", + ) + parser.add_argument( + "--gan_feature_layers", nargs="+", type=int, default=[21, 29], + help="fake_score block indices feeding the discriminator head, one register token per layer.", + ) + parser.add_argument( + "--gan_block_ffn_dim", type=int, default=2048, + help="FFN dimension of the discriminator cross-attention blocks.", + ) + parser.add_argument( + "--gan_block_num_heads", type=int, default=12, + help="Number of heads of the discriminator cross-attention blocks.", + ) + parser.add_argument( + "--relativistic_discriminator", action="store_true", + help="Use relativistic GAN losses based on real/fake logit differences.", + ) + parser.add_argument( + "--r1_weight", type=float, default=0.0, + help="Finite-difference R1 penalty weight on noisy real latents (APT form, not divided by sigma).", + ) + parser.add_argument( + "--r1_sigma", type=float, default=0.01, + help="Perturbation standard deviation for the R1 penalty.", + ) + parser.add_argument( + "--r2_weight", type=float, default=0.0, + help="Finite-difference R2 penalty weight on noisy generated latents (0.5 * squared logit difference / sigma squared)." + ) + parser.add_argument( + "--r2_sigma", type=float, default=0.01, + help="Perturbation standard deviation for the R2 penalty.", + ) + + args = parser.parse_args() + env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) + if env_local_rank != -1 and env_local_rank != args.local_rank: + args.local_rank = env_local_rank + + return args + + +def main(): + args = parse_args() + use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)) + if any(not math.isfinite(w) or w < 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)): + raise ValueError("GAN and regularization weights must be finite and non-negative.") + if use_gan and (args.use_deepspeed or args.gradient_accumulation_steps != 1): + raise ValueError("The GAN branch does not support DeepSpeed and requires --gradient_accumulation_steps=1.") + if use_gan and (args.r1_weight > 0 and (not math.isfinite(args.r1_sigma) or args.r1_sigma <= 0) + or args.r2_weight > 0 and (not math.isfinite(args.r2_sigma) or args.r2_sigma <= 0)): + raise ValueError("Enabled R1/R2 penalties require finite, positive perturbation standard deviations.") + + if args.report_to == "wandb" and args.hub_token is not None: + raise ValueError( + "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." + " Please use `huggingface-cli login` to authenticate with the Hub." + ) + + logging_dir = os.path.join(args.output_dir, args.logging_dir) + + config = OmegaConf.load(args.config_path) + accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) + + accelerator = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + accelerator_fake_score_transformer3d = Accelerator( + gradient_accumulation_steps=args.gradient_accumulation_steps, + mixed_precision=args.mixed_precision, + log_with=args.report_to, + project_config=accelerator_project_config, + ) + + deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None + fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None + if use_gan and deepspeed_plugin is not None: + raise ValueError("The GAN branch does not support DeepSpeed.") + if deepspeed_plugin is not None: + zero_stage = int(deepspeed_plugin.zero_stage) + fsdp_stage = 0 + print(f"Using DeepSpeed Zero stage: {zero_stage}") + + args.use_deepspeed = True + if zero_stage == 3: + print(f"Auto set save_state to True because zero_stage == 3") + args.save_state = True + elif fsdp_plugin is not None: + from torch.distributed.fsdp import ShardingStrategy + zero_stage = 0 + if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD: + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2. + fsdp_stage = 3 + elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP: + fsdp_stage = 2 + else: + fsdp_stage = 0 + print(f"Using FSDP stage: {fsdp_stage}") + if args.fsdp_reduce_fp32 and fsdp_plugin.mixed_precision_policy is not None: + fsdp_plugin.mixed_precision_policy = dataclasses.replace(fsdp_plugin.mixed_precision_policy, reduce_dtype=torch.float32) + print(f"FSDP mixed precision policy: {fsdp_plugin.mixed_precision_policy}") + + args.use_fsdp = True + if fsdp_stage == 3: + print(f"Auto set save_state to True because fsdp_stage == 3") + args.save_state = True + else: + zero_stage = 0 + fsdp_stage = 0 + print("DeepSpeed is not enabled.") + + if accelerator.is_main_process: + writer = SummaryWriter(log_dir=logging_dir) + + # Make one log on every process with the configuration for debugging. + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + logger.info(accelerator.state, main_process_only=False) + if accelerator.is_local_main_process: + datasets.utils.logging.set_verbosity_warning() + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + datasets.utils.logging.set_verbosity_error() + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + # If passed along, set the training seed now. + if args.seed is not None: + set_seed(args.seed) + rng = np.random.default_rng(np.random.PCG64(args.seed + accelerator.process_index)) + torch_rng = torch.Generator(accelerator.device).manual_seed(args.seed + accelerator.process_index) + else: + rng = None + torch_rng = None + index_rng = np.random.default_rng(np.random.PCG64(43)) + print(f"Init rng with seed {args.seed + accelerator.process_index}. Process_index is {accelerator.process_index}") + + # Handle the repository creation + if accelerator.is_main_process: + if args.output_dir is not None: + os.makedirs(args.output_dir, exist_ok=True) + + # For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora transformer3d) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = torch.float32 + if accelerator.mixed_precision == "fp16": + weight_dtype = torch.float16 + args.mixed_precision = accelerator.mixed_precision + elif accelerator.mixed_precision == "bf16": + weight_dtype = torch.bfloat16 + args.mixed_precision = accelerator.mixed_precision + + # Load scheduler, tokenizer and models. + noise_scheduler = FlowMatchEulerDiscreteScheduler( + **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) + ) + + # Get Tokenizer + tokenizer = AutoTokenizer.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), + ) + + def deepspeed_zero_init_disabled_context_manager(): + """ + returns either a context list that includes one that will disable zero.Init or an empty context list + """ + deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None + if deepspeed_plugin is None: + return [] + + return [deepspeed_plugin.zero3_init_context_manager(enable=False)] + + # Currently Accelerate doesn't know how to handle multiple models under Deepspeed ZeRO stage 3. + # For this to work properly all models must be run through `accelerate.prepare`. But accelerate + # will try to assign the same optimizer with the same weights to all models during + # `deepspeed.initialize`, which of course doesn't work. + # + # For now the following workaround will partially support Deepspeed ZeRO-3, by excluding the 2 + # frozen models from being partitioned during `zero.Init` which gets called during + # `from_pretrained` So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding + # across multiple gpus and only UNet2DConditionModel will get ZeRO sharded. + with ContextManagers(deepspeed_zero_init_disabled_context_manager()): + # Get Text encoder + text_encoder = WanT5EncoderModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, 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, + ) + text_encoder = text_encoder.eval() + # Get Vae + vae = AutoencoderKLWan.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')), + additional_kwargs=OmegaConf.to_container(config['vae_kwargs']), + ) + vae.eval() + # Get Clip Image Encoder + if args.train_mode != "normal": + clip_image_encoder = CLIPModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')), + ) + clip_image_encoder = clip_image_encoder.eval() + else: + clip_image_encoder = None + + # Get Transformer. + # Keep all three transformers in fp32 — the same lesson learned in Stage 1/2: + # accelerate's mixed_precision="bf16" autocasts forward to bf16 while keeping + # master weights & Adam moments in fp32. If params live in bf16, every update + # (LR*grad ~ 1e-5 for LR=2e-6) falls below bf16 mantissa precision and is + # rounded to zero. CF Stage 3 FSDP also keeps fp32 master weights (no + # compute_dtype set on the MixedPrecision policy). + generator_transformer3d = WanTransformer3DModel_SelfForcing.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + # CF Stage 3 uses a 14B non-causal teacher as real_score; pass an explicit + # `--real_score_pretrained_model_name_or_path` to point at Wan2.1-T2V-14B. + real_score_base = args.real_score_pretrained_model_name_or_path or args.pretrained_model_name_or_path + real_score_transformer3d = WanTransformer3DModel.from_pretrained( + os.path.join(real_score_base, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + fake_score_transformer3d = WanTransformer3DModel.from_pretrained( + os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + + # Freeze vae and text_encoder and set generator_transformer3d to trainable + vae.requires_grad_(False) + text_encoder.requires_grad_(False) + generator_transformer3d.requires_grad_(False) + real_score_transformer3d.requires_grad_(False) + fake_score_transformer3d.requires_grad_(False) + if args.train_mode != "normal": + clip_image_encoder.requires_grad_(False) + + # CF Stage 3 init + if args.transformer_path is not None: + print(f"From checkpoint: {args.transformer_path}") + if args.transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.transformer_path) + else: + state_dict = torch.load(args.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 + 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 = generator_transformer3d.load_state_dict(state_dict, strict=False) + m, u = real_score_transformer3d.load_state_dict(state_dict, strict=False) + m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + if args.ode_transformer_path is not None: + print(f"From checkpoint: {args.ode_transformer_path}") + if args.ode_transformer_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.ode_transformer_path) + else: + state_dict = torch.load(args.ode_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 + 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 = generator_transformer3d.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + assert len(u) == 0 + + if args.vae_path is not None: + print(f"From checkpoint: {args.vae_path}") + if args.vae_path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + state_dict = load_file(args.vae_path) + else: + state_dict = torch.load(args.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)}") + assert len(u) == 0 + + # A good trainable modules is showed below now. + # For 3D Patch: trainable_modules = ['ff.net', 'pos_embed', 'attn2', 'proj_out', 'timepositionalencoding', 'h_position', 'w_position'] + # For 2D Patch: trainable_modules = ['ff.net', 'attn2', 'timepositionalencoding', 'h_position', 'w_position'] + generator_transformer3d.train() + fake_score_transformer3d.train() + if accelerator.is_main_process: + accelerator.print( + f"Trainable modules '{args.trainable_modules}'." + ) + for name, param in generator_transformer3d.named_parameters(): + for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + param.requires_grad = True + break + for name, param in fake_score_transformer3d.named_parameters(): + for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + param.requires_grad = True + break + + if use_gan: + if (not args.gan_feature_layers or any(i < 0 or i >= len(fake_score_transformer3d.blocks) for i in args.gan_feature_layers) + or args.gan_block_ffn_dim <= 0 or args.gan_block_num_heads <= 0 + or fake_score_transformer3d.dim % args.gan_block_num_heads != 0): + raise ValueError("GAN feature layers must exist and the positive head count must divide the fake_score hidden dimension.") + # Keep trainable discriminator master weights in fp32; forward uses autocast. + add_wan_gan_discriminator( + fake_score_transformer3d, + args.gan_feature_layers, + args.gan_block_ffn_dim, + args.gan_block_num_heads, + torch.float32, + seed=args.seed or 0, + ) + + # Create EMA for the generator_transformer3d. + if args.use_ema: + if zero_stage == 3: + raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.") + + if fsdp_plugin is not None: + if getattr(fsdp_plugin, "fsdp_version", 1) != 1: + raise NotImplementedError("FSDPEMA 仅支持 FSDP1。") + ema_module = WanTransformer3DModel_SelfForcing.from_config(generator_transformer3d.config) + ema_generator_transformer3d = FSDPEMA( + ema_module, source=generator_transformer3d, + accelerator=accelerator, fsdp_plugin=fsdp_plugin, + decay=args.ema_decay, + ) + else: + ema_generator_transformer3d = EMAModel( + generator_transformer3d.parameters(), + model_cls=WanTransformer3DModel_SelfForcing, + model_config=generator_transformer3d.config, + ) + + # `accelerate` 0.16.0 will have better support for customized saving + if version.parse(accelerate.__version__) >= version.parse("0.16.0"): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + if fsdp_plugin is not None or zero_stage == 3: + def save_model_hook(models, weights, output_dir): + accelerate_state_dict = accelerator.get_state_dict(models[-1], unwrap=True) + if accelerator.is_main_process: + from safetensors.torch import save_file + + safetensor_save_path = os.path.join(output_dir, f"diffusion_pytorch_model.safetensors") + accelerate_state_dict = { + k: v.to(dtype=torch.float32 if k.startswith("gan_head.") else weight_dtype) + for k, v in accelerate_state_dict.items() + } + save_file(accelerate_state_dict, safetensor_save_path, metadata={"format": "pt"}) + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + + def load_model_hook(models, input_dir): + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + else: + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if accelerator.is_main_process: + model = accelerator.unwrap_model(models[0]) + save_dir = os.path.join(output_dir, "transformer") + if hasattr(model, "gan_head"): + from safetensors.torch import save_file + model.save_config(save_dir) + save_file(model.state_dict(), os.path.join(save_dir, "diffusion_pytorch_model.safetensors"), metadata={"format": "pt"}) + else: + model.save_pretrained(save_dir) + if not args.use_deepspeed: + weights.pop() + + with open(os.path.join(output_dir, "sampler_pos_start.pkl"), 'wb') as file: + pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) + + def load_model_hook(models, input_dir): + for i in range(len(models)): + # pop models so that they are not loaded again + model = models.pop() + + # GAN heads are recreated from the same CLI configuration before resume. + if hasattr(model, "gan_head"): + from safetensors.torch import load_file + model.load_state_dict(load_file(os.path.join(input_dir, "transformer", "diffusion_pytorch_model.safetensors")), strict=True) + else: + load_model = WanTransformer3DModel.from_pretrained( + input_dir, + subfolder="transformer", + low_cpu_mem_usage=True, + ) + model.register_to_config(**load_model.config) + model.load_state_dict(load_model.state_dict()) + del load_model + + pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + loaded_number, _ = pickle.load(file) + batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0) + print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.") + + accelerator.register_save_state_pre_hook(save_model_hook) + accelerator.register_load_state_pre_hook(load_model_hook) + accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook) + accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook) + + if args.use_ema: + def save_ema_hook(models, weights, output_dir): + if isinstance(ema_generator_transformer3d, FSDPEMA) or accelerator.is_main_process: + ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema")) + + def load_ema_hook(models, input_dir): + ema_path = os.path.join(input_dir, "transformer_ema") + if isinstance(ema_generator_transformer3d, FSDPEMA): + ema_generator_transformer3d.load_pretrained(ema_path) + else: + _, ema_kwargs = WanTransformer3DModel_SelfForcing.load_config(ema_path, return_unused_kwargs=True) + load_model = WanTransformer3DModel_SelfForcing.from_pretrained( + input_dir, subfolder="transformer_ema", + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), + low_cpu_mem_usage=True, + torch_dtype=torch.float32, + ) + load_model = EMAModel(load_model.parameters(), model_cls=WanTransformer3DModel_SelfForcing, model_config=load_model.config) + load_model.load_state_dict(ema_kwargs) + ema_generator_transformer3d.load_state_dict(load_model.state_dict()) + ema_generator_transformer3d.to(accelerator.device) + del load_model + + accelerator.register_save_state_pre_hook(save_ema_hook) + accelerator.register_load_state_pre_hook(load_ema_hook) + + if args.gradient_checkpointing: + generator_transformer3d.enable_gradient_checkpointing() + fake_score_transformer3d.enable_gradient_checkpointing() + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if args.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + if args.scale_lr: + args.learning_rate = ( + args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes + ) + + # Initialize the optimizer + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`" + ) + + optimizer_cls = bnb.optim.AdamW8bit + elif args.use_came: + try: + from came_pytorch import CAME + except Exception: + raise ImportError( + "Please install came_pytorch to use CAME. You can do so by running `pip install came_pytorch`" + ) + + optimizer_cls = CAME + else: + optimizer_cls = torch.optim.AdamW + + trainable_params = list(filter(lambda p: p.requires_grad, generator_transformer3d.parameters())) + trainable_params_optim = [ + {'params': [], 'lr': args.learning_rate}, + {'params': [], 'lr': args.learning_rate / 2}, + ] + in_already = [] + for name, param in generator_transformer3d.named_parameters(): + high_lr_flag = False + if name in in_already: + continue + for trainable_module_name in args.trainable_modules: + if trainable_module_name in name: + in_already.append(name) + high_lr_flag = True + trainable_params_optim[0]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate}") + break + if high_lr_flag: + continue + for trainable_module_name in args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + in_already.append(name) + trainable_params_optim[1]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate / 2}") + break + + fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) + fake_trainable_params_optim = [ + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, + ] + in_already = [] + for name, param in fake_score_transformer3d.named_parameters(): + high_lr_flag = False + if name in in_already: + continue + for trainable_module_name in args.trainable_modules: + if trainable_module_name in name: + in_already.append(name) + high_lr_flag = True + fake_trainable_params_optim[0]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate_critic}") + break + if high_lr_flag: + continue + for trainable_module_name in args.trainable_modules_low_learning_rate: + if trainable_module_name in name: + in_already.append(name) + fake_trainable_params_optim[1]['params'].append(param) + if accelerator.is_main_process: + print(f"Set {name} to lr : {args.learning_rate_critic / 2}") + break + + if use_gan: + fake_trainable_params_optim[0]['params'].extend( + param for name, param in fake_score_transformer3d.named_parameters() + if name.startswith("gan_head.") and name not in in_already + ) + + if args.use_came: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + # weight_decay=args.adam_weight_decay, + betas=(0.9, 0.999, 0.9999), + eps=(1e-30, 1e-16) + ) + critic_optimizer = optimizer_cls( + fake_trainable_params_optim, + lr=args.learning_rate_critic, + # weight_decay=args.adam_weight_decay, + betas=(0.9, 0.999, 0.9999), + eps=(1e-30, 1e-16) + ) + else: + optimizer = optimizer_cls( + trainable_params_optim, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + critic_optimizer = optimizer_cls( + fake_trainable_params_optim, + lr=args.learning_rate_critic, + betas=(args.adam_beta1, args.adam_beta2), + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + ) + + # Get the training dataset + sample_n_frames_bucket_interval = vae.config.temporal_compression_ratio + + if args.fix_sample_size is not None and args.enable_bucket: + args.video_sample_size = max(max(args.fix_sample_size), args.video_sample_size) + args.image_sample_size = max(max(args.fix_sample_size), args.image_sample_size) + args.training_with_video_token_length = False + args.random_hw_adapt = False + + # Get the dataset + if args.train_mode != "normal" or args.use_teacher_forcing or use_gan: + train_dataset = ImageVideoDataset( + args.train_data_meta, args.train_data_dir, + video_sample_size=args.video_sample_size, video_sample_stride=args.video_sample_stride, video_sample_n_frames=args.video_sample_n_frames, + video_repeat=args.video_repeat, + image_sample_size=args.image_sample_size, + enable_bucket=args.enable_bucket, enable_inpaint=True if args.train_mode != "normal" else False, + ) + else: + train_dataset = TextDataset( + args.train_data_meta + ) + + def get_length_to_frame_num(token_length): + if args.image_sample_size > args.video_sample_size: + sample_sizes = list(range(args.video_sample_size, args.image_sample_size + 1, 128)) + + if sample_sizes[-1] != args.image_sample_size: + sample_sizes.append(args.image_sample_size) + else: + sample_sizes = [args.image_sample_size] + + length_to_frame_num = { + sample_size: min(token_length / sample_size / sample_size, args.video_sample_n_frames) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 for sample_size in sample_sizes + } + + return length_to_frame_num + + if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing or use_gan: + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = AspectRatioBatchImageVideoSampler( + sampler=RandomSampler(train_dataset, generator=batch_sampler_generator), dataset=train_dataset.dataset, + batch_size=args.train_batch_size, train_folder = args.train_data_dir, drop_last=True, + aspect_ratios=aspect_ratio_sample_size, + ) + + def collate_fn(examples): + # Get token length + target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size + length_to_frame_num = get_length_to_frame_num(target_token_length) + + # Create new output + new_examples = {} + new_examples["target_token_length"] = target_token_length + new_examples["pixel_values"] = [] + new_examples["text"] = [] + # Used in Inpaint mode + if args.train_mode != "normal": + new_examples["mask_pixel_values"] = [] + new_examples["mask"] = [] + new_examples["clip_pixel_values"] = [] + + # Get downsample ratio in image and videos + pixel_value = examples[0]["pixel_values"] + data_type = examples[0]["data_type"] + f, h, w, c = np.shape(pixel_value) + if data_type == 'image': + random_downsample_ratio = 1 if not args.random_hw_adapt else get_random_downsample_ratio(args.image_sample_size, image_ratio=[args.image_sample_size / args.video_sample_size], rng=rng) + + aspect_ratio_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.image_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + if args.random_hw_adapt: + if args.training_with_video_token_length: + local_min_size = np.min(np.array([np.mean(np.array([np.shape(example["pixel_values"])[1], np.shape(example["pixel_values"])[2]])) for example in examples])) + # The video will be resized to a lower resolution than its own. + choice_list = [length for length in list(length_to_frame_num.keys()) if length < local_min_size * 1.25] + if len(choice_list) == 0: + choice_list = list(length_to_frame_num.keys()) + if rng is None: + local_video_sample_size = np.random.choice(choice_list) + else: + local_video_sample_size = rng.choice(choice_list) + batch_video_length = length_to_frame_num[local_video_sample_size] + random_downsample_ratio = args.video_sample_size / local_video_sample_size + else: + random_downsample_ratio = get_random_downsample_ratio( + args.video_sample_size, rng=rng) + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + else: + random_downsample_ratio = 1 + batch_video_length = args.video_sample_n_frames + sample_n_frames_bucket_interval + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + aspect_ratio_random_crop_sample_size = {key : [x / 512 * args.video_sample_size / random_downsample_ratio for x in ASPECT_RATIO_RANDOM_CROP_512[key]] for key in ASPECT_RATIO_RANDOM_CROP_512.keys()} + + if args.fix_sample_size is not None: + fix_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + elif args.random_ratio_crop: + if rng is None: + random_sample_size = aspect_ratio_random_crop_sample_size[ + np.random.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + else: + random_sample_size = aspect_ratio_random_crop_sample_size[ + rng.choice(list(aspect_ratio_random_crop_sample_size.keys()), p = ASPECT_RATIO_RANDOM_CROP_PROB) + ] + random_sample_size = [int(x / 16) * 16 for x in random_sample_size] + else: + closest_size, closest_ratio = get_closest_ratio(h, w, ratios=aspect_ratio_sample_size) + closest_size = [int(x / 16) * 16 for x in closest_size] + + min_example_length = min( + [example["pixel_values"].shape[0] for example in examples] + ) + batch_video_length = int(min(batch_video_length, min_example_length)) + + # Magvae needs the number of frames to be 4n + 1. + batch_video_length = (batch_video_length - 1) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 + + # KV cache training requires latent frames divisible by num_frame_per_block + if args.use_kv_cache_training: + k = (batch_video_length - 1) // sample_n_frames_bucket_interval + if args.independent_first_frame: + # latent_frames - 1 = k must be divisible by num_frame_per_block + k = (k // args.num_frame_per_block) * args.num_frame_per_block + else: + # latent_frames = k + 1 must be divisible by num_frame_per_block + k = ((k + 1) // args.num_frame_per_block) * args.num_frame_per_block - 1 + batch_video_length = k * sample_n_frames_bucket_interval + 1 + + if batch_video_length <= 0: + batch_video_length = 1 + + for example in examples: + if args.fix_sample_size is not None: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + fix_sample_size = list(map(lambda x: int(x), fix_sample_size)) + transform = transforms.Compose([ + transforms.Resize(fix_sample_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(fix_sample_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + elif args.random_ratio_crop: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + b, c, h, w = pixel_values.size() + th, tw = random_sample_size + if th / tw > h / w: + nh = int(th) + nw = int(w / h * nh) + else: + nw = int(tw) + nh = int(h / w * nw) + + transform = transforms.Compose([ + transforms.Resize([nh, nw]), + transforms.CenterCrop([int(x) for x in random_sample_size]), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + else: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + + # Get adapt hw for resize + closest_size = list(map(lambda x: int(x), closest_size)) + if closest_size[0] / h > closest_size[1] / w: + resize_size = closest_size[0], int(w * closest_size[0] / h) + else: + resize_size = int(h * closest_size[1] / w), closest_size[1] + + transform = transforms.Compose([ + transforms.Resize(resize_size, interpolation=transforms.InterpolationMode.BILINEAR), # Image.BICUBIC + transforms.CenterCrop(closest_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), + ]) + + new_examples["pixel_values"].append(transform(pixel_values)[:batch_video_length]) + new_examples["text"].append(example["text"]) + + if args.train_mode != "normal": + mask = get_random_mask(new_examples["pixel_values"][-1].size(), image_start_only=True) + mask_pixel_values = new_examples["pixel_values"][-1] * (1 - mask) + # Wan 2.1 use 0 for masked pixels + # + torch.ones_like(new_examples["pixel_values"][-1]) * -1 * mask + new_examples["mask_pixel_values"].append(mask_pixel_values) + new_examples["mask"].append(mask) + + clip_pixel_values = new_examples["pixel_values"][-1][0].permute(1, 2, 0).contiguous() + clip_pixel_values = (clip_pixel_values * 0.5 + 0.5) * 255 + new_examples["clip_pixel_values"].append(clip_pixel_values) + + # Limit the number of frames to the same + new_examples["pixel_values"] = torch.stack([example for example in new_examples["pixel_values"]]) + if args.train_mode != "normal": + new_examples["mask_pixel_values"] = torch.stack([example for example in new_examples["mask_pixel_values"]]) + new_examples["mask"] = torch.stack([example for example in new_examples["mask"]]) + new_examples["clip_pixel_values"] = torch.stack([example for example in new_examples["clip_pixel_values"]]) + + # Encode prompts when enable_text_encoder_in_dataloader=True + if args.enable_text_encoder_in_dataloader: + prompt_ids = tokenizer( + new_examples['text'], + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=prompt_attention_mask.to("cpu"))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + new_examples['encoder_attention_mask'] = prompt_ids.attention_mask + new_examples['encoder_hidden_states'] = prompt_embeds + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to("cpu"), attention_mask=neg_prompt_attention_mask.to("cpu"))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + new_examples['neg_encoder_attention_mask'] = neg_prompt_ids.attention_mask + new_examples['neg_encoder_hidden_states'] = neg_prompt_embeds + + return new_examples + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=collate_fn, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + elif args.train_mode == "normal": + def collate_fn(examples): + new_examples = {} + new_examples["text"] = [] + for example in examples: + new_examples["text"].append(example["text"]) + + # Encode prompts when enable_text_encoder_in_dataloader=True + if args.enable_text_encoder_in_dataloader: + prompt_ids = tokenizer( + new_examples['text'], + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to("cpu"), attention_mask=prompt_attention_mask.to("cpu"))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + new_examples['encoder_attention_mask'] = prompt_ids.attention_mask + new_examples['encoder_hidden_states'] = prompt_embeds + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + max_length=args.tokenizer_max_length, + padding="max_length", + add_special_tokens=True, + truncation=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to("cpu"), attention_mask=neg_prompt_attention_mask.to("cpu"))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + new_examples['neg_encoder_attention_mask'] = neg_prompt_ids.attention_mask + new_examples['neg_encoder_hidden_states'] = neg_prompt_embeds + + return new_examples + + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = BatchSampler(RandomSampler(train_dataset, generator=batch_sampler_generator), batch_size=args.train_batch_size, drop_last=True) + + # DataLoaders creation: + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + collate_fn=collate_fn, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + else: + # DataLoaders creation: + batch_sampler_generator = torch.Generator().manual_seed(args.seed) + batch_sampler = ImageVideoSampler(RandomSampler(train_dataset, generator=batch_sampler_generator), train_dataset, args.train_batch_size) + train_dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_sampler=batch_sampler, + persistent_workers=True if args.dataloader_num_workers != 0 else False, + num_workers=args.dataloader_num_workers, + ) + + # Scheduler and math around the number of training steps. + overrode_max_train_steps = False + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if args.max_train_steps is None: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + overrode_max_train_steps = True + + lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + fake_score_lr_scheduler = get_scheduler( + args.lr_scheduler, + optimizer=critic_optimizer, + num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, + num_training_steps=args.max_train_steps * accelerator.num_processes, + ) + + # Prepare everything with our `accelerator`. + generator_transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + generator_transformer3d, optimizer, train_dataloader, lr_scheduler + ) + fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler= accelerator_fake_score_transformer3d.prepare( + fake_score_transformer3d, critic_optimizer, fake_score_lr_scheduler + ) + if fsdp_stage != 0 or zero_stage != 0: + from functools import partial + + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) + real_score_transformer3d = shard_fn(real_score_transformer3d) + if fsdp_stage != 0 or zero_stage != 0: + from functools import partial + + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) + text_encoder = shard_fn(text_encoder) + + # Move text_encode and vae to gpu and cast to weight_dtype + vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + real_score_transformer3d.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device if not args.low_vram else "cpu") + if args.train_mode != "normal": + clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype) + + if args.use_ema: + ema_generator_transformer3d.to(accelerator.device) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) + if overrode_max_train_steps: + args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) + + # We need to initialize the trackers we use, and also store our configuration. + # The trackers initializes automatically on the main process. + if accelerator.is_main_process: + tracker_config = dict(vars(args)) + keys_to_pop = [k for k, v in tracker_config.items() if isinstance(v, list)] + for k in keys_to_pop: + tracker_config.pop(k) + print(f"Removed tracker_config['{k}']") + accelerator.init_trackers(args.tracker_project_name, tracker_config) + + # Function for unwrapping if model was compiled with `torch.compile`. + def unwrap_model(model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + # Train! + total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps + + logger.info("***** Running training *****") + logger.info(f" Num examples = {len(train_dataset)}") + logger.info(f" Num Epochs = {args.num_train_epochs}") + logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {args.max_train_steps}") + global_step = 0 + first_epoch = 0 + + # Potentially load in the weights and states from a previous save + if args.resume_from_checkpoint: + if args.resume_from_checkpoint != "latest": + path = os.path.basename(args.resume_from_checkpoint) + else: + # Get the most recent checkpoint + dirs = os.listdir(args.output_dir) + dirs = [d for d in dirs if d.startswith("checkpoint")] + dirs = sorted(dirs, key=lambda x: int(x.split("-")[1])) + path = dirs[-1] if len(dirs) > 0 else None + + if path is None: + accelerator.print( + f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run." + ) + args.resume_from_checkpoint = None + initial_global_step = 0 + else: + global_step = int(path.split("-")[1]) + + initial_global_step = global_step + + pkl_path = os.path.join(os.path.join(args.output_dir, path), "sampler_pos_start.pkl") + if os.path.exists(pkl_path): + with open(pkl_path, 'rb') as file: + _, first_epoch = pickle.load(file) + else: + first_epoch = global_step // num_update_steps_per_epoch + print(f"Load pkl from {pkl_path}. Get first_epoch = {first_epoch}.") + + accelerator.print(f"Resuming from checkpoint {path}") + fake_score_path = os.path.join(path, "fake_score") + accelerator.load_state(os.path.join(args.output_dir, path)) + accelerator_fake_score_transformer3d.load_state(os.path.join(args.output_dir, fake_score_path)) + else: + initial_global_step = 0 + + progress_bar = PauseAwareTqdm( + range(0, args.max_train_steps), + initial=initial_global_step, + desc="Steps", + # Only show the progress bar once on each machine. + disable=not accelerator.is_local_main_process, + ) + + if args.multi_stream and args.train_mode != "normal": + # create extra cuda streams to speedup inpaint vae computation + vae_stream_1 = torch.cuda.Stream() + vae_stream_2 = torch.cuda.Stream() + else: + vae_stream_1 = None + vae_stream_2 = None + + idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling) + min_score_timestep = int(0.02 * noise_scheduler.config.num_train_timesteps) + max_score_timestep = int(0.98 * noise_scheduler.config.num_train_timesteps) + + def randomize_denoising_step_indices( + denoising_step_indices_list, + train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=0.3, + tail_margin=1, + ): + indices = list(denoising_step_indices_list) + n = len(indices) + tail_margin = max(int(tail_margin), 1) + + # The head stays fixed at the pure-noise start; the remaining steps jitter + # symmetrically around their base values so the expected schedule is unchanged. + result = [indices[0]] + for i in range(1, n): + gap_upper = indices[i - 1] - indices[i] + if i + 1 < n: + gap_lower = indices[i] - indices[i + 1] + max_jitter = int(min(gap_upper, gap_lower) * jitter_ratio) + else: + # Tail step: no lower neighbor (the base value sits at the clean end), + # so the jitter budget comes from the upward gap and the downward side + # is clamped by tail_margin. Keeping the tail off the schedule's cleanest + # position guarantees Decoupled DMD's tau_CA always has a cleaner slot. + max_jitter = int(gap_upper * jitter_ratio) + + if max_jitter > 0: + # NB: torch_rng may live on CUDA while randint's default output is CPU; + # use the global CPU RNG here, the result is broadcast from rank 0 anyway. + jitter = torch.randint( + -max_jitter, max_jitter + 1, (1,) + ).item() + else: + jitter = 0 + + value = indices[i] + jitter + # Strict monotonicity by construction (no post-hoc clamp repair). + value = min(value, result[i - 1] - 1) + if i == n - 1: + value = max(value, tail_margin) + result.append(value) + + result = [max(1, min(train_sampling_steps, x)) for x in result] + result = torch.tensor(result) + + if dist.is_initialized(): + result = result.to(accelerator.device) + dist.broadcast(result, src=0) + result = result.cpu() + return result + + # Official One-Forcing only creates the EMA shadow once global_step reaches + # ema_start_step; a resume past that point loads the shadow from the checkpoint + # instead of re-seeding it, so the flag starts true in that case. + ema_started = global_step >= args.ema_start_step + + for epoch in range(first_epoch, args.num_train_epochs): + if global_step >= args.max_train_steps: + break + train_dmd_loss = 0.0 + # Number of generator backward contributions since the last log flush; the + # generator only backprops every gen_update_interval batches, so its metrics + # must be averaged by contribution count, not by gradient_accumulation_steps. + train_gen_log_count = 0 + train_denoising_loss = 0.0 + train_gan_g_loss = 0.0 + train_gan_d_loss = 0.0 + batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) + for step, batch in enumerate(train_dataloader): + # Data batch sanity check + if args.train_mode != "normal" and epoch == first_epoch and step == 0: + pixel_values, texts = batch['pixel_values'].cpu(), batch['text'] + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + os.makedirs(os.path.join(args.output_dir, "sanity_check"), exist_ok=True) + for idx, (pixel_value, text) in enumerate(zip(pixel_values, texts)): + pixel_value = pixel_value[None, ...] + gif_name = '-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'{global_step}-{idx}' + save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/{gif_name[:10]}.mp4", rescale=True) + + clip_pixel_values, mask_pixel_values, texts = batch['clip_pixel_values'].cpu(), batch['mask_pixel_values'].cpu(), batch['text'] + mask_pixel_values = rearrange(mask_pixel_values, "b f c h w -> b c f h w") + for idx, (clip_pixel_value, pixel_value, text) in enumerate(zip(clip_pixel_values, mask_pixel_values, texts)): + pixel_value = pixel_value[None, ...] + Image.fromarray(np.uint8(clip_pixel_value)).save(f"{args.output_dir}/sanity_check/clip_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.png") + save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.mp4", rescale=True) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + if args.train_mode != "normal" or args.use_teacher_forcing or use_gan: + # Convert images to latent space + pixel_values = batch["pixel_values"].to(weight_dtype) + + # Increase the batch size when the length of the latent sequence of the current sample is small + if args.auto_tile_batch_size and args.training_with_video_token_length and zero_stage != 3: + if args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 16 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + pixel_values = torch.tile(pixel_values, (4, 1, 1, 1, 1)) + if args.enable_text_encoder_in_dataloader: + batch['encoder_hidden_states'] = torch.tile(batch['encoder_hidden_states'], (4, 1, 1)) + batch['encoder_attention_mask'] = torch.tile(batch['encoder_attention_mask'], (4, 1)) + batch['neg_encoder_hidden_states'] = torch.tile(batch['neg_encoder_hidden_states'], (4, 1, 1)) + batch['neg_encoder_attention_mask'] = torch.tile(batch['neg_encoder_attention_mask'], (4, 1)) + else: + batch['text'] = batch['text'] * 4 + elif args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 4 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + pixel_values = torch.tile(pixel_values, (2, 1, 1, 1, 1)) + if args.enable_text_encoder_in_dataloader: + batch['encoder_hidden_states'] = torch.tile(batch['encoder_hidden_states'], (2, 1, 1)) + batch['encoder_attention_mask'] = torch.tile(batch['encoder_attention_mask'], (2, 1)) + batch['neg_encoder_hidden_states'] = torch.tile(batch['neg_encoder_hidden_states'], (2, 1, 1)) + batch['neg_encoder_attention_mask'] = torch.tile(batch['neg_encoder_attention_mask'], (2, 1)) + else: + batch['text'] = batch['text'] * 2 + if args.train_mode != "normal": + clip_pixel_values = batch["clip_pixel_values"].to(weight_dtype) + mask_pixel_values = batch["mask_pixel_values"].to(weight_dtype) + mask = batch["mask"].to(weight_dtype) + # Increase the batch size when the length of the latent sequence of the current sample is small + if args.auto_tile_batch_size and args.training_with_video_token_length and zero_stage != 3: + if args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 16 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + clip_pixel_values = torch.tile(clip_pixel_values, (4, 1, 1, 1)) + mask_pixel_values = torch.tile(mask_pixel_values, (4, 1, 1, 1, 1)) + mask = torch.tile(mask, (4, 1, 1, 1, 1)) + elif args.video_sample_n_frames * args.token_sample_size * args.token_sample_size // 4 >= pixel_values.size()[1] * pixel_values.size()[3] * pixel_values.size()[4]: + clip_pixel_values = torch.tile(clip_pixel_values, (2, 1, 1, 1)) + mask_pixel_values = torch.tile(mask_pixel_values, (2, 1, 1, 1, 1)) + mask = torch.tile(mask, (2, 1, 1, 1, 1)) + + if args.random_frame_crop: + def _create_special_list(length): + if length == 1: + return [1.0] + if length >= 2: + last_element = 0.90 + remaining_sum = 1.0 - last_element + other_elements_value = remaining_sum / (length - 1) + special_list = [other_elements_value] * (length - 1) + [last_element] + return special_list + select_frames = [_tmp for _tmp in list(range(sample_n_frames_bucket_interval + 1, args.video_sample_n_frames + sample_n_frames_bucket_interval, sample_n_frames_bucket_interval))] + select_frames_prob = np.array(_create_special_list(len(select_frames))) + + if len(select_frames) != 0: + if rng is None: + temp_n_frames = np.random.choice(select_frames, p = select_frames_prob) + else: + temp_n_frames = rng.choice(select_frames, p = select_frames_prob) + else: + temp_n_frames = 1 + + # Magvae needs the number of frames to be 4n + 1. + temp_n_frames = (temp_n_frames - 1) // sample_n_frames_bucket_interval + 1 + + pixel_values = pixel_values[:, :temp_n_frames, :, :] + if args.train_mode != "normal": + mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :] + mask = mask[:, :temp_n_frames, :, :] + + # Keep all node same token length to accelerate the traning when resolution grows. + if args.keep_all_node_same_token_length: + if args.token_sample_size > 256: + numbers_list = list(range(256, args.token_sample_size + 1, 128)) + + if numbers_list[-1] != args.token_sample_size: + numbers_list.append(args.token_sample_size) + else: + numbers_list = [256] + numbers_list = [_number * _number * args.video_sample_n_frames for _number in numbers_list] + + actual_token_length = index_rng.choice(numbers_list) + actual_video_length = (min( + actual_token_length / pixel_values.size()[-1] / pixel_values.size()[-2], args.video_sample_n_frames + ) - 1) // sample_n_frames_bucket_interval * sample_n_frames_bucket_interval + 1 + actual_video_length = int(max(actual_video_length, 1)) + + # Magvae needs the number of frames to be 4n + 1. + actual_video_length = (actual_video_length - 1) // sample_n_frames_bucket_interval + 1 + + pixel_values = pixel_values[:, :actual_video_length, :, :] + if args.train_mode != "normal": + mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :] + mask = mask[:, :actual_video_length, :, :] + + if args.low_vram: + torch.cuda.empty_cache() + vae.to(accelerator.device) + if args.train_mode != "normal": + clip_image_encoder.to(accelerator.device) + real_score_transformer3d = real_score_transformer3d.to("cpu") + if not args.enable_text_encoder_in_dataloader: + text_encoder.to("cpu") + + with torch.no_grad(): + # This way is quicker when batch grows up + def _batch_encode_vae(pixel_values): + pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") + bs = args.vae_mini_batch + new_pixel_values = [] + for i in range(0, pixel_values.shape[0], bs): + pixel_values_bs = pixel_values[i : i + bs] + pixel_values_bs = vae.encode(pixel_values_bs)[0] + pixel_values_bs = pixel_values_bs.sample() + new_pixel_values.append(pixel_values_bs) + return torch.cat(new_pixel_values, dim = 0) + if args.use_teacher_forcing or use_gan: + clean_latents = _batch_encode_vae(pixel_values) + else: + clean_latents = None + + if args.train_mode != "normal": + # Encode inpaint latents. + mask_latents = _batch_encode_vae(mask_pixel_values) + if vae_stream_2 is not None: + torch.cuda.current_stream().wait_stream(vae_stream_2) + + mask = rearrange(mask, "b f c h w -> b c f h w") + mask = torch.concat( + [ + torch.repeat_interleave(mask[:, :, 0:1], repeats=4, dim=2), + mask[:, :, 1:] + ], dim=2 + ) + mask = mask.view(mask.shape[0], mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]) + mask = mask.transpose(1, 2) + mask = resize_mask(1 - mask, mask_latents) + + inpaint_latents = torch.concat([mask, mask_latents], dim=1) + + clip_context = [] + for clip_pixel_value in clip_pixel_values: + clip_image = Image.fromarray(np.uint8(clip_pixel_value.float().cpu().numpy())) + clip_image = TF.to_tensor(clip_image).sub_(0.5).div_(0.5).to(clip_image_encoder.device, weight_dtype) + _clip_context = clip_image_encoder([clip_image[:, None, :, :]]) + clip_context.append(_clip_context) + clip_context = torch.cat(clip_context) + + if args.use_teacher_forcing or use_gan: + target_shape = clean_latents.size() + else: + target_shape = mask_latents.size() + else: + text = batch['text'] + if args.fix_sample_size is not None: + local_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size] + num_frames = args.video_sample_n_frames + else: + if args.random_hw_adapt and args.training_with_video_token_length: + # Get token length + target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size + length_to_frame_num = get_length_to_frame_num(target_token_length) + + if rng is None: + local_length = np.random.choice(list(length_to_frame_num.keys())) + else: + local_length = rng.choice(list(length_to_frame_num.keys())) + num_frames = length_to_frame_num[local_length] + + aspect_ratio_sample_size = {key : [x / 512 * local_length for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + if rng is None: + aspect_ratio_key = np.random.choice(list(aspect_ratio_sample_size.keys())) + else: + aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys())) + local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] + else: + num_frames = args.video_sample_n_frames + + aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + if rng is None: + aspect_ratio_key = np.random.choice(list(aspect_ratio_sample_size.keys())) + else: + aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys())) + local_sample_size = aspect_ratio_sample_size[aspect_ratio_key] + local_sample_size = [int(x / 16) * 16 for x in local_sample_size] + + # Compute latent frame count + latent_num_frames = int((num_frames - 1) // vae.temporal_compression_ratio + 1) + + # Align latent_num_frames to num_frame_per_block for KV cache training + if args.use_kv_cache_training: + if args.independent_first_frame: + # latent_frames - 1 must be divisible by num_frame_per_block + k = latent_num_frames - 1 + k = (k // args.num_frame_per_block) * args.num_frame_per_block + latent_num_frames = k + 1 + else: + # latent_frames must be divisible by num_frame_per_block + latent_num_frames = (latent_num_frames // args.num_frame_per_block) * args.num_frame_per_block + latent_num_frames = max(latent_num_frames, args.num_frame_per_block) + + target_shape = ( + len(text), + vae.latent_channels, + latent_num_frames, + int(local_sample_size[0] // vae.spatial_compression_ratio), + int(local_sample_size[1] // vae.spatial_compression_ratio), + ) + clean_latents = None + + if args.low_vram: + vae.to('cpu') + real_score_transformer3d = real_score_transformer3d.to("cpu") + if args.train_mode != "normal": + clip_image_encoder.to('cpu') + torch.cuda.empty_cache() + if not args.enable_text_encoder_in_dataloader: + text_encoder.to(accelerator.device) + + if args.enable_text_encoder_in_dataloader: + prompt_embeds = batch['encoder_hidden_states'].to(device=accelerator.device) + neg_prompt_embeds = batch['neg_encoder_hidden_states'].to(device=accelerator.device) + else: + with torch.no_grad(): + prompt_ids = tokenizer( + batch['text'], + padding="max_length", + max_length=args.tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt" + ) + text_input_ids = prompt_ids.input_ids + prompt_attention_mask = prompt_ids.attention_mask + + seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long() + prompt_embeds = text_encoder(text_input_ids.to(accelerator.device), attention_mask=prompt_attention_mask.to(accelerator.device))[0] + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] + + neg_txt = [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text'] + ] + neg_prompt_ids = tokenizer( + neg_txt, + padding="max_length", + max_length=args.tokenizer_max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt" + ) + neg_text_input_ids = neg_prompt_ids.input_ids + neg_prompt_attention_mask = neg_prompt_ids.attention_mask + + neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long() + neg_prompt_embeds = text_encoder(neg_text_input_ids.to(accelerator.device), attention_mask=neg_prompt_attention_mask.to(accelerator.device))[0] + neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)] + + if args.low_vram: + generator_transformer3d = generator_transformer3d.to(accelerator.device) + real_score_transformer3d = real_score_transformer3d.to(accelerator.device) + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + if not args.enable_text_encoder_in_dataloader: + text_encoder.to('cpu') + torch.cuda.empty_cache() + + generator_update = step % args.gen_update_interval == 0 + # Enter the generator's accumulation context only on batches that actually + # backprop through the generator. Entering it on every batch would advance + # the accumulation counter gen_update_interval times faster than real + # generator gradients are produced; whenever gcd(gradient_accumulation_steps, + # gen_update_interval) > 1 the sync flag would then never coincide with a + # generator-update batch and optimizer.step() would silently never fire. + generator_accumulate_ctx = ( + accelerator.accumulate(generator_transformer3d) if generator_update else contextlib.nullcontext() + ) + with generator_accumulate_ctx: + def get_sigmas(timesteps, n_dim=4, dtype=torch.float32): + sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype) + schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device) + timesteps = timesteps.to(accelerator.device) + + step_indices = [ + torch.argmin(torch.abs(schedule_timesteps - t)).item() + for t in timesteps + ] + step_indices = torch.tensor(step_indices, device=accelerator.device) + sigma = sigmas[step_indices].flatten() + + while len(sigma.shape) < n_dim: + sigma = sigma.unsqueeze(-1) + return sigma + + def add_noise(latents, noise, timesteps): + sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype) + return (1.0 - sigmas) * latents + sigmas * noise + + def generate_and_sync_list(num_denoising_steps, device): + indices = torch.randint(low=0, high=num_denoising_steps, size=(1,), generator=torch_rng, device=device) + if dist.is_initialized(): + dist.broadcast(indices, src=0) + return indices.tolist() + + def convert_flow_pred_to_x0( + scheduler, + flow_pred: torch.Tensor, + xt: torch.Tensor, + timestep: torch.Tensor + ) -> torch.Tensor: + """ + Convert flow matching's prediction to x0 prediction. + Supports both 4D [B, C, H, W] and 5D [B, C, F, H, W] inputs. + """ + original_dtype = flow_pred.dtype + device = flow_pred.device + + flow_pred = flow_pred.double() + xt = xt.double() + timesteps = scheduler.timesteps.to(device).double() + sigmas = scheduler.sigmas.to(device).double() + timestep = timestep.to(device).double() + + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id] + + ndim = flow_pred.ndim + if ndim == 4: + sigma_t = sigma_t.view(-1, 1, 1, 1) + elif ndim == 5: + sigma_t = sigma_t.view(-1, 1, 1, 1, 1) + else: + raise ValueError(f"Expected 4D or 5D input, got {ndim}D tensor.") + + x0_pred = xt - sigma_t * flow_pred + return x0_pred.to(original_dtype) + + # Create discrete denoising steps (per-step, with optional randomization) + if getattr(args, 'randomize_step_indices', False): + random_indices = randomize_denoising_step_indices( + args.denoising_step_indices_list, + args.train_sampling_steps, + torch_rng, + accelerator, + jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30), + ) + else: + random_indices = torch.tensor(args.denoising_step_indices_list) + + denoising_step_list = noise_scheduler.timesteps[args.train_sampling_steps - random_indices] + + # --- Main Training Logic --- + bsz, channel, num_frames, height, width = target_shape + if generator_update: # generator_update computed before the accumulate ctx above + if args.use_kv_cache_training: + # Calculate frame_seq_length + patch_h, patch_w = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + frame_seq_length = (target_shape[3] * target_shape[4]) // (patch_h * patch_w) + + # Determine block structure with variable-length support + if not args.independent_first_frame: + assert num_frames % args.num_frame_per_block == 0 + max_num_blocks = num_frames // args.num_frame_per_block + assert args.score_num_frames % args.num_frame_per_block == 0 + min_num_blocks = args.score_num_frames // args.num_frame_per_block + else: + assert (num_frames - 1) % args.num_frame_per_block == 0 + max_num_blocks = (num_frames - 1) // args.num_frame_per_block + if args.score_num_frames > 1: + assert (args.score_num_frames - 1) % args.num_frame_per_block == 0 + min_num_blocks = (args.score_num_frames - 1) // args.num_frame_per_block + else: + min_num_blocks = 0 + + # Random sample number of blocks (Self-Forcing variable-length training) + if args.min_length_prob_bias > 0.0 and max_num_blocks > min_num_blocks: + # Weighted sampling: give min_num_blocks a higher probability + num_options = max_num_blocks - min_num_blocks + 1 + bias = min(args.min_length_prob_bias, 0.99) + remaining_prob = (1.0 - bias) / (num_options - 1) + probs = [remaining_prob] * num_options + probs[0] = bias # min_num_blocks gets the bias + probs_tensor = torch.tensor(probs, device=accelerator.device) + block_indices = torch.multinomial(probs_tensor, 1, generator=torch_rng) + num_generated_blocks = (min_num_blocks + block_indices).item() + else: + num_generated_blocks = torch.randint( + min_num_blocks, max_num_blocks + 1, (1,), + generator=torch_rng, device=accelerator.device + ).item() + if dist.is_initialized(): + _sync = torch.tensor([num_generated_blocks], device=accelerator.device) + dist.broadcast(_sync, src=0) + num_generated_blocks = _sync.item() + + all_num_frames = [args.num_frame_per_block] * num_generated_blocks + if args.independent_first_frame: + all_num_frames = [1] + all_num_frames + + num_generated_frames = sum(all_num_frames) + + # Initialize KV cache + num_layers = generator_transformer3d.config.num_layers + num_heads = generator_transformer3d.config.num_heads + head_dim = generator_transformer3d.config.dim // num_heads + text_len = 512 # T5 sequence length + + kv_cache = initialize_kv_cache_for_training( + batch_size=bsz, + num_frames=num_frames, + frame_seq_length=frame_seq_length, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=weight_dtype, + device=accelerator.device + ) + + crossattn_cache = initialize_crossattn_cache_for_training( + batch_size=bsz, + text_len=text_len, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=weight_dtype, + device=accelerator.device + ) + + # Block-by-block generation + generator_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + current_start_frame = 0 + num_input_frames = 0 # T2V mode + + # Use actual batch size from generator_noise (may differ due to SP) + actual_bsz = generator_noise.shape[0] + output_pred = torch.zeros_like(generator_noise) + + # Decide whether to use teacher forcing for this video (once per video, not per block) + use_teacher_forcing_step = ( + args.use_teacher_forcing and + torch.rand(1, generator=torch_rng, device=accelerator.device).item() < args.teacher_forcing_prob + ) + + # Same exit step across all blocks (matches original Self-Forcing default) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=accelerator.device)[0] + + # Only blocks in the last score_num_frames get gradient at exit step + # (matches Self-Forcing: start_gradient_frame_index = num_output_frames - 21) + start_gradient_frame_index = num_generated_frames - args.score_num_frames + + for block_idx, current_num_frames in enumerate(all_num_frames): + # Extract noise for current block + start_idx = current_start_frame - num_input_frames + end_idx = start_idx + current_num_frames + noisy_input = generator_noise[:, :, start_idx:end_idx] + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + [bsz, current_num_frames], + current_timestep, + device=noisy_input.device, + dtype=torch.int64 + ) + + # Only enable gradient for final step AND block within score window + if not is_final_step or current_start_frame < start_gradient_frame_index: + context_manager = torch.no_grad() + else: + context_manager = contextlib.nullcontext() + + with context_manager: + # Convert noisy_input to list format + noisy_input_list = [noisy_input[i] for i in range(bsz)] + + # Use full seq_len (consistent with inference code) + full_seq_len = frame_seq_length * num_frames + + generator_pred_block = generator_transformer3d( + x=noisy_input_list, + context=prompt_embeds, + t=timestep, + seq_len=full_seq_len, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=current_start_frame * frame_seq_length, + cache_start=None, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + # Stack list output to tensor: [B, C, F, H, W] + if isinstance(generator_pred_block, list): + generator_pred_block = torch.stack(generator_pred_block, dim=0) + + # Flatten timestep for convert_flow_pred_to_x0: [B, F] -> [B*F] + generator_pred_block = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=generator_pred_block, + xt=noisy_input, + timestep=timestep[:, 0] + ) + + if is_final_step: + break + + # Add noise for next step + next_timestep = denoising_step_list[index + 1] * torch.ones( + bsz, dtype=torch.long, device=noisy_input.device + ) + noisy_input = add_noise( + generator_pred_block, + torch.randn(generator_pred_block.shape, dtype=generator_pred_block.dtype, device=generator_pred_block.device, generator=torch_rng), + next_timestep + ) + + # Record output + output_pred[:, :, current_start_frame:current_start_frame + current_num_frames] = generator_pred_block + + # Update KV cache with clean context (consistent with inference: feed denoised_pred directly) + if block_idx < len(all_num_frames) - 1: + context_timestep = torch.ones([bsz, current_num_frames], device=accelerator.device, dtype=torch.int64) * args.context_noise + + # Use clean latents for teacher forcing, otherwise use denoised prediction directly + if use_teacher_forcing_step and clean_latents is not None: + context_input = clean_latents[:, :, start_idx:end_idx] + else: + context_input = generator_pred_block + + context_input_list = [context_input[i] for i in range(bsz)] + + # Use full seq_len (consistent with inference code) + full_seq_len = frame_seq_length * num_frames + + with torch.no_grad(): + generator_transformer3d( + x=context_input_list, + context=prompt_embeds, + t=context_timestep, + seq_len=full_seq_len, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=current_start_frame * frame_seq_length, + cache_start=None, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + current_start_frame += current_num_frames + + # Final output — slice generated frames (may be < num_frames for variable-length) + generator_pred_full = output_pred[:, :, :num_generated_frames] + + need_gradient_mask = (num_generated_frames != args.score_num_frames) + + # Slice for score computation: last score_num_frames frames + if num_generated_frames > args.score_num_frames: + # Re-encode boundary for cleaner score input + generator_pred_for_score, score_num_frames, _ = slice_for_score( + generator_pred_full, vae, weight_dtype, + score_num_frames=args.score_num_frames, + independent_first_frame=args.independent_first_frame, + ) + else: + generator_pred_for_score = generator_pred_full + score_num_frames = num_generated_frames + + score_mask = None + if need_gradient_mask: + score_mask = torch.ones_like(generator_pred_for_score, dtype=torch.bool) + if args.independent_first_frame: + score_mask[:, :, :1] = False + else: + score_mask[:, :, :args.num_frame_per_block] = False + + # generator_pred = the sliced version for DMD loss + generator_pred = generator_pred_for_score + seq_len = frame_seq_length * score_num_frames # Score always on fixed window + + else: + # === Block mask training (flex attention, no KV cache) === + # Block mask training: use flex attention to process entire video at once + # Note: for long videos, use KV cache mode instead + score_mask = None # Block mask mode: no gradient mask needed + if num_frames > args.score_num_frames: + raise ValueError( + f"Block mask mode does not support variable-length training " + f"(video produces {num_frames} latent frames > score_num_frames={args.score_num_frames}). " + f"Use --use_kv_cache_training for long video training." + ) + + patch_h_bm, patch_w_bm = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + frame_seqlen_bm = (height * width) // (patch_h_bm * patch_w_bm) + + # Standard backward simulation training + generator_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=generator_noise.device)[0] + + # Precompute seq_len once (same for all steps) + seq_len = frame_seqlen_bm * num_frames + + # Decide whether to use teacher forcing for this step + use_teacher_forcing_step = ( + args.use_teacher_forcing and + torch.rand(1, generator=torch_rng, device=accelerator.device).item() < args.teacher_forcing_prob + ) + + # Create appropriate block mask based on teacher forcing decision + if use_teacher_forcing_step and clean_latents is not None: + # Teacher forcing: clean + noisy sequence mask + accelerator.unwrap_model(generator_transformer3d).create_teacher_forcing_mask( + device=accelerator.device, + num_frames=num_frames, + frame_seqlen=frame_seqlen_bm, + num_frame_per_block=args.num_frame_per_block, + ) + # Prepare clean_x and aug_t for teacher forcing + clean_x = [clean_latents[i] for i in range(clean_latents.size(0))] + aug_t = torch.zeros(bsz, device=accelerator.device, dtype=torch.int64) + else: + # Standard causal mask + accelerator.unwrap_model(generator_transformer3d).create_block_mask_for_training( + num_frames=num_frames, + frame_seqlen=frame_seqlen_bm, + num_frame_per_block=args.num_frame_per_block, + independent_first_frame=args.independent_first_frame, + device=accelerator.device + ) + clean_x = None + aug_t = None + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + generator_noise.shape[:1], + current_timestep, + device=generator_noise.device, + dtype=torch.int64 + ) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + context_manager = torch.no_grad() if not is_final_step else contextlib.nullcontext() + + with context_manager: + # Convert to list format for transformer + generator_noise_list = [generator_noise[i] for i in range(bsz)] + clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None + + # Use block_mask for causal training + generator_pred = generator_transformer3d( + x=generator_noise_list, + context=prompt_embeds, + t=timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + clean_x=clean_x_list, + aug_t=aug_t, + ) + generator_pred = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=generator_pred, + xt=generator_noise, + timestep=timestep + ) + + if is_final_step: + break + + next_timestep = denoising_step_list[index + 1] * torch.ones( + generator_noise.shape[:1], dtype=torch.long, device=generator_noise.device + ) + generator_noise = add_noise( + generator_pred, + torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng), + next_timestep + ) + + # Official One-Forcing masks the DMD gradient to gradient_num_blocks trailing + # (or leading) frames of the score window; 0 keeps every scored frame unmasked. + gradient_window_frames = args.gradient_num_blocks * args.num_frame_per_block + if 0 < gradient_window_frames < generator_pred.shape[2]: + if score_mask is None: + score_mask = torch.ones_like(generator_pred, dtype=torch.bool) + if args.gradient_window_position == "first": + score_mask[:, :, gradient_window_frames:] = False + else: + score_mask[:, :, :generator_pred.shape[2] - gradient_window_frames] = False + + # Common code for both KV cache and block mask training + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + if args.score_timestep_clamp: + generator_timestep = generator_timestep.clamp(min_score_timestep, max_score_timestep) + generator_denoised_input = add_noise( + generator_pred, + torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng), + generator_timestep + ).detach().to(accelerator.device, dtype=weight_dtype) + + # Compute fake score + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device), torch.no_grad(): + fake_score_main_cond = fake_score_transformer3d( + x=generator_denoised_input, + context=prompt_embeds, + t=generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + fake_score_main_cond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_main_cond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + if args.fake_guidance_scale != 0.0: + fake_score_main_uncond = fake_score_transformer3d( + x=generator_denoised_input, + context=neg_prompt_embeds, + t=generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + fake_score_main_uncond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_main_uncond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + fake_score_main = fake_score_main_uncond + ( + fake_score_main_cond - fake_score_main_uncond + ) * args.fake_guidance_scale + else: + fake_score_main = fake_score_main_cond + + # Compute real score + real_score_main_cond = real_score_transformer3d( + x=generator_denoised_input, + context=prompt_embeds, + t=generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + real_score_main_cond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=real_score_main_cond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + real_score_main_uncond = real_score_transformer3d( + x=generator_denoised_input, + context=neg_prompt_embeds, + t=generator_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + real_score_main_uncond = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=real_score_main_uncond, + xt=generator_denoised_input, + timestep=generator_timestep + ) + + real_score_main = real_score_main_uncond + ( + real_score_main_cond - real_score_main_uncond + ) * args.real_guidance_scale + + # DMD loss + fake_to_real_grad = fake_score_main - real_score_main + generator_to_real_norm = generator_pred - real_score_main + normalizer = torch.abs(generator_to_real_norm).mean(dim=[1, 2, 3, 4], keepdim=True) + fake_to_real_grad = fake_to_real_grad / normalizer + fake_to_real_grad = torch.nan_to_num(fake_to_real_grad) + + # Apply gradient mask: only compute loss on unmasked elements + # (matches Self-Forcing dmd.py: F.mse_loss(x[mask], target[mask])) + if score_mask is not None: + dmd_loss = 0.5 * F.mse_loss( + generator_pred.double()[score_mask], + (generator_pred.double() - fake_to_real_grad.double()).detach()[score_mask], + reduction="mean" + ) + else: + dmd_loss = 0.5 * F.mse_loss( + generator_pred.double(), + (generator_pred.double() - fake_to_real_grad.double()).detach(), + reduction="mean" + ) + + avg_dmd_loss = accelerator.gather(dmd_loss.repeat(args.train_batch_size)).mean() + train_dmd_loss += avg_dmd_loss.item() + train_gen_log_count += 1 + + # --- One-Forcing Generator GAN Loss --- + generator_loss = dmd_loss + if args.gan_g_weight > 0: + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + gan_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + if args.score_timestep_clamp: + gan_timestep = gan_timestep.clamp(min_score_timestep, max_score_timestep) + gan_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng) + gan_inputs = [add_noise(generator_pred, gan_noise, gan_timestep)] + if args.relativistic_discriminator: + real_latents = clean_latents[:, :, :generator_pred.size(2)].to(generator_pred.dtype) + gan_inputs.append(add_noise(real_latents, gan_noise, gan_timestep)) + num_gan_inputs = len(gan_inputs) + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _, gan_logits = fake_score_transformer3d( + x=torch.cat(gan_inputs).to(dtype=weight_dtype), + context=repeat_batch(prompt_embeds, num_gan_inputs), + t=gan_timestep.repeat(num_gan_inputs), + seq_len=seq_len, + y=repeat_batch(inpaint_latents, num_gan_inputs) if args.train_mode != "normal" else None, + clip_fea=repeat_batch(clip_context, num_gan_inputs) if args.train_mode != "normal" else None, + gan_classify=True, + ) + gan_logits = gan_logits.float().chunk(num_gan_inputs) + if args.relativistic_discriminator: + gan_g_loss = F.softplus(-(gan_logits[0] - gan_logits[1])).mean() + else: + gan_g_loss = F.softplus(-gan_logits[0]).mean() + gan_g_loss = gan_g_loss * args.gan_g_weight + train_gan_g_loss += accelerator.gather(gan_g_loss.detach().repeat(args.train_batch_size)).mean().item() + generator_loss = generator_loss + gan_g_loss + + if args.low_vram: + real_score_transformer3d = real_score_transformer3d.to("cpu") + if args.gan_g_weight == 0: + fake_score_transformer3d = fake_score_transformer3d.to("cpu") + torch.cuda.empty_cache() + + accelerator.backward(generator_loss) + if accelerator.sync_gradients: + accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + if args.gan_g_weight > 0: + # Discard discriminator gradients from the generator GAN backward. + critic_optimizer.zero_grad() + + if args.low_vram: + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + torch.cuda.empty_cache() + + with accelerator_fake_score_transformer3d.accumulate(fake_score_transformer3d): + # --- Fake Critic Denoising Loss --- + + if args.use_kv_cache_training: + # KV cache mode: block-by-block generation + fake_score_critic_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + + # Calculate frame_seq_length + frame_seq_length = (target_shape[3] * target_shape[4]) // (patch_h * patch_w) + + # Determine block structure (variable-length, mirrors generator branch) + if not args.independent_first_frame: + max_num_blocks_critic = num_frames // args.num_frame_per_block + min_num_blocks_critic = args.score_num_frames // args.num_frame_per_block + else: + max_num_blocks_critic = (num_frames - 1) // args.num_frame_per_block + if args.score_num_frames > 1: + min_num_blocks_critic = (args.score_num_frames - 1) // args.num_frame_per_block + else: + min_num_blocks_critic = 0 + + # Random sample number of blocks (mirrors generator's variable-length training) + if args.min_length_prob_bias > 0.0 and max_num_blocks_critic > min_num_blocks_critic: + num_options = max_num_blocks_critic - min_num_blocks_critic + 1 + bias = min(args.min_length_prob_bias, 0.99) + remaining_prob = (1.0 - bias) / (num_options - 1) + probs = [remaining_prob] * num_options + probs[0] = bias # min_num_blocks_critic gets the bias + probs_tensor = torch.tensor(probs, device=accelerator.device) + block_indices = torch.multinomial(probs_tensor, 1, generator=torch_rng) + num_generated_blocks_critic = (min_num_blocks_critic + block_indices).item() + else: + num_generated_blocks_critic = torch.randint( + min_num_blocks_critic, max_num_blocks_critic + 1, (1,), + generator=torch_rng, device=accelerator.device + ).item() + if dist.is_initialized(): + _sync = torch.tensor([num_generated_blocks_critic], device=accelerator.device) + dist.broadcast(_sync, src=0) + num_generated_blocks_critic = _sync.item() + + all_num_frames = [args.num_frame_per_block] * num_generated_blocks_critic + if args.independent_first_frame: + all_num_frames = [1] + all_num_frames + num_generated_frames_critic = sum(all_num_frames) + + # Initialize KV cache + num_layers = generator_transformer3d.config.num_layers + num_heads = generator_transformer3d.config.num_heads + head_dim = generator_transformer3d.config.dim // num_heads + text_len = 512 + + critic_kv_cache = initialize_kv_cache_for_training( + batch_size=bsz, + num_frames=num_frames, + frame_seq_length=frame_seq_length, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=weight_dtype, + device=accelerator.device + ) + + critic_crossattn_cache = initialize_crossattn_cache_for_training( + batch_size=bsz, + text_len=text_len, + num_layers=num_layers, + num_heads=num_heads, + head_dim=head_dim, + dtype=weight_dtype, + device=accelerator.device + ) + + current_start_frame = 0 + num_input_frames = 0 + output_pred = torch.zeros_like(fake_score_critic_noise) + + # Decide whether to use teacher forcing for this video + use_teacher_forcing_step = ( + args.use_teacher_forcing and + torch.rand(1, generator=torch_rng, device=accelerator.device).item() < args.teacher_forcing_prob + ) + + # Same exit step across all blocks (matches original Self-Forcing default) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=accelerator.device)[0] + + for block_idx, current_num_frames in enumerate(all_num_frames): + start_idx = current_start_frame - num_input_frames + end_idx = start_idx + current_num_frames + noisy_input = fake_score_critic_noise[:, :, start_idx:end_idx] + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + [bsz, current_num_frames], + current_timestep, + device=noisy_input.device, + dtype=torch.int64 + ) + + context_manager = torch.no_grad() + + with context_manager: + noisy_input_list = [noisy_input[i] for i in range(bsz)] + + # Use full seq_len (consistent with inference code) + full_seq_len = frame_seq_length * num_frames + + fake_score_denoised_pred_block = generator_transformer3d( + x=noisy_input_list, + context=prompt_embeds, + t=timestep, + seq_len=full_seq_len, + kv_cache=critic_kv_cache, + crossattn_cache=critic_crossattn_cache, + current_start=current_start_frame * frame_seq_length, + cache_start=None, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + # Stack list output to tensor: [B, C, F, H, W] + if isinstance(fake_score_denoised_pred_block, list): + fake_score_denoised_pred_block = torch.stack(fake_score_denoised_pred_block, dim=0) + + fake_score_denoised_pred_block = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_denoised_pred_block, + xt=noisy_input, + timestep=timestep[:, 0] + ) + + if is_final_step: + break + + next_timestep = denoising_step_list[index + 1] * torch.ones( + bsz, dtype=torch.long, device=noisy_input.device + ) + noisy_input = add_noise( + fake_score_denoised_pred_block, + torch.randn(fake_score_denoised_pred_block.shape, dtype=fake_score_denoised_pred_block.dtype, device=fake_score_denoised_pred_block.device, generator=torch_rng), + next_timestep + ) + + output_pred[:, :, current_start_frame:current_start_frame + current_num_frames] = fake_score_denoised_pred_block + + # Update KV cache with clean context (consistent with inference: feed denoised_pred directly) + if block_idx < len(all_num_frames) - 1: + context_timestep = torch.ones([bsz, current_num_frames], device=accelerator.device, dtype=torch.int64) * args.context_noise + + # Use clean latents for teacher forcing, otherwise use denoised prediction directly + if use_teacher_forcing_step and clean_latents is not None: + context_input = clean_latents[:, :, start_idx:end_idx] + else: + context_input = fake_score_denoised_pred_block + + context_input_list = [context_input[i] for i in range(bsz)] + + # Use full seq_len (consistent with inference code) + full_seq_len = frame_seq_length * num_frames + + with torch.no_grad(): + generator_transformer3d( + x=context_input_list, + context=prompt_embeds, + t=context_timestep, + seq_len=full_seq_len, + kv_cache=critic_kv_cache, + crossattn_cache=critic_crossattn_cache, + current_start=current_start_frame * frame_seq_length, + cache_start=None, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + current_start_frame += current_num_frames + + fake_score_denoised_pred_full = output_pred[:, :, :num_generated_frames_critic] + + # Slice for critic score: last score_num_frames frames + if num_generated_frames_critic > args.score_num_frames: + fake_score_denoised_pred, critic_score_num_frames, _ = slice_for_score( + fake_score_denoised_pred_full, vae, weight_dtype, + score_num_frames=args.score_num_frames, + independent_first_frame=args.independent_first_frame, + ) + else: + fake_score_denoised_pred = fake_score_denoised_pred_full + critic_score_num_frames = num_generated_frames_critic + + seq_len = frame_seq_length * critic_score_num_frames + + else: + with torch.no_grad(): + # Block mask mode: use flex attention to process entire video at once + + patch_h_bm, patch_w_bm = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + frame_seqlen_bm = (height * width) // (patch_h_bm * patch_w_bm) + seq_len = frame_seqlen_bm * num_frames + + fake_score_critic_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype) + num_denoising_steps = len(denoising_step_list) + final_step_index = generate_and_sync_list(num_denoising_steps, device=fake_score_critic_noise.device)[0] + + # Decide whether to use teacher forcing for this step + use_teacher_forcing_step = ( + args.use_teacher_forcing and + torch.rand(1, generator=torch_rng, device=accelerator.device).item() < args.teacher_forcing_prob + ) + + # Create appropriate block mask based on teacher forcing decision + if use_teacher_forcing_step and clean_latents is not None: + # Teacher forcing: clean + noisy sequence mask + accelerator.unwrap_model(generator_transformer3d).create_teacher_forcing_mask( + device=accelerator.device, + num_frames=num_frames, + frame_seqlen=frame_seqlen_bm, + num_frame_per_block=args.num_frame_per_block, + ) + clean_x = [clean_latents[i] for i in range(clean_latents.size(0))] + aug_t = torch.zeros(bsz, device=accelerator.device, dtype=torch.int64) + else: + # Standard causal mask + accelerator.unwrap_model(generator_transformer3d).create_block_mask_for_training( + num_frames=num_frames, + frame_seqlen=frame_seqlen_bm, + num_frame_per_block=args.num_frame_per_block, + independent_first_frame=args.independent_first_frame, + device=accelerator.device + ) + clean_x = None + aug_t = None + + for index, current_timestep in enumerate(denoising_step_list): + is_final_step = (index == final_step_index) + timestep = torch.full( + fake_score_critic_noise.shape[:1], + current_timestep, + device=fake_score_critic_noise.device, + dtype=torch.int64 + ) + + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + # Convert to list format for transformer + fake_score_critic_noise_list = [fake_score_critic_noise[i] for i in range(bsz)] + clean_x_list = [clean_latents[i] for i in range(bsz)] if clean_x is not None else None + + fake_score_denoised_pred = generator_transformer3d( + x=fake_score_critic_noise_list, + context=prompt_embeds, + t=timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + clean_x=clean_x_list, + aug_t=aug_t, + ) + fake_score_denoised_pred = convert_flow_pred_to_x0( + scheduler=noise_scheduler, + flow_pred=fake_score_denoised_pred, + xt=fake_score_critic_noise, + timestep=timestep + ) + + if is_final_step: + break + + next_timestep = denoising_step_list[index + 1] * torch.ones( + fake_score_critic_noise.shape[:1], + dtype=torch.long, + device=fake_score_critic_noise.device + ) + + fake_score_critic_noise = add_noise( + fake_score_denoised_pred, + torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng), + next_timestep + ) + + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + critic_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + if args.score_timestep_clamp: + critic_timestep = critic_timestep.clamp(min_score_timestep, max_score_timestep) + critic_noise = torch.randn(fake_score_denoised_pred.shape, dtype=fake_score_denoised_pred.dtype, device=fake_score_denoised_pred.device, generator=torch_rng) + + fake_score_denoised_input = add_noise( + fake_score_denoised_pred, + critic_noise, + critic_timestep + ) + + # Generated and real latents share critic_noise / critic_timestep. Prepare the + # real input and the R1/R2 perturbations once; they are consumed by the critic + # forward below and by the R1/R2 penalties after the main backward(s). + real_noisy = None + r1_perturbation = None + r2_perturbation = None + if use_gan: + real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype) + real_noisy = add_noise(real_latents, critic_noise, critic_timestep) + del real_latents + if args.r1_weight > 0: + r1_perturbation = args.r1_sigma * torch.randn( + real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng + ) + if args.r2_weight > 0: + r2_perturbation = args.r2_sigma * torch.randn( + fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype, + device=fake_score_denoised_input.device, generator=torch_rng + ) + + def custom_mse_loss(noise_pred, target, weighting=None, threshold=50): + noise_pred = noise_pred.float() + target = target.float() + diff = noise_pred - target + mse_loss = F.mse_loss(noise_pred, target, reduction='none') + mask = (diff.abs() <= threshold).float() + masked_loss = mse_loss * mask + if weighting is not None: + masked_loss = masked_loss * weighting + final_loss = masked_loss.mean() + return final_loss + + # --- One-Forcing Discriminator GAN Loss --- + if use_gan and not args.relativistic_discriminator: + # Memory-optimized critic: the non-relativistic gan_d_loss = + # softplus(-real) + softplus(fake) is separable, so run fake and real as two + # sequential 1x-batch forwards and backward the fake pass first. Backwarding + # fake frees its gradient-checkpointing block inputs before the real forward, + # dropping the packed fake+real+R1+R2 forward peak (up to 4x batch) down to + # 1x batch here; R1/R2 re-forward their own 2x batch below. Both backward(s) + # accumulate into fake_score.grad, giving the identical update as the packed + # single-forward critic. + # -- Fake pass (1x batch): denoising output + fake discriminator logit -- + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d( + x=fake_score_denoised_input, + context=prompt_embeds, + t=critic_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + gan_classify=True, + ) + denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred) + avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean() + train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps + + fake_logit = fake_gan_logits.float() + fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight + critic_loss = denoising_loss + fake_d_loss + accelerator_fake_score_transformer3d.backward(critic_loss) + # Fake graph freed: the transformer blocks' stored inputs are released here. + fake_d_detached = fake_d_loss.detach() + del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output + + # -- Real pass (1x batch): discriminator logit only -- + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _, real_gan_logits = fake_score_transformer3d( + x=real_noisy, + context=prompt_embeds, + t=critic_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + gan_classify=True, + ) + real_logit = real_gan_logits.float() + real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight + accelerator_fake_score_transformer3d.backward(real_d_loss) + # Real graph freed. Combine the two detached halves for logging. + gan_d_loss = fake_d_detached + real_d_loss.detach() + train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item() + del fake_d_detached, real_gan_logits, real_logit, real_d_loss + else: + # Combined path: 2x-batch forward for relativistic GAN (both logits coupled by + # softplus(-(real - fake))), or 1x-batch forward when GAN is disabled. + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + if use_gan: + fake_score_denoised_output, gan_logits = fake_score_transformer3d( + x=torch.cat([fake_score_denoised_input, real_noisy]), + context=repeat_batch(prompt_embeds, 2), + t=critic_timestep.repeat(2), + seq_len=seq_len, + y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None, + clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None, + gan_classify=True, + ) + fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0] + gan_logits = gan_logits.float().chunk(2) + else: + fake_score_denoised_output = fake_score_transformer3d( + x=fake_score_denoised_input, + context=prompt_embeds, + t=critic_timestep, + seq_len=seq_len, + y=inpaint_latents if args.train_mode != "normal" else None, + clip_fea=clip_context if args.train_mode != "normal" else None, + ) + + denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred) + avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean() + train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps + + critic_loss = denoising_loss + if use_gan: + fake_logit, real_logit = gan_logits[0], gan_logits[1] + if args.relativistic_discriminator: + gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean() + else: + gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean() + gan_d_loss = gan_d_loss * args.gan_d_weight + train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item() + critic_loss = critic_loss + gan_d_loss + + accelerator_fake_score_transformer3d.backward(critic_loss) + del critic_loss + + # R1/R2 gradient penalties: each re-forwards its own 2x batch after the main + # backward(s) and finite-differences the perturbed logit against the clean one, + # matching the packed-formula value while keeping the peak at 2x batch. + if use_gan and args.r1_weight > 0: + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _, r1_logits = fake_score_transformer3d( + x=torch.cat([real_noisy + r1_perturbation, real_noisy]), + context=repeat_batch(prompt_embeds, 2), + t=critic_timestep.repeat(2), + seq_len=seq_len, + y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None, + clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None, + gan_classify=True, + ) + r1_logits = r1_logits.float().chunk(2) + r1_loss = args.r1_weight * (r1_logits[0] - r1_logits[1]).pow(2).mean() + accelerator_fake_score_transformer3d.backward(r1_loss) + del r1_loss + + if use_gan and args.r2_weight > 0: + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + _, r2_logits = fake_score_transformer3d( + x=torch.cat([fake_score_denoised_input + r2_perturbation, fake_score_denoised_input]), + context=repeat_batch(prompt_embeds, 2), + t=critic_timestep.repeat(2), + seq_len=seq_len, + y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None, + clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None, + gan_classify=True, + ) + r2_logits = r2_logits.float().chunk(2) + r2_loss = 0.5 * args.r2_weight * ((r2_logits[0] - r2_logits[1]) / args.r2_sigma).pow(2).mean() + accelerator_fake_score_transformer3d.backward(r2_loss) + del r2_loss + + if accelerator_fake_score_transformer3d.sync_gradients: + accelerator_fake_score_transformer3d.clip_grad_norm_(fake_trainable_params, args.max_grad_norm) + critic_optimizer.step() + fake_score_lr_scheduler.step() + critic_optimizer.zero_grad() + + if args.low_vram: + fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) + generator_transformer3d = generator_transformer3d.to(accelerator.device) + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + + # Official One-Forcing averages the generator EMA only on steps where the + # generator actually ran an optimizer update (every gen_update_interval steps). + # Gating by generator_update stops critic-only rounds from repeatedly pulling + # the shadow toward an unchanged generator. (GAN forces grad_accum=1, so + # generator_update is True exactly on every gen_update_interval-th global step.) + if args.use_ema and generator_update and ema_started: + ema_generator_transformer3d.step(generator_transformer3d.parameters()) + progress_bar.update(1) + global_step += 1 + if args.use_ema and not ema_started and global_step >= args.ema_start_step: + # Official creates the shadow at ema_start_step from the live weights of + # that step; re-seeding here drops the warm-start mass a step-0 shadow + # would otherwise carry through a short run. + ema_generator_transformer3d.reinit_from(generator_transformer3d.parameters()) + ema_started = True + logs = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)} + if use_gan: + logs.update({"train_gan_g_loss": train_gan_g_loss / max(train_gen_log_count, 1), "train_gan_d_loss": train_gan_d_loss}) + accelerator.log(logs, step=global_step) + train_dmd_loss = 0.0 + train_gen_log_count = 0 + train_denoising_loss = 0.0 + train_gan_g_loss = 0.0 + train_gan_d_loss = 0.0 + + if global_step % args.checkpointing_steps == 0: + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + # _before_ saving state, check if this save would set us over the `checkpoints_total_limit` + if args.checkpoints_total_limit is not None: + checkpoints = os.listdir(args.output_dir) + checkpoints = [d for d in checkpoints if d.startswith("checkpoint")] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + + # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= args.checkpoints_total_limit: + num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint) + shutil.rmtree(removing_checkpoint) + + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + # Keep the checkpoint out of the progress bar rate: a minute-long save would + # otherwise land in the next step's interval and be shown as a slow step. The + # save also stages the whole state in host RAM (safetensors materializes every + # tensor as bytes) and leaves the freed blocks in the allocator caches, so the + # cache flushes run inside the same window. + with progress_bar.paused(): + fake_score_save_path = os.path.join(save_path, "fake_score") + accelerator.save_state(save_path) + accelerator_fake_score_transformer3d.save_state(fake_score_save_path) + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + logger.info(f"Saved state to {save_path}") + + if args.validation_prompts is not None and global_step % args.validation_steps == 0: + with progress_bar.paused(): + # Validation scores the live generator: official One-Forcing runs no + # in-training EMA validation and its infer.sh scores `generator` by + # default, so swapping the shadow in here only hid the live weights + # behind a lagging average. + log_validation( + vae, + text_encoder, + tokenizer, + clip_image_encoder, + generator_transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + + logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + + if global_step >= args.max_train_steps: + break + + if global_step >= args.max_train_steps: + break + + if args.validation_prompts is not None and epoch % args.validation_epochs == 0: + with progress_bar.paused(): + # Live-generator validation, same rationale as the step-wise site above. + log_validation( + vae, + text_encoder, + tokenizer, + clip_image_encoder, + generator_transformer3d, + args, + config, + accelerator, + weight_dtype, + global_step, + ) + + # Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever + # something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto + # the smoothed rate, so the worker warm-up and the first dataloader fetch do not dilute it. + progress_bar.close() + + # Create the pipeline using the trained modules and save it. + accelerator.wait_for_everyone() + if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") + fake_score_save_path = os.path.join(save_path, "fake_score") + accelerator.save_state(save_path) + accelerator_fake_score_transformer3d.save_state(fake_score_save_path) + logger.info(f"Saved state to {save_path}") + + accelerator.end_training() + + +if __name__ == "__main__": + main() diff --git a/scripts/wan2.1_one_forcing/train_one_forcing.sh b/scripts/wan2.1_one_forcing/train_one_forcing.sh new file mode 100644 index 0000000..5509513 --- /dev/null +++ b/scripts/wan2.1_one_forcing/train_one_forcing.sh @@ -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 "." diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md index fd71d6e..f79abc4 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md @@ -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 diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md index 4224fb6..66ffcdc 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL_zh-CN.md @@ -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 相同 diff --git a/scripts/wan2.1_self_forcing/train_distill.sh b/scripts/wan2.1_self_forcing/train_distill.sh index b46f570..4be78c0 100644 --- a/scripts/wan2.1_self_forcing/train_distill.sh +++ b/scripts/wan2.1_self_forcing/train_distill.sh @@ -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" \ diff --git a/videox_fun/models/attention_utils.py b/videox_fun/models/attention_utils.py index 969b3a9..0330e2f 100644 --- a/videox_fun/models/attention_utils.py +++ b/videox_fun/models/attention_utils.py @@ -5,6 +5,7 @@ import warnings import torch import torch.nn as nn +import torch.nn.functional as F try: import flash_attn_interface diff --git a/videox_fun/models/wan_tae.py b/videox_fun/models/wan_tae.py index 8a15ba6..f5df288 100644 --- a/videox_fun/models/wan_tae.py +++ b/videox_fun/models/wan_tae.py @@ -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 权重放在与 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) diff --git a/videox_fun/pipeline/pipeline_wan_self_forcing.py b/videox_fun/pipeline/pipeline_wan_self_forcing.py index 4336b12..5e44da8 100644 --- a/videox_fun/pipeline/pipeline_wan_self_forcing.py +++ b/videox_fun/pipeline/pipeline_wan_self_forcing.py @@ -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 diff --git a/videox_fun/utils/fsdp_utils.py b/videox_fun/utils/fsdp_utils.py index b3a15be..d85e724 100644 --- a/videox_fun/utils/fsdp_utils.py +++ b/videox_fun/utils/fsdp_utils.py @@ -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