From e7afc604e1795c11cebca7dcd1ef75683ce519d6 Mon Sep 17 00:00:00 2001 From: bubbliiiing <3323290568@qq.com> Date: Fri, 23 Oct 2026 13:10:15 +0800 Subject: [PATCH] Update self forcing gan and causal forcing precision --- examples/fastvideo/predict_t2v_minimax_h3.py | 331 +++ examples/fastvideo/predict_t2v_wan2.1.py | 313 +++ scripts/fastvideo/README_TRAIN_DISTILL.md | 648 +++++ .../fastvideo/README_TRAIN_DISTILL_zh-CN.md | 648 +++++ scripts/fastvideo/train_distill_wan2.1.py | 2411 +++++++++++++++++ scripts/fastvideo/train_distill_wan2.1.sh | 48 + .../README_TRAIN_AR_DIFFUSION.md | 2 + .../README_TRAIN_AR_DIFFUSION_zh-CN.md | 2 + ...README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md | 2 + ..._TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md | 2 + .../README_TRAIN_CAUSAL_DMD_DISTILL.md | 50 +- .../README_TRAIN_CAUSAL_DMD_DISTILL_zh-CN.md | 74 +- .../train_ar_diffusion.py | 13 + .../train_causal_consistency_distill.py | 15 +- .../wan2.1_causal_forcing/train_causal_dmd.py | 483 +++- .../README_TRAIN_DISTILL.md | 23 + .../README_TRAIN_DISTILL_zh-CN.md | 23 + .../wan2.1_self_forcing/README_TRAIN_ODE.md | 3 + .../README_TRAIN_ODE_zh-CN.md | 3 + scripts/wan2.1_self_forcing/train_distill.py | 344 ++- scripts/wan2.1_self_forcing/train_ode.py | 10 + videox_fun/dist/fsdp.py | 13 +- videox_fun/models/__init__.py | 5 + videox_fun/models/attention_kernel.py | 565 ++++ videox_fun/models/attention_utils.py | 501 +++- .../fastvideo_minimax_h3_transformer3d.py | 256 ++ .../models/fastvideo_wan_transformer3d.py | 204 ++ videox_fun/models/wan_gan_discriminator.py | 108 + videox_fun/models/wan_tae.py | 368 ++- videox_fun/models/wan_transformer3d.py | 10 +- 30 files changed, 7265 insertions(+), 213 deletions(-) create mode 100644 examples/fastvideo/predict_t2v_minimax_h3.py create mode 100644 examples/fastvideo/predict_t2v_wan2.1.py create mode 100644 scripts/fastvideo/README_TRAIN_DISTILL.md create mode 100644 scripts/fastvideo/README_TRAIN_DISTILL_zh-CN.md create mode 100644 scripts/fastvideo/train_distill_wan2.1.py create mode 100644 scripts/fastvideo/train_distill_wan2.1.sh create mode 100644 videox_fun/models/fastvideo_minimax_h3_transformer3d.py create mode 100644 videox_fun/models/fastvideo_wan_transformer3d.py create mode 100644 videox_fun/models/wan_gan_discriminator.py diff --git a/examples/fastvideo/predict_t2v_minimax_h3.py b/examples/fastvideo/predict_t2v_minimax_h3.py new file mode 100644 index 0000000..c13866e --- /dev/null +++ b/examples/fastvideo/predict_t2v_minimax_h3.py @@ -0,0 +1,331 @@ +import os +import sys + +import torch + +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 (AutoencoderKLMiniMaxH3, + AutoencoderKLMiniMaxH3Audio, + FastVideoMiniMaxH3Transformer3DModel, + Qwen2TokenizerFast, + Qwen3VLForConditionalGeneration, + Qwen3VLProcessor) +from videox_fun.pipeline import MiniMaxH3Pipeline +from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode, + convert_model_weight_to_float8, merge_lora, + save_videos_with_audio_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. +# The MiniMax-H3 stack (~19B transformer + the ~62GB Qwen3-VL conditioner) does not fit a single +# GPU at full load, so keep an offloading mode here. +GPU_memory_mode = "model_group_offload" +# Multi GPUs config +# NOTE: the VSA sparse kernel has no sequence-parallel path, so a FastVideo-H3 student runs +# single-GPU only; keep ulysses_degree = ring_degree = 1 (enable_multi_gpus_inference() raises +# otherwise, rather than silently falling back to the dense multi-GPU processor). +ulysses_degree = 1 +ring_degree = 1 +# Use FSDP to save more GPU memory in multi gpus. Unused at ulysses_degree = ring_degree = 1. +fsdp_dit = False +fsdp_text_encoder = False +# Compile will give a speedup in fixed resolution and need a little GPU memory. +# The compile_dit is not compatible with sequential_cpu_offload. +compile_dit = False + +# model path +model_name = "models/Diffusion_Transformer/MiniMax-H3" + +# ---- FastVideo VSA sparse-student config -------------------------------------------------- +# Official FastH3 8-step DMD student (FastVideo/FastVideo-FastH3-8-Step-V2): a MiniMax-H3 backbone +# distilled with data-free DMD2 + VSA-H3 at 80% sparsity, carrying a trained `to_gate_compress` gate +# per block. `model_name` still supplies the pipeline components (text encoder / VAEs / tokenizer / +# schedulers) from the dense base; the transformer comes from `transformer_path` below. +# student_topk is the keep-ratio of key blocks per query (0.20 -> 80% sparsity, the official +# vsa_sparsity=0.8); set it to 1.0 to fall back to dense attention. +student_topk = 0.20 + +# Transformer weights. Accepts either the official FastH3 diffusers transformer folder (`.../transformer`, +# sharded; its gate keys `attn.to_gate_compress.*` are remapped to this repo's `attn.attn_op.to_gate_compress.*` +# on load), or a single-file native-layout state_dict from this repo's own training (loaded AFTER +# set_student_topk so the trained gate overwrites the zero init). Set to None to run the sparse student on +# the dense base weights (zero-init gate). +transformer_path = "models/Diffusion_Transformer/FastVideo-FastH3-8-Step-V2/transformer" +vae_path = None +lora_path = None +pdd_lora_path = None + +# Other params +# MiniMax-H3 generates at a fixed 24 fps, only accepts multiples of 32 as height / width, and snaps video_length up +# to the next 17 * n + 5 the video VAE can decode (the duration has to stay between 5 and 15 seconds). +# Leave sample_size as None to use MiniMax-H3's own 16:9 canvas (768x1344). +sample_size = [704, 1280] +video_length = 124 +fps = 24 + +# 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 red fox trotting through a snowy pine forest, snow crunching underfoot" +seed = 43 +# Number of denoising steps, i.e. of model evaluations. The official FastH3-8-Step-V2 runs eight +# transformer forwards (fastvideo_inference.json: transformer_forwards=8). +num_inference_steps = 8 +# The released checkpoint is guidance-distilled: leave guidance_scale at 1 to run one forward pass per step +# with no CFG. A value above 1 enables classifier-free guidance with a negative_prompt, running two passes. +guidance_scale = 1 +# The exponential sigma shifts of the two schedules. The official 8-step DMD grid +# (dmd_denoising_steps [999,874,...,125]) is *uniform*; here a uniform sigma grid is produced by +# shift=1.0 (set_shift warps the grid, so the checkpoint's video shift 10 / audio 3 do not map onto +# flow_shift directly). flow_shift=1.0 reproduces the trained video schedule; audio keeps shift 3. +flow_shift = 1.0 +audio_flow_shift = 3.0 +lora_weight = 0.55 +save_path = "samples/fasth3-8step-videos-t2v" + +device = set_multi_gpus_devices(ulysses_degree, ring_degree) + +# Transformer: build the VSA sparse student, then REBUILD the VSA kernels at the requested keep-ratio. +# set_student_topk recreates attn_op (resetting to_gate_compress to its zero init), so it must run +# BEFORE any trained gate is loaded. `transformer_path` may be the official FastH3 diffusers transformer +# folder (sharded) or a single-file native-layout state_dict; when it is a folder the backbone comes from +# it directly, otherwise from the dense `model_name` base (whose transformer lives in `transformer/`). +is_transformer_dir = transformer_path is not None and os.path.isdir(transformer_path) +transformer = FastVideoMiniMaxH3Transformer3DModel.from_pretrained( + transformer_path if is_transformer_dir else model_name, + subfolder="" if is_transformer_dir else "transformer", + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, +) +transformer.set_student_topk(student_topk) + +# The official FastH3 checkpoint is sharded and names its gate `attn.to_gate_compress.*`; this repo +# nests it under `attn.attn_op.to_gate_compress.*`. from_pretrained above loaded the (distilled) +# backbone and dropped the gate keys as unexpected, so load just the gate here, remapped, and assert +# every block got one (a silently zero gate would leave the compressed branch untrained -> artifacts). +if is_transformer_dir and student_topk < 1.0: + import json + from safetensors import safe_open + tdir = transformer_path + weight_map = json.load(open(os.path.join(tdir, "diffusion_pytorch_model.safetensors.index.json")))["weight_map"] + shard_to_gate_keys = {} + for key, shard in weight_map.items(): + if "to_gate_compress" in key: + shard_to_gate_keys.setdefault(shard, []).append(key) + gate_sd = {} + for shard, keys in shard_to_gate_keys.items(): + with safe_open(os.path.join(tdir, shard), framework="pt") as f: + for key in keys: + gate_sd[key.replace("attn.to_gate_compress", "attn.attn_op.to_gate_compress")] = f.get_tensor(key) + m, u = transformer.load_state_dict(gate_sd, strict=False) + gate_loaded = sum(1 for k in gate_sd if k not in m) + print(f"FastH3 gate: {gate_loaded}/{len(gate_sd)} loaded, unexpected {len(u)}", flush=True) + assert len(u) == 0 and gate_loaded == len(gate_sd), ( + f"FastH3 gate remap incomplete: loaded {gate_loaded}/{len(gate_sd)}, unexpected {len(u)}.") + +# Single-file native-layout state_dict (gate already under attn_op): overwrite the base backbone wholesale. +if transformer_path is not None and not is_transformer_dir: + 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 isinstance(state_dict, dict) and "state_dict" in state_dict else state_dict + + m, u = transformer.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + attn_op_missing = [k for k in m if "attn_op." in k] + if attn_op_missing and student_topk < 1.0: + # A trained VSA student should provide every attn_op param; missing ones mean the compressed + # branch would run on freshly-init (zero) gate values. Fail loudly instead of silently. + raise RuntimeError( + f"{len(attn_op_missing)} VSA attn_op keys are missing from the checkpoint " + f"(student_topk={student_topk}); e.g. {attn_op_missing[:4]}. Check the checkpoint / topk.") + # `strict=False` accepts a file whose keys belong to another model โ€” a LoRA checkpoint, say โ€” by + # loading nothing at all, so an unexpected key is a hard error. + assert len(u) == 0, ( + f"{transformer_path} holds {len(u)} key(s) the transformer does not have, e.g. {list(u)[:3]}. A LoRA " + "checkpoint belongs in `lora_path`, not `transformer_path`." + ) + +pdd_config = None +if pdd_lora_path is not None: + if lora_path is not None: + raise ValueError("`lora_path` and `pdd_lora_path` cannot be used together.") + from videox_fun.models.minimax_h3_pdd import (load_pdd_lora, + pdd_num_inference_steps, + pdd_step_callback) + pdd_config = load_pdd_lora(transformer, pdd_lora_path) + num_inference_steps = pdd_num_inference_steps(pdd_config, num_inference_steps, teacher_default=40) + +# Video VAE. The released weights are float32 and the decode runs under float16 autocast, so the VAE is not +# downcast even when the rest of the pipeline is bfloat16 (this is also how the training scripts load it). +vae = AutoencoderKLMiniMaxH3.from_pretrained( + model_name, + subfolder="vae", + low_cpu_mem_usage=True, +) + +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)}") + +# Audio VAE, waveform in / waveform out: MiniMax-H3 has no separate vocoder. Float32 as released, like the video VAE. +audio_vae = AutoencoderKLMiniMaxH3Audio.from_pretrained( + model_name, + subfolder="audio_vae", + low_cpu_mem_usage=True, +) + +# Get Tokenizer and Processor +tokenizer = Qwen2TokenizerFast.from_pretrained(os.path.join(model_name, "tokenizer")) +processor = Qwen3VLProcessor.from_pretrained(os.path.join(model_name, "processor")) + +# Get Text encoder. MiniMax-H3 reads the unnormalized hidden state after the 50th decoder layer of Qwen3-VL. +text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + os.path.join(model_name, "text_encoder"), + low_cpu_mem_usage=True, + torch_dtype=weight_dtype, +) +text_encoder = text_encoder.eval() + +# Get Schedulers. MiniMax-H3 steps the video and the audio latents down two schedules inside one transformer call. +scheduler = MiniMaxH3Scheduler.from_pretrained(model_name, subfolder="scheduler") +audio_scheduler = MiniMaxH3Scheduler.from_pretrained(model_name, subfolder="audio_scheduler") + +pipeline = MiniMaxH3Pipeline( + vae=vae, + audio_vae=audio_vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + processor=processor, + transformer=transformer, + scheduler=scheduler, + audio_scheduler=audio_scheduler, +) + +# The float32 modules of the mixed-precision checkpoint stay untouched by the float8 quantization. +fp8_exclude_module_name = [ + "proj_in", "audio_proj_in", "context_embedder", "time_embedder", "time_proj", + "token_refiner", "norm_out", "proj_out", "audio_proj_out", +] +use_qfloat8 = "qfloat8" in GPU_memory_mode +if use_qfloat8: + convert_model_weight_to_float8(transformer, exclude_module_name=fp8_exclude_module_name, device=device) + +# The pure *cpu_offload modes (qfloat8 variants excluded, unsupported under FSDP). Under sequence parallelism the +# sharded modules offload through FSDP's own CPUOffload, enabled here at wrap time; the pipeline then routes the +# whole-model offload call to the non-FSDP VAE / audio VAE automatically, so no per-mode branch is needed below. +wants_offload = GPU_memory_mode in ("model_cpu_offload", "model_group_offload", "sequential_cpu_offload") + +if ulysses_degree > 1 or ring_degree > 1: + from functools import partial + # NOTE: raises for the VSA student (no sequence-parallel path); keep ulysses_degree=ring_degree=1. + transformer.enable_multi_gpus_inference() + if fsdp_dit: + fp32_modules = [m for m in transformer.modules() + if any(p.dtype == torch.float32 for p in m.parameters(recurse=False))] + shard_fn = partial(shard_model, device_id=device, param_dtype=None, cast_dtype=False, + module_to_wrapper=list(transformer.transformer_blocks), + ignored_modules=fp32_modules, offload_to_cpu=wants_offload) + 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, + module_to_wrapper=list(text_encoder.model.language_model.layers), + offload_to_cpu=wants_offload) + pipeline.text_encoder.model = shard_fn(pipeline.text_encoder.model) + print("Add FSDP TEXT ENCODER") + +if compile_dit: + for i in range(len(pipeline.transformer.transformer_blocks)): + pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.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. +# The FP8 conversion above has already run (before the FSDP sharding, on purpose); only the dequant +# wrapper and the memory placement are left, which is what the preconverted tag installs. +apply_gpu_memory_mode(pipeline, GPU_memory_mode, device, weight_dtype, + quant_tag="qfloat8_preconverted" if GPU_memory_mode.endswith("_and_qfloat8") else None, + exclude_module_name=[]) + +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) + +pdd_callback = None if pdd_config is None else pdd_step_callback( + transformer, scheduler, audio_scheduler, pdd_config, num_inference_steps +) + +with torch.no_grad(): + output = pipeline( + prompt=prompt, + height=None if sample_size is None else sample_size[0], + width=None if sample_size is None else sample_size[1], + num_frames=video_length, + num_inference_steps=num_inference_steps, + flow_shift=flow_shift, + audio_flow_shift=audio_flow_shift, + guidance_scale=guidance_scale, + generator=generator, + output_type="pt", + callback_on_step_end=pdd_callback, + ) +print(f"[{os.environ.get('RANK', '0')}] generation done, decoding", flush=True) + +if lora_path is not None: + pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) + +sample = output.videos +audio = output.audio +audio_sample_rate = output.sampling_rate + +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) + video_path = os.path.join(save_path, prefix + ".mp4") + save_videos_with_audio_grid(sample, audio, video_path, fps=fps, audio_sample_rate=audio_sample_rate) + +if ulysses_degree * ring_degree > 1: + import torch.distributed as dist + if dist.get_rank() == 0: + save_results() + # Keep every rank alive until the saving rank finishes; an early exit of one rank makes the elastic launcher + # terminate the others. + dist.barrier() +else: + save_results() diff --git a/examples/fastvideo/predict_t2v_wan2.1.py b/examples/fastvideo/predict_t2v_wan2.1.py new file mode 100644 index 0000000..6a0d872 --- /dev/null +++ b/examples/fastvideo/predict_t2v_wan2.1.py @@ -0,0 +1,313 @@ +import os +import re +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, + FastVideoWanTransformer3DModel, WanT5EncoderModel) +from videox_fun.pipeline import WanPipeline +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_cpu_offload" +# 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. +# NOTE: the VSA sparse kernel has no sequence-parallel path, so a FastWan student runs single-GPU +# only; keep ulysses_degree = ring_degree = 1 (enable_multi_gpus_inference() raises otherwise). +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 + +# Riflex config +enable_riflex = False +# Index of intrinsic frequency +riflex_k = 6 + +# Config and model path +config_path = "config/wan2.1/wan_civitai.yaml" +# Base model (Wan-AI/Wan2.1-T2V-1.3B): provides config.json + the dense backbone + VAE + T5. +model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B" + +# ---- FastVideo VSA sparse-student config -------------------------------------------------- +# Official released checkpoint: FastVideo/FastWan2.1-T2V-1.3B-Diffusers (VSA sparse attention + +# DMD few-step distillation). student_topk is the keep-ratio of key blocks per query and MUST +# equal the value the checkpoint was trained at: topk=0.20 -> 80% sparsity (--VSA-sparsity 0.8). +student_topk = 0.20 + +# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++". +# FastWan ships UniPCMultistepScheduler(flow_shift=3.0); "Flow_Unipc" + shift=3.0 mirrors it and +# keeps every denoising step at a high noise level (DMD-style), unlike "Flow" whose last step ~0. +sampler_name = "Flow_Unipc" +# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics. 3.0 is the official +# flow_shift and the recommended value for 480p. +shift = 3.0 + +# Official released checkpoint (diffusers layout: transformer/diffusion_pytorch_model.safetensors, +# diffusers key naming + the VSA to_gate_compress gate). Download: +# HF_ENDPOINT=https://hf-mirror.com huggingface-cli download \ +# FastVideo/FastWan2.1-T2V-1.3B-Diffusers --local-dir +transformer_path = "models/Diffusion_Transformer/FastWan2.1-T2V-1.3B-Diffusers/transformer/diffusion_pytorch_model.safetensors" +vae_path = None +lora_path = None + +# Other params +sample_size = [480, 848] +video_length = 81 +fps = 16 + +# 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 cinematic aerial shot of a snow-covered mountain village at dawn, warm light spilling from wooden windows, smoke curling from the chimneys while a river of mist drifts through the valley below." +negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" +guidance_scale = 1.0 # the DMD student is CFG-free +seed = 1024 +num_inference_steps = 3 +lora_weight = 0.55 +save_path = "samples/fastvideo-wan2.1-official" + + +def remap_fastwan_diffusers_to_vxf(state_dict): + """Official FastWan diffusers transformer -> VideoX-Fun civitai WanTransformer3DModel layout. + + FastWan ships diffusers naming (attn1/attn2, ffn.net, condition_embedder, proj_out, + scale_shift_table) plus the VSA gate ``to_gate_compress``; VideoX-Fun keeps the native Wan + civitai naming (self_attn/cross_attn, ffn.0/2, time_/text_embedding, head.head, modulation) + and nests the gate under ``self_attn.attn_op``. Only the cross-attn norm is affine, so the + single diffusers ``norm2`` maps to civitai ``norm3``; ``patch_embedding`` is already 5D and + passes through. Turns the 885 released keys into an exact 1:1 match (825 backbone + 30x2 gate). + """ + TOP = { + "condition_embedder.time_embedder.linear_1": "time_embedding.0", + "condition_embedder.time_embedder.linear_2": "time_embedding.2", + "condition_embedder.time_proj": "time_projection.1", + "condition_embedder.text_embedder.linear_1": "text_embedding.0", + "condition_embedder.text_embedder.linear_2": "text_embedding.2", + "proj_out": "head.head", + "scale_shift_table": "head.modulation", + } + out = {} + for k, v in state_dict.items(): + m = re.match(r"blocks\.(\d+)\.(.+)$", k) + if m: + blk, rest = m.group(1), m.group(2) + rest = rest.replace("attn1.", "self_attn.").replace("attn2.", "cross_attn.") + rest = (rest.replace(".to_out.0", ".o").replace(".to_q", ".q") + .replace(".to_k", ".k").replace(".to_v", ".v")) + rest = rest.replace("ffn.net.0.proj", "ffn.0").replace("ffn.net.2", "ffn.2") + if rest.startswith("norm2."): + rest = "norm3." + rest[len("norm2."):] + if rest == "scale_shift_table": + rest = "modulation" + if rest.startswith("to_gate_compress"): + rest = "self_attn.attn_op." + rest + nk = "blocks." + blk + "." + rest + else: + nk = k + for src, dst in TOP.items(): + if k == src or k.startswith(src + "."): + nk = dst + k[len(src):] + break + out[nk] = v + return out + + +device = set_multi_gpus_devices(ulysses_degree, ring_degree) +config = OmegaConf.load(config_path) + +# Build the sparse student from the base Wan2.1 config, then REBUILD the VSA kernels at the +# released keep-ratio BEFORE loading weights: set_student_topk recreates attn_op (resetting +# to_gate_compress to its zero init), so the trained checkpoint values must be loaded after. +transformer = FastVideoWanTransformer3DModel.from_pretrained( + os.path.join(model_name, 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=weight_dtype, +) +transformer.set_student_topk(student_topk) + +if transformer_path is not None: + 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 isinstance(state_dict, dict) and "state_dict" in state_dict else state_dict + + # Official diffusers release -> VideoX-Fun civitai key layout (route the VSA gate under attn_op). + state_dict = remap_fastwan_diffusers_to_vxf(state_dict) + + m, u = transformer.load_state_dict(state_dict, strict=False) + print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") + attn_op_missing = [k for k in m if "attn_op." in k] + if attn_op_missing: + # A VSA student MUST provide every attn_op param; missing ones mean the compressed branch + # would run on freshly-init (zero) gate values. Fail loudly instead of silently. + raise RuntimeError( + f"{len(attn_op_missing)} VSA attn_op keys are missing from the checkpoint " + f"(student_topk={student_topk}); e.g. {attn_op_missing[:4]}. Check the remap / topk.") + if u: + print(f"WARNING: {len(u)} unexpected keys ignored; e.g. {list(u)[:4]}") + # For the official FastWan2.1-T2V-1.3B checkpoint this is expected to be 0 missing / 0 unexpected. + +# 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 = WanPipeline( + 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. +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 + + if enable_riflex: + pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames) + + 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, + ).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) + print(f"Saved video to: {video_path}") + +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/fastvideo/README_TRAIN_DISTILL.md b/scripts/fastvideo/README_TRAIN_DISTILL.md new file mode 100644 index 0000000..4058db5 --- /dev/null +++ b/scripts/fastvideo/README_TRAIN_DISTILL.md @@ -0,0 +1,648 @@ +# FastVideo VSA Distill Training Guide + +This document provides a complete workflow for distilling Wan2.1 into FastWan2.1 with VSA (Video Sparse Attention), including environment configuration, data preparation, distributed training, and inference testing. + +> **Note**: FastVideo VSA combines a block-sparse video attention (a learned gate keeps ~20% of key blocks per query) with DMD few-step distillation, so the student runs CFG-free in ~3 steps while attending sparsely. The training recipe, data pipeline, and distributed backends are identical to the TurboDiffusion distill guide; only the sparse student model class differs. + +--- + +## Table of Contents +- [1. Environment Configuration](#1-environment-configuration) +- [2. Data Preparation](#2-data-preparation) + - [2.1 Quick Test Dataset](#21-quick-test-dataset) + - [2.2 Dataset Structure](#22-dataset-structure) + - [2.3 metadata.json Format](#23-metadatajson-format) + - [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage) +- [3. Distill Training](#3-distill-training) + - [3.1 Download Pretrained Model](#31-download-pretrained-model) + - [3.2 Quick Start (DeepSpeed-Zero-2)](#32-quick-start-deepspeed-zero-2) + - [3.3 Common Training Parameters](#33-common-training-parameters) + - [3.4 Training Validation](#34-training-validation) + - [3.5 Training with FSDP](#35-training-with-fsdp) + - [3.6 Other Backends](#36-other-backends) + - [3.7 Multi-Machine Distributed Training](#37-multi-machine-distributed-training) +- [4. Inference Testing](#4-inference-testing) + - [4.1 Inference Parameters](#41-inference-parameters) + - [4.2 Text-to-Video (T2V) Inference](#42-text-to-video-t2v-inference) + - [4.3 Single-GPU Only](#43-single-gpu-only) +- [5. Additional Resources](#5-additional-resources) + +--- + +## 1. Environment Configuration + +**Method 1: Using requirements.txt** + +```bash +pip install -r requirements.txt +``` + +**Method 2: Manual Dependency 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 +``` + +> ๐Ÿ’ก The VSA block-sparse attention runs on a Triton kernel; Triton ships with PyTorch (>=2.1) on Linux, so no extra install is needed. + +**Method 3: Using Docker** + +When using Docker, please ensure that the GPU driver and CUDA environment are correctly installed on your machine, then execute the following commands: + +``` +# pull image +docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun + +# enter image +docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun +``` + +--- + +## 2. Data Preparation + +### 2.1 Quick Test Dataset + +We provide a test dataset containing several training samples. + +```bash +# Download official demo dataset +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 2.2 Dataset Structure + +``` +๐Ÿ“ฆ datasets/ +โ”œโ”€โ”€ ๐Ÿ“‚ my_dataset/ +โ”‚ โ”œโ”€โ”€ ๐Ÿ“‚ train/ +โ”‚ โ”‚ โ”œโ”€โ”€ ๐Ÿ“„ video001.mp4 +โ”‚ โ”‚ โ”œโ”€โ”€ ๐Ÿ“„ video002.mp4 +โ”‚ โ”‚ โ””โ”€โ”€ ๐Ÿ“„ ... +โ”‚ โ””โ”€โ”€ ๐Ÿ“„ metadata.json +``` + +### 2.3 metadata.json Format + +**Relative Path Format** (example): +```json +[ + { + "file_path": "train/video001.mp4", + "text": "A beautiful sunset over the ocean, golden hour lighting", + "type": "video", + "width": 1024, + "height": 1024 + }, + { + "file_path": "train/video002.mp4", + "text": "A person walking through a forest, cinematic view", + "type": "video", + "width": 1328, + "height": 1328 + } +] +``` + +**Absolute Path Format**: +```json +[ + { + "file_path": "/mnt/data/videos/sunset.mp4", + "text": "A beautiful sunset over the ocean", + "type": "video", + "width": 1024, + "height": 1024 + } +] +``` + +**Key Field Descriptions**: +- `file_path`: Video path (relative or absolute path) +- `text`: Video description (English prompt) +- `type`: Data type, fixed as `"video"` +- `width` / `height`: Video dimensions (**recommended** to provide for bucket training. If not provided, they will be automatically read during training, which may slow down training when data is stored on slow systems like OSS). + - You can use `scripts/process_json_add_width_and_height.py` to extract width and height fields for JSON files without them, supporting both images and videos. + - Usage: `python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Videos-Demo/metadata.json --output_file datasets/X-Fun-Videos-Demo/metadata_add_width_height.json`. + +### 2.4 Relative vs Absolute Path Usage + +**Relative Path**: + +If your data uses relative paths, configure in the training script: + +```bash +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +``` + +**Absolute Path**: + +If your data uses absolute paths, configure in the training script: + +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" +``` + +> ๐Ÿ’ก **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths. + +--- + +## 3. Distill Training + +The distillation is a DMD (distribution matching distillation) run with three networks: a sparse **generator** (the FastWan student, `FastVideoWanTransformer3DModel`), a sparse **fake_score** critic (same class), and a frozen **dense** `real_score` teacher (`WanTransformer3DModel`). The student and critic carry the VSA block-sparse self-attention; the teacher keeps full attention. + +### 3.1 Download Pretrained Model + +```bash +# Create model directory +mkdir -p models/Diffusion_Transformer + +# Download Wan2.1 official weights +modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B +``` + +### 3.2 Quick Start (DeepSpeed-Zero-2) + +After downloading the data as per **2.1 Quick Test Dataset** and the weights as per **3.1 Download Pretrained Model**, you can directly copy and run the quick start command. + +DeepSpeed-Zero-2 and FSDP are recommended for training. Here we use DeepSpeed-Zero-2 as an example. + +The difference between DeepSpeed-Zero-2 and FSDP is whether the model weights are sharded. **If you experience insufficient GPU memory when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP for training. + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +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. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + +accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +### 3.3 Common Training Parameters + +**Key Parameter Descriptions**: + +| Parameter | Description | Example Value | +|-----|------|-------| +| `--config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` | +| `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | +| `--train_data_dir` | Training data directory | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | +| `--train_batch_size` | Batch size per GPU | 1 | +| `--image_sample_size` | Maximum image training resolution | 640 | +| `--video_sample_size` | Maximum video training resolution | 640 | +| `--token_sample_size` | Token sample size | 640 | +| `--video_sample_stride` | Video sampling stride | 2 | +| `--video_sample_n_frames` | Video sampling frames | 81 | +| `--gradient_accumulation_steps` | Gradient accumulation steps (effectively increases batch) | 1 | +| `--dataloader_num_workers` | DataLoader worker processes | 8 | +| `--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 | +| `--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_fastvideo_wan2.1_1.3b` | +| `--gradient_checkpointing` | Enable gradient checkpointing | - | +| `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` | +| `--adam_weight_decay` | AdamW weight decay | 3e-2 | +| `--adam_epsilon` | AdamW epsilon | 1e-10 | +| `--vae_mini_batch` | VAE encoding mini-batch size | 1 | +| `--max_grad_norm` | Gradient clipping threshold | 0.05 | +| `--enable_bucket` | Enable bucket training, no cropping, group by resolution | - | +| `--random_hw_adapt` | Auto-scale images/videos to random size in `[min_size, max_size]` range | - | +| `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - | +| `--uniform_sampling` | Uniform timestep sampling | - | +| `--low_vram` | Low VRAM mode | - | +| `--train_mode` | Training mode: `normal` (standard) or `i2v` (image-to-video) | `normal` | +| `--resume_from_checkpoint` | Resume training path, use `"latest"` to auto-select latest checkpoint | None | +| `--validation_steps` | Run validation every N steps | 2000 | +| `--validation_epochs` | Run validation every N epochs | 5 | +| `--validation_prompts` | Prompts for video generation validation | `"A dog shaking head..."` | +| `--trainable_modules` | Trainable modules (`"."` means all modules) | `"."` | + +**Distill-Specific Parameters**: + +| Parameter | Description | Example Value | +|-----|------|-------| +| `--denoising_step_indices_list` | Denoising step list (core distill parameter) | `1000 750 500 250` | +| `--real_guidance_scale` | Real guidance scale for scoring | 6.0 | +| `--fake_guidance_scale` | Fake guidance scale for scoring | 0.0 | +| `--gen_update_interval` | Generator update interval | 5 | +| `--negative_prompt` | Negative prompt for distillation | Chinese negative prompt | +| `--validation_paths` | Validation image paths for I2V mode | Image path list | +| `--train_sampling_steps` | Training sampling steps | 1000 | + +**VSA Sparsity**: + +- The student/critic self-attention is VSA block-sparse. The keep-ratio `student_topk` defaults to `0.20` (i.e. 80% sparsity) inside `FastVideoWanTransformer3DModel`; it is a model-class default, not a CLI flag. +- The learned compression gate `to_gate_compress` is created zero-initialized and becomes trainable automatically under `--trainable_modules "."` (its parameter names contain a dot, like every other module). +- The `real_score` teacher stays dense (`WanTransformer3DModel`) and frozen, so the sparsity is distilled into the student rather than copied from the teacher. +- Keep the inference-time `student_topk` equal to the value the checkpoint was trained at; changing it rebuilds the attention op and resets the learned gate. + +**Sample Size Configuration Guide**: +- `video_sample_size` represents the resolution size of videos; when `random_hw_adapt` is True, it represents the minimum value between video and image resolutions. +- `image_sample_size` represents the resolution size of images; when `random_hw_adapt` is True, it represents the maximum value between video and image resolutions. +- `token_sample_size` represents the resolution corresponding to the maximum token length when `training_with_video_token_length` is True. +- Due to potential confusion in configuration, **if you don't require arbitrary resolution for finetuning**, it is recommended to set `video_sample_size`, `image_sample_size`, and `token_sample_size` to the same fixed value, such as **(320, 480, 512, 640, 960)**. + - **All set to 320** represents **240P**. + - **All set to 480** represents **320P**. + - **All set to 640** represents **480P**. + - **All set to 960** represents **720P**. + +**Other Parameter Explanations**: +- `enable_bucket` is used to enable bucket training. When enabled, the model does not crop the images and videos at the center, but instead, it trains the entire images and videos after grouping them into buckets based on resolution. +- `random_hw_adapt` is used to enable automatic height and width scaling for images and videos. When `random_hw_adapt` is enabled, the training images will have their height and width set to `image_sample_size` as the maximum and `min(video_sample_size, 512)` as the minimum. +- `training_with_video_token_length` specifies training the model according to token length. For training images and videos, the height and width will be set to `image_sample_size` as the maximum and `video_sample_size` as the minimum. +- `train_mode` is used to specify the training mode, which can be either normal or i2v. If you only wish to achieve text-to-video generation, keep `normal`. +- `resume_from_checkpoint` is used to set the training should be resumed from a previous checkpoint. Use a path or `"latest"` to automatically select the last available checkpoint. + +### 3.4 Training Validation + +You can configure validation parameters to periodically generate test videos during training, allowing you to monitor training progress and model quality. + +**Validation Parameter Descriptions**: + +| Parameter | Description | Recommended Value | +|------|------|--------| +| `--validation_steps` | Run validation every N steps | 2000 | +| `--validation_epochs` | Run validation every N epochs | 5 | +| `--validation_prompts` | Prompts for video generation validation | English prompts | +| `--validation_paths` | Validation image paths for I2V mode (i2v mode only) | `"asset/1.png"` | + +**Normal Mode Example** (T2V Validation): + +```bash + --validation_steps=2000 \ + --validation_epochs=5 \ + --validation_prompts="A dog shaking head. The video is of high quality, and the view is very clear." +``` + +**Notes**: +- Validation videos will be saved to the `output_dir` directory. +- Multiple prompts validation format: `--validation_prompts "prompt1" "prompt2" "prompt3"`. +- Distill model validation uses the few steps defined in `denoising_step_indices_list` and runs CFG-free (`guidance_scale=1.0`), matching inference. + +### 3.5 Training with FSDP + +**If you experience insufficient GPU memory when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP for training. `scripts/fastvideo/train_distill_wan2.1.sh` is a ready-to-run FSDP launcher (it wraps `WanAttentionBlock` with `FULL_SHARD`). + +```sh +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +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. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + +accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +### 3.6 Other Backends + +#### 3.6.1 Training with DeepSpeed-Zero-3 + +DeepSpeed Zero-3 is not highly recommended at the moment. In this repository, using FSDP has fewer errors and is more stable. + +DeepSpeed Zero-3 is suitable for 14B Wan at high resolutions. After training, you can use the following command to get the final model: +```bash +python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization +``` + +Training shell command is as follows: +```sh +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +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/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +#### 3.6.2 Training Without DeepSpeed and FSDP + +**This approach is not recommended because without memory-saving backends, it easily causes insufficient GPU memory**. This is only provided as a reference shell for training. + +```sh +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +NCCL_DEBUG=INFO + +accelerate launch --mixed_precision="bf16" scripts/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +### 3.7 Multi-Machine Distributed Training + +**Suitable for**: Ultra-large datasets, faster training speed + +#### 3.7.1 Environment Configuration + +Assuming 2 machines, each with 8 GPUs: + +**Machine 0 (Master)**: +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +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 +export MASTER_PORT=10086 +export WORLD_SIZE=2 # Total number of machines +export NUM_PROCESS=16 # Total processes = machines ร— 8 +export RANK=0 # Current machine rank (0 or 1) +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/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +**Machine 1 (Worker)**: +```bash +export MASTER_ADDR="192.168.1.100" # Same as Master +export MASTER_PORT=10086 +export WORLD_SIZE=2 +export NUM_PROCESS=16 +export RANK=1 # Note this is 1 +NCCL_DEBUG=INFO + +# Use the same accelerate launch command as Machine 0 +``` + +#### 3.7.2 Multi-Machine Training Notes + +- **Network Requirements**: + - RDMA/InfiniBand recommended (high performance) + - Without RDMA, add environment variables: + ```bash + export NCCL_IB_DISABLE=1 + export NCCL_P2P_DISABLE=1 + ``` +- **Data Synchronization**: All machines must be able to access the same data paths (NFS/shared storage) + +--- + +## 4. Inference Testing + +### 4.1 Inference Parameters + +**Key Parameter Descriptions** (from `examples/fastvideo/predict_t2v_wan2.1.py`): + +| Parameter | Description | Example Value | +|------|------|-------| +| `GPU_memory_mode` | GPU memory mode, see table below for options | `model_cpu_offload` | +| `ulysses_degree` | Ulysses parallelism degree (VSA is single-GPU only, keep 1) | 1 | +| `ring_degree` | Ring parallelism degree (VSA is single-GPU only, keep 1) | 1 | +| `config_path` | Config file path | `config/wan2.1/wan_civitai.yaml` | +| `model_name` | Base model path (config + dense backbone + VAE + T5) | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` | +| `student_topk` | VSA keep-ratio; must equal the trained value | `0.20` | +| `sampler_name` | Sampler type: `Flow`, `Flow_Unipc`, `Flow_DPM++` | `Flow_Unipc` | +| `shift` | Noise schedule shift (official FastWan flow_shift) | `3.0` | +| `transformer_path` | Path to the FastWan student Transformer weights | trained checkpoint / official FastWan | +| `vae_path` | Path to trained VAE weights | `None` | +| `lora_path` | LoRA weights path | `None` | +| `sample_size` | Generated video resolution `[height, width]` | `[480, 848]` | +| `video_length` | Number of generated frames | `81` | +| `fps` | Frames per second | `16` | +| `weight_dtype` | Model weight precision, use `torch.float16` for GPUs without bf16 support | `torch.bfloat16` | +| `prompt` | Positive prompt describing generated content | `"A cinematic aerial shot..."` | +| `negative_prompt` | Negative prompt to avoid certain content | Chinese negative prompt | +| `guidance_scale` | Guidance strength (the DMD student is CFG-free) | 1.0 | +| `seed` | Random seed for reproducibility | 1024 | +| `num_inference_steps` | Number of inference steps (FastWan uses 3) | 3 | +| `save_path` | Path to save generated videos | `samples/fastvideo-wan2.1-official` | + +**GPU Memory Mode Descriptions**: + +| Mode | Description | Memory Usage | +|------|------|---------| +| `model_full_load` | Full model loaded to GPU | Highest | +| `model_full_load_and_qfloat8` | Full load + FP8 quantization | High | +| `model_cpu_offload` | Offload model to CPU after use | Medium | +| `model_cpu_offload_and_qfloat8` | CPU offload + FP8 quantization | Medium-Low | +| `model_group_offload` | Layer groups switch between CPU/CUDA | Low | +| `sequential_cpu_offload` | Layer-by-layer offload (slowest) | Lowest | + +### 4.2 Text-to-Video (T2V) Inference + +Run single-GPU inference: + +```bash +python examples/fastvideo/predict_t2v_wan2.1.py +``` + +Edit `examples/fastvideo/predict_t2v_wan2.1.py` according to your needs. For first-time inference, focus on the following key parameters. For other parameters, refer to the inference parameter descriptions above. + +```python +# Choose based on GPU memory +GPU_memory_mode = "model_cpu_offload" +# Base model path (config + dense backbone + VAE + T5) +model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +# VSA keep-ratio; must equal the value the checkpoint was trained at +student_topk = 0.20 +# Path to the FastWan student weights (your trained checkpoint, or the official FastWan release) +transformer_path = "models/Personalized_Model/FastWan2.1-T2V-1.3B.pth" +# FastWan sampler + shift (CFG-free few-step student) +sampler_name = "Flow_Unipc" +shift = 3.0 +guidance_scale = 1.0 +num_inference_steps = 3 +# Write based on your generation content +prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn..." +# ... +``` + +> ๐Ÿ’ก The official `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` checkpoint ships in diffusers key naming; `examples/fastvideo/predict_t2v_wan2.1.py` remaps it to the VideoX-Fun civitai layout (including the VSA `to_gate_compress` gate) before loading. A checkpoint produced by this training guide is already in the native layout. + +### 4.3 Single-GPU Only + +The VSA block-sparse kernel has no sequence-parallel path, so a FastWan student runs on a **single GPU**. Keep `ulysses_degree = ring_degree = 1`; calling `enable_multi_gpus_inference()` raises otherwise. To fit larger resolutions on one card, lower the memory footprint via `GPU_memory_mode` (e.g. `model_cpu_offload` or `sequential_cpu_offload`) rather than splitting across GPUs. + +--- + +## 5. Additional Resources + +- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun +- **FastVideo GitHub**: https://github.com/hao-ai-lab/FastVideo +- **VSA Paper (Video Sparse Attention)**: https://arxiv.org/abs/2505.13389 diff --git a/scripts/fastvideo/README_TRAIN_DISTILL_zh-CN.md b/scripts/fastvideo/README_TRAIN_DISTILL_zh-CN.md new file mode 100644 index 0000000..b895666 --- /dev/null +++ b/scripts/fastvideo/README_TRAIN_DISTILL_zh-CN.md @@ -0,0 +1,648 @@ +# FastVideo VSA ่’ธ้ฆ่ฎญ็ปƒๆŒ‡ๅ— + +ๆœฌๆ–‡ๆกฃๆไพ›ไบ†ๅฐ† Wan2.1 ่’ธ้ฆไธบ FastWan2.1๏ผˆ้‡‡็”จ VSA๏ผŒVideo Sparse Attention๏ผ‰็š„ๅฎŒๆ•ดๅทฅไฝœๆต๏ผŒๅŒ…ๆ‹ฌ็Žฏๅขƒ้…็ฝฎใ€ๆ•ฐๆฎๅ‡†ๅค‡ใ€ๅˆ†ๅธƒๅผ่ฎญ็ปƒๅ’ŒๆŽจ็†ๆต‹่ฏ•ใ€‚ + +> **ๆณจๆ„**: FastVideo VSA ๅฐ†ๅ—็จ€็–่ง†้ข‘ๆณจๆ„ๅŠ›๏ผˆ็”ฑๅฏๅญฆไน ้—จๆŽงไธบๆฏไธช query ไฟ็•™็บฆ 20% ็š„ key ๅ—๏ผ‰ไธŽ DMD ๅฐ‘ๆญฅ่’ธ้ฆ็ป“ๅˆ๏ผŒไฝฟๅญฆ็”Ÿๆจกๅž‹ๅœจ็บฆ 3 ๆญฅใ€ๆ—  CFG ็š„ๆกไปถไธ‹ไปฅ็จ€็–ๆณจๆ„ๅŠ›ๆŽจ็†ใ€‚่ฎญ็ปƒ้…ๆ–นใ€ๆ•ฐๆฎ็ฎก็บฟไธŽๅˆ†ๅธƒๅผๅŽ็ซฏไธŽ TurboDiffusion ่’ธ้ฆๆŒ‡ๅ—ๅฎŒๅ…จไธ€่‡ด๏ผŒไป…็จ€็–ๅญฆ็”Ÿๆจกๅž‹็ฑปไธๅŒใ€‚ + +--- + +## ็›ฎๅฝ• +- [ไธ€ใ€็Žฏๅขƒ้…็ฝฎ](#ไธ€็Žฏๅขƒ้…็ฝฎ) +- [ไบŒใ€ๆ•ฐๆฎๅ‡†ๅค‡](#ไบŒๆ•ฐๆฎๅ‡†ๅค‡) + - [2.1 ๅฟซ้€Ÿๆต‹่ฏ•ๆ•ฐๆฎ้›†](#21-ๅฟซ้€Ÿๆต‹่ฏ•ๆ•ฐๆฎ้›†) + - [2.2 ๆ•ฐๆฎ้›†็ป“ๆž„](#22-ๆ•ฐๆฎ้›†็ป“ๆž„) + - [2.3 metadata.json ๆ ผๅผ](#23-metadatajson-ๆ ผๅผ) + - [2.4 ็›ธๅฏน่ทฏๅพ„ไธŽ็ปๅฏน่ทฏๅพ„ไฝฟ็”จๆ–นๆกˆ](#24-็›ธๅฏน่ทฏๅพ„ไธŽ็ปๅฏน่ทฏๅพ„ไฝฟ็”จๆ–นๆกˆ) +- [ไธ‰ใ€่’ธ้ฆ่ฎญ็ปƒ](#ไธ‰่’ธ้ฆ่ฎญ็ปƒ) + - [3.1 ไธ‹่ฝฝ้ข„่ฎญ็ปƒๆจกๅž‹](#31-ไธ‹่ฝฝ้ข„่ฎญ็ปƒๆจกๅž‹) + - [3.2 ๅฟซ้€Ÿๅผ€ๅง‹๏ผˆDeepSpeed-Zero-2๏ผ‰](#32-ๅฟซ้€Ÿๅผ€ๅง‹deepspeed-zero-2) + - [3.3 ๅธธ่ง่ฎญ็ปƒๅ‚ๆ•ฐ](#33-ๅธธ่ง่ฎญ็ปƒๅ‚ๆ•ฐ) + - [3.4 ่ฎญ็ปƒ้ชŒ่ฏ](#34-่ฎญ็ปƒ้ชŒ่ฏ) + - [3.5 ไฝฟ็”จ FSDP ่ฎญ็ปƒ](#35-ไฝฟ็”จ-fsdp-่ฎญ็ปƒ) + - [3.6 ๅ…ถไป–ๅŽ็ซฏ](#36-ๅ…ถไป–ๅŽ็ซฏ) + - [3.7 ๅคšๆœบๅˆ†ๅธƒๅผ่ฎญ็ปƒ](#37-ๅคšๆœบๅˆ†ๅธƒๅผ่ฎญ็ปƒ) +- [ๅ››ใ€ๆŽจ็†ๆต‹่ฏ•](#ๅ››ๆŽจ็†ๆต‹่ฏ•) + - [4.1 ๆŽจ็†ๅ‚ๆ•ฐ่งฃๆž](#41-ๆŽจ็†ๅ‚ๆ•ฐ่งฃๆž) + - [4.2 ๆ–‡ๆœฌ็”Ÿๆˆ่ง†้ข‘๏ผˆT2V๏ผ‰ๆŽจ็†](#42-ๆ–‡ๆœฌ็”Ÿๆˆ่ง†้ข‘t2vๆŽจ็†) + - [4.3 ไป…ๆ”ฏๆŒๅ•ๅก](#43-ไป…ๆ”ฏๆŒๅ•ๅก) +- [ไบ”ใ€ๆ›ดๅคš่ต„ๆบ](#ไบ”ๆ›ดๅคš่ต„ๆบ) + +--- + +## ไธ€ใ€็Žฏๅขƒ้…็ฝฎ + +**ๆ–นๅผไธ€๏ผšไฝฟ็”จ requirements.txt** + +```bash +pip install -r requirements.txt +``` + +**ๆ–นๅผไบŒ๏ผšๆ‰‹ๅŠจๅฎ‰่ฃ…ไพ่ต–** + +```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 +``` + +> ๐Ÿ’ก VSA ๅ—็จ€็–ๆณจๆ„ๅŠ›่ฟ่กŒๅœจ Triton ๅ†…ๆ ธไธŠ๏ผ›Triton ้š PyTorch๏ผˆ>=2.1๏ผ‰ๅœจ Linux ไธŠไธ€ๅนถๆไพ›๏ผŒๆ— ้œ€้ขๅค–ๅฎ‰่ฃ…ใ€‚ + +**ๆ–นๅผไธ‰๏ผšไฝฟ็”จ Docker** + +ไฝฟ็”จ Docker ๆ—ถ๏ผŒ่ฏท็กฎไฟๆœบๅ™จไธŠๅทฒๆญฃ็กฎๅฎ‰่ฃ… GPU ้ฉฑๅŠจไธŽ CUDA ็Žฏๅขƒ๏ผŒ็„ถๅŽๆ‰ง่กŒไปฅไธ‹ๅ‘ฝไปค๏ผš + +``` +# ๆ‹‰ๅ–้•œๅƒ +docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun + +# ่ฟ›ๅ…ฅๅฎนๅ™จ +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 +``` + +--- + +## ไบŒใ€ๆ•ฐๆฎๅ‡†ๅค‡ + +### 2.1 ๅฟซ้€Ÿๆต‹่ฏ•ๆ•ฐๆฎ้›† + +ๆˆ‘ไปฌๆไพ›ไบ†ไธ€ไธชๅŒ…ๅซ่‹ฅๅนฒ่ฎญ็ปƒๆ ทๆœฌ็š„ๆต‹่ฏ•ๆ•ฐๆฎ้›†ใ€‚ + +```bash +# ไธ‹่ฝฝๅฎ˜ๆ–นๆผ”็คบๆ•ฐๆฎ้›† +modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo +``` + +### 2.2 ๆ•ฐๆฎ้›†็ป“ๆž„ + +``` +๐Ÿ“ฆ datasets/ +โ”œโ”€โ”€ ๐Ÿ“‚ my_dataset/ +โ”‚ โ”œโ”€โ”€ ๐Ÿ“‚ train/ +โ”‚ โ”‚ โ”œโ”€โ”€ ๐Ÿ“„ video001.mp4 +โ”‚ โ”‚ โ”œโ”€โ”€ ๐Ÿ“„ video002.mp4 +โ”‚ โ”‚ โ””โ”€โ”€ ๐Ÿ“„ ... +โ”‚ โ””โ”€โ”€ ๐Ÿ“„ metadata.json +``` + +### 2.3 metadata.json ๆ ผๅผ + +**็›ธๅฏน่ทฏๅพ„ๆ ผๅผ**๏ผˆ็คบไพ‹๏ผ‰๏ผš +```json +[ + { + "file_path": "train/video001.mp4", + "text": "A beautiful sunset over the ocean, golden hour lighting", + "type": "video", + "width": 1024, + "height": 1024 + }, + { + "file_path": "train/video002.mp4", + "text": "A person walking through a forest, cinematic view", + "type": "video", + "width": 1328, + "height": 1328 + } +] +``` + +**็ปๅฏน่ทฏๅพ„ๆ ผๅผ**๏ผš +```json +[ + { + "file_path": "/mnt/data/videos/sunset.mp4", + "text": "A beautiful sunset over the ocean", + "type": "video", + "width": 1024, + "height": 1024 + } +] +``` + +**ๅ…ณ้”ฎๅญ—ๆฎต่ฏดๆ˜Ž**๏ผš +- `file_path`๏ผš่ง†้ข‘่ทฏๅพ„๏ผˆ็›ธๅฏนๆˆ–็ปๅฏน่ทฏๅพ„๏ผ‰ +- `text`๏ผš่ง†้ข‘ๆ่ฟฐ๏ผˆ่‹ฑๆ–‡ prompt๏ผ‰ +- `type`๏ผšๆ•ฐๆฎ็ฑปๅž‹๏ผŒๅ›บๅฎšไธบ `"video"` +- `width` / `height`๏ผš่ง†้ข‘ๅฐบๅฏธ๏ผˆ**ๅปบ่ฎฎ**ๆไพ›ไปฅไพฟๅˆ†ๆกถ่ฎญ็ปƒใ€‚่‹ฅไธๆไพ›๏ผŒ่ฎญ็ปƒๆ—ถไผš่‡ชๅŠจ่ฏปๅ–๏ผŒๅฝ“ๆ•ฐๆฎๅญ˜ๆ”พๅœจ OSS ็ญ‰่พƒๆ…ข็ณป็ปŸไธŠๆ—ถไผšๆ‹–ๆ…ข่ฎญ็ปƒ๏ผ‰ใ€‚ + - ๅฏไฝฟ็”จ `scripts/process_json_add_width_and_height.py` ไธบ็ผบๅฐ‘ๅฎฝ้ซ˜ๅญ—ๆฎต็š„ JSON ๆ–‡ไปถๆๅ– width ๅ’Œ height๏ผŒๅŒๆ—ถๆ”ฏๆŒๅ›พ็‰‡ไธŽ่ง†้ข‘ใ€‚ + - ็”จๆณ•๏ผš`python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Videos-Demo/metadata.json --output_file datasets/X-Fun-Videos-Demo/metadata_add_width_height.json`ใ€‚ + +### 2.4 ็›ธๅฏน่ทฏๅพ„ไธŽ็ปๅฏน่ทฏๅพ„ไฝฟ็”จๆ–นๆกˆ + +**็›ธๅฏน่ทฏๅพ„**๏ผš + +่‹ฅๆ•ฐๆฎไฝฟ็”จ็›ธๅฏน่ทฏๅพ„๏ผŒๅœจ่ฎญ็ปƒ่„šๆœฌไธญ้…็ฝฎ๏ผš + +```bash +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +``` + +**็ปๅฏน่ทฏๅพ„**๏ผš + +่‹ฅๆ•ฐๆฎไฝฟ็”จ็ปๅฏน่ทฏๅพ„๏ผŒๅœจ่ฎญ็ปƒ่„šๆœฌไธญ้…็ฝฎ๏ผš + +```bash +export DATASET_NAME="" +export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json" +``` + +> ๐Ÿ’ก **ๅปบ่ฎฎ**๏ผšๆ•ฐๆฎ้›†่พƒๅฐไธ”ๅญ˜ๆ”พๅœจๆœฌๅœฐๆ—ถไฝฟ็”จ็›ธๅฏน่ทฏๅพ„๏ผ›ๆ•ฐๆฎ้›†ๅญ˜ๆ”พๅœจๅค–้ƒจๅญ˜ๅ‚จ๏ผˆๅฆ‚ NASใ€OSS๏ผ‰ๆˆ–่ทจๅคšๅฐๆœบๅ™จๅ…ฑไบซๆ—ถไฝฟ็”จ็ปๅฏน่ทฏๅพ„ใ€‚ + +--- + +## ไธ‰ใ€่’ธ้ฆ่ฎญ็ปƒ + +่’ธ้ฆๆ˜ฏไธ€ๆฌก DMD๏ผˆdistribution matching distillation๏ผ‰่ฎญ็ปƒ๏ผŒๅŒ…ๅซไธ‰ไธช็ฝ‘็ปœ๏ผš็จ€็–็š„ **generator**๏ผˆFastWan ๅญฆ็”Ÿ๏ผŒ`FastVideoWanTransformer3DModel`๏ผ‰ใ€็จ€็–็š„ **fake_score** critic๏ผˆๅŒไธ€ไธช็ฑป๏ผ‰๏ผŒไปฅๅŠๅ†ป็ป“็š„ **dense** `real_score` ๆ•™ๅธˆ๏ผˆ`WanTransformer3DModel`๏ผ‰ใ€‚ๅญฆ็”ŸไธŽ critic ๆบๅธฆ VSA ๅ—็จ€็–่‡ชๆณจๆ„ๅŠ›๏ผŒๆ•™ๅธˆไฟๆŒๅ…จๆณจๆ„ๅŠ›ใ€‚ + +### 3.1 ไธ‹่ฝฝ้ข„่ฎญ็ปƒๆจกๅž‹ + +```bash +# ๅˆ›ๅปบๆจกๅž‹็›ฎๅฝ• +mkdir -p models/Diffusion_Transformer + +# ไธ‹่ฝฝ Wan2.1 ๅฎ˜ๆ–นๆƒ้‡ +modelscope download --model Wan-AI/Wan2.1-T2V-1.3B --local_dir models/Diffusion_Transformer/Wan2.1-T2V-1.3B +``` + +### 3.2 ๅฟซ้€Ÿๅผ€ๅง‹๏ผˆDeepSpeed-Zero-2๏ผ‰ + +ๆŒ‰็…ง **2.1 ๅฟซ้€Ÿๆต‹่ฏ•ๆ•ฐๆฎ้›†** ไธ‹่ฝฝๆ•ฐๆฎใ€ๆŒ‰็…ง **3.1 ไธ‹่ฝฝ้ข„่ฎญ็ปƒๆจกๅž‹** ไธ‹่ฝฝๆƒ้‡ๅŽ๏ผŒๅณๅฏ็›ดๆŽฅๅคๅˆถ่ฟ่กŒๅฟซ้€Ÿๅผ€ๅง‹ๅ‘ฝไปคใ€‚ + +ๆŽจ่ไฝฟ็”จ DeepSpeed-Zero-2 ๆˆ– FSDP ่ฎญ็ปƒ๏ผŒ่ฟ™้‡Œไปฅ DeepSpeed-Zero-2 ไธบไพ‹ใ€‚ + +DeepSpeed-Zero-2 ไธŽ FSDP ็š„ๅŒบๅˆซๅœจไบŽๆจกๅž‹ๆƒ้‡ๆ˜ฏๅฆๅˆ†็‰‡ใ€‚**่‹ฅไฝฟ็”จ DeepSpeed-Zero-2 ๅคšๅก่ฎญ็ปƒๆ—ถๆ˜พๅญ˜ไธ่ถณ**๏ผŒๅฏๅˆ‡ๆขๅˆฐ FSDP ่ฎญ็ปƒใ€‚ + +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +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. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + +accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +### 3.3 ๅธธ่ง่ฎญ็ปƒๅ‚ๆ•ฐ + +**ๅ…ณ้”ฎๅ‚ๆ•ฐ่ฏดๆ˜Ž**๏ผš + +| ๅ‚ๆ•ฐ | ๆ่ฟฐ | ็คบไพ‹ๅ€ผ | +|-----|------|-------| +| `--config_path` | ้…็ฝฎๆ–‡ไปถ่ทฏๅพ„ | `config/wan2.1/wan_civitai.yaml` | +| `--pretrained_model_name_or_path` | ้ข„่ฎญ็ปƒๆจกๅž‹่ทฏๅพ„ | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B/` | +| `--train_data_dir` | ่ฎญ็ปƒๆ•ฐๆฎ็›ฎๅฝ• | `datasets/X-Fun-Videos-Demo/` | +| `--train_data_meta` | ่ฎญ็ปƒๆ•ฐๆฎๅ…ƒไฟกๆฏๆ–‡ไปถ | `datasets/X-Fun-Videos-Demo/metadata_add_width_height.json` | +| `--train_batch_size` | ๆฏๅผ ๅก็š„ batch size | 1 | +| `--image_sample_size` | ๅ›พ็‰‡่ฎญ็ปƒๆœ€ๅคงๅˆ†่พจ็Ž‡ | 640 | +| `--video_sample_size` | ่ง†้ข‘่ฎญ็ปƒๆœ€ๅคงๅˆ†่พจ็Ž‡ | 640 | +| `--token_sample_size` | Token ้‡‡ๆ ทๅฐบๅฏธ | 640 | +| `--video_sample_stride` | ่ง†้ข‘้‡‡ๆ ทๆญฅ้•ฟ | 2 | +| `--video_sample_n_frames` | ่ง†้ข‘้‡‡ๆ ทๅธงๆ•ฐ | 81 | +| `--gradient_accumulation_steps` | ๆขฏๅบฆ็ดฏ็งฏๆญฅๆ•ฐ๏ผˆ็ญ‰ๆ•ˆๅขžๅคง batch๏ผ‰ | 1 | +| `--dataloader_num_workers` | DataLoader ๅทฅไฝœ่ฟ›็จ‹ๆ•ฐ | 8 | +| `--num_train_epochs` | ่ฎญ็ปƒ่ฝฎๆ•ฐ | 100 | +| `--checkpointing_steps` | ๆฏ N ๆญฅไฟๅญ˜ checkpoint | 50 | +| `--learning_rate` | ๅˆๅง‹ๅญฆไน ็Ž‡๏ผˆgenerator๏ผ‰ | 2e-06 | +| `--learning_rate_critic` | ๅˆๅง‹ๅญฆไน ็Ž‡๏ผˆcritic๏ผ‰ | 2e-06 | +| `--lr_scheduler` | ๅญฆไน ็Ž‡่ฐƒๅบฆๅ™จ | `constant_with_warmup` | +| `--lr_warmup_steps` | ๅญฆไน ็Ž‡ warmup ๆญฅๆ•ฐ | 100 | +| `--seed` | ้šๆœบ็งๅญ | 42 | +| `--output_dir` | ่พ“ๅ‡บ็›ฎๅฝ• | `output_dir_fastvideo_wan2.1_1.3b` | +| `--gradient_checkpointing` | ๅฏ็”จๆขฏๅบฆๆฃ€ๆŸฅ็‚น | - | +| `--mixed_precision` | ๆททๅˆ็ฒพๅบฆ๏ผš`fp16/bf16` | `bf16` | +| `--adam_weight_decay` | AdamW ๆƒ้‡่กฐๅ‡ | 3e-2 | +| `--adam_epsilon` | AdamW epsilon | 1e-10 | +| `--vae_mini_batch` | VAE ็ผ–็  mini-batch ๅคงๅฐ | 1 | +| `--max_grad_norm` | ๆขฏๅบฆ่ฃๅ‰ช้˜ˆๅ€ผ | 0.05 | +| `--enable_bucket` | ๅฏ็”จๅˆ†ๆกถ่ฎญ็ปƒ๏ผŒไธ่ฃๅ‰ช๏ผŒๆŒ‰ๅˆ†่พจ็Ž‡ๅˆ†็ป„ | - | +| `--random_hw_adapt` | ๅฐ†ๅ›พ็‰‡/่ง†้ข‘่‡ชๅŠจ็ผฉๆ”พๅˆฐ `[min_size, max_size]` ่Œƒๅ›ดๅ†…็š„้šๆœบๅฐบๅฏธ | - | +| `--training_with_video_token_length` | ๅŸบไบŽ token ้•ฟๅบฆ่ฎญ็ปƒ๏ผŒๆ”ฏๆŒไปปๆ„ๅˆ†่พจ็Ž‡ | - | +| `--uniform_sampling` | ๅ‡ๅŒ€ๆ—ถ้—ดๆญฅ้‡‡ๆ ท | - | +| `--low_vram` | ไฝŽๆ˜พๅญ˜ๆจกๅผ | - | +| `--train_mode` | ่ฎญ็ปƒๆจกๅผ๏ผš`normal`๏ผˆๆ ‡ๅ‡†๏ผ‰ๆˆ– `i2v`๏ผˆๅ›พ็”Ÿ่ง†้ข‘๏ผ‰ | `normal` | +| `--resume_from_checkpoint` | ็ปญ่ฎญ่ทฏๅพ„๏ผŒไฝฟ็”จ `"latest"` ่‡ชๅŠจ้€‰ๆ‹ฉๆœ€ๆ–ฐ checkpoint | None | +| `--validation_steps` | ๆฏ N ๆญฅ่ฟ่กŒไธ€ๆฌก้ชŒ่ฏ | 2000 | +| `--validation_epochs` | ๆฏ N ่ฝฎ่ฟ่กŒไธ€ๆฌก้ชŒ่ฏ | 5 | +| `--validation_prompts` | ็”จไบŽ่ง†้ข‘็”Ÿๆˆ้ชŒ่ฏ็š„ prompt | `"A dog shaking head..."` | +| `--trainable_modules` | ๅฏ่ฎญ็ปƒๆจกๅ—๏ผˆ`"."` ่กจ็คบๅ…จ้ƒจๆจกๅ—๏ผ‰ | `"."` | + +**่’ธ้ฆไธ“็”จๅ‚ๆ•ฐ**๏ผš + +| ๅ‚ๆ•ฐ | ๆ่ฟฐ | ็คบไพ‹ๅ€ผ | +|-----|------|-------| +| `--denoising_step_indices_list` | ๅŽปๅ™ชๆญฅๅˆ—่กจ๏ผˆ่’ธ้ฆๆ ธๅฟƒๅ‚ๆ•ฐ๏ผ‰ | `1000 750 500 250` | +| `--real_guidance_scale` | ๆ‰“ๅˆ†็”จ็š„ real guidance scale | 6.0 | +| `--fake_guidance_scale` | ๆ‰“ๅˆ†็”จ็š„ fake guidance scale | 0.0 | +| `--gen_update_interval` | generator ๆ›ดๆ–ฐ้—ด้š” | 5 | +| `--negative_prompt` | ่’ธ้ฆ็”จ็š„่ดŸๅ‘ prompt | ไธญๆ–‡่ดŸๅ‘ prompt | +| `--validation_paths` | I2V ๆจกๅผ็š„้ชŒ่ฏๅ›พ็‰‡่ทฏๅพ„ | ๅ›พ็‰‡่ทฏๅพ„ๅˆ—่กจ | +| `--train_sampling_steps` | ่ฎญ็ปƒ้‡‡ๆ ทๆญฅๆ•ฐ | 1000 | + +**VSA ็จ€็–ๅบฆ**๏ผš + +- ๅญฆ็”Ÿ/critic ็š„่‡ชๆณจๆ„ๅŠ›ไธบ VSA ๅ—็จ€็–ใ€‚ไฟ็•™ๆฏ” `student_topk` ๅœจ `FastVideoWanTransformer3DModel` ๅ†…้ป˜่ฎคไธบ `0.20`๏ผˆๅณ 80% ็จ€็–ๅบฆ๏ผ‰๏ผ›ๅฎƒๆ˜ฏๆจกๅž‹็ฑป้ป˜่ฎคๅ€ผ๏ผŒ่€Œ้žๅ‘ฝไปค่กŒๆ——ๆ ‡ใ€‚ +- ๅฏๅญฆไน ็š„ๅŽ‹็ผฉ้—จๆŽง `to_gate_compress` ไปฅ้›ถๅˆๅง‹ๅŒ–ๅˆ›ๅปบ๏ผŒๅนถๅœจ `--trainable_modules "."` ไธ‹่‡ชๅŠจๅ˜ไธบๅฏ่ฎญ็ปƒ๏ผˆๅ…ถๅ‚ๆ•ฐๅไธŽๅ…ถๅฎƒๆจกๅ—ไธ€ๆ ทๅŒ…ๅซ็‚นๅท๏ผ‰ใ€‚ +- `real_score` ๆ•™ๅธˆไฟๆŒ dense๏ผˆ`WanTransformer3DModel`๏ผ‰ไธ”ๅ†ป็ป“๏ผŒๅ› ๆญค็จ€็–ๆ€งๆ˜ฏ่ขซ่’ธ้ฆ่ฟ›ๅญฆ็”Ÿ๏ผŒ่€Œ้žไปŽๆ•™ๅธˆๆ‹ท่ดใ€‚ +- ๆŽจ็†ๆ—ถ็š„ `student_topk` ้กปไธŽ checkpoint ่ฎญ็ปƒๆ—ถ็š„ๅ–ๅ€ผไธ€่‡ด๏ผ›ไฟฎๆ”นๅฎƒไผš้‡ๅปบๆณจๆ„ๅŠ›็ฎ—ๅญๅนถ้‡็ฝฎๅทฒ่ฎญ็ปƒๅฅฝ็š„้—จๆŽงใ€‚ + +**้‡‡ๆ ทๅฐบๅฏธ้…็ฝฎๆŒ‡ๅ—**๏ผš +- `video_sample_size` ่กจ็คบ่ง†้ข‘็š„ๅˆ†่พจ็Ž‡ๅฐบๅฏธ๏ผ›ๅฝ“ `random_hw_adapt` ไธบ True ๆ—ถ๏ผŒ่กจ็คบ่ง†้ข‘ไธŽๅ›พ็‰‡ๅˆ†่พจ็Ž‡ไน‹้—ด็š„ๆœ€ๅฐๅ€ผใ€‚ +- `image_sample_size` ่กจ็คบๅ›พ็‰‡็š„ๅˆ†่พจ็Ž‡ๅฐบๅฏธ๏ผ›ๅฝ“ `random_hw_adapt` ไธบ True ๆ—ถ๏ผŒ่กจ็คบ่ง†้ข‘ไธŽๅ›พ็‰‡ๅˆ†่พจ็Ž‡ไน‹้—ด็š„ๆœ€ๅคงๅ€ผใ€‚ +- `token_sample_size` ่กจ็คบ `training_with_video_token_length` ไธบ True ๆ—ถๆœ€ๅคง token ้•ฟๅบฆๅฏนๅบ”็š„ๅˆ†่พจ็Ž‡ใ€‚ +- ็”ฑไบŽ้…็ฝฎๅฎนๆ˜“ๆททๆท†๏ผŒ**่‹ฅๅพฎ่ฐƒไธ้œ€่ฆไปปๆ„ๅˆ†่พจ็Ž‡**๏ผŒๅปบ่ฎฎๅฐ† `video_sample_size`ใ€`image_sample_size`ใ€`token_sample_size` ่ฎพไธบๅŒไธ€ๅ›บๅฎšๅ€ผ๏ผŒไพ‹ๅฆ‚ **(320, 480, 512, 640, 960)**ใ€‚ + - **ๅ…จ้ƒจ่ฎพไธบ 320** ไปฃ่กจ **240P**ใ€‚ + - **ๅ…จ้ƒจ่ฎพไธบ 480** ไปฃ่กจ **320P**ใ€‚ + - **ๅ…จ้ƒจ่ฎพไธบ 640** ไปฃ่กจ **480P**ใ€‚ + - **ๅ…จ้ƒจ่ฎพไธบ 960** ไปฃ่กจ **720P**ใ€‚ + +**ๅ…ถๅฎƒๅ‚ๆ•ฐ่ฏดๆ˜Ž**๏ผš +- `enable_bucket` ็”จไบŽๅฏ็”จๅˆ†ๆกถ่ฎญ็ปƒใ€‚ๅฏ็”จๅŽๆจกๅž‹ไธๅœจไธญๅฟƒ่ฃๅ‰ชๅ›พ็‰‡ไธŽ่ง†้ข‘๏ผŒ่€Œๆ˜ฏๆŒ‰ๅˆ†่พจ็Ž‡ๅˆ†ๆกถๅŽๅฏนๅฎŒๆ•ดๅ›พ็‰‡ไธŽ่ง†้ข‘่ฟ›่กŒ่ฎญ็ปƒใ€‚ +- `random_hw_adapt` ็”จไบŽๅฏนๅ›พ็‰‡ไธŽ่ง†้ข‘่‡ชๅŠจ้ซ˜ๅฎฝ็ผฉๆ”พใ€‚ๅฏ็”จๅŽ๏ผŒ่ฎญ็ปƒๅ›พ็‰‡็š„้ซ˜ๅฎฝไปฅ `image_sample_size` ไธบๆœ€ๅคงใ€`min(video_sample_size, 512)` ไธบๆœ€ๅฐใ€‚ +- `training_with_video_token_length` ๆŒ‡ๅฎšๆŒ‰ token ้•ฟๅบฆ่ฎญ็ปƒใ€‚่ฎญ็ปƒๅ›พ็‰‡ไธŽ่ง†้ข‘็š„้ซ˜ๅฎฝไปฅ `image_sample_size` ไธบๆœ€ๅคงใ€`video_sample_size` ไธบๆœ€ๅฐใ€‚ +- `train_mode` ็”จไบŽๆŒ‡ๅฎš่ฎญ็ปƒๆจกๅผ๏ผŒๅฏไธบ normal ๆˆ– i2vใ€‚่‹ฅๅช้œ€ๆ–‡ๆœฌ็”Ÿๆˆ่ง†้ข‘๏ผŒไฟๆŒ `normal`ใ€‚ +- `resume_from_checkpoint` ็”จไบŽ่ฎพ็ฝฎไปŽไน‹ๅ‰็š„ checkpoint ็ปญ่ฎญใ€‚ไฝฟ็”จ่ทฏๅพ„ๆˆ– `"latest"` ่‡ชๅŠจ้€‰ๆ‹ฉๆœ€่ฟ‘ๅฏ็”จ็š„ checkpointใ€‚ + +### 3.4 ่ฎญ็ปƒ้ชŒ่ฏ + +ไฝ ๅฏไปฅ้…็ฝฎ้ชŒ่ฏๅ‚ๆ•ฐ๏ผŒๅœจ่ฎญ็ปƒๆœŸ้—ดๅ‘จๆœŸๆ€ง็”Ÿๆˆๆต‹่ฏ•่ง†้ข‘๏ผŒไปฅไพฟ็›‘ๆŽง่ฎญ็ปƒ่ฟ›ๅบฆไธŽๆจกๅž‹่ดจ้‡ใ€‚ + +**้ชŒ่ฏๅ‚ๆ•ฐ่ฏดๆ˜Ž**๏ผš + +| ๅ‚ๆ•ฐ | ๆ่ฟฐ | ๆŽจ่ๅ€ผ | +|------|------|--------| +| `--validation_steps` | ๆฏ N ๆญฅ่ฟ่กŒไธ€ๆฌก้ชŒ่ฏ | 2000 | +| `--validation_epochs` | ๆฏ N ่ฝฎ่ฟ่กŒไธ€ๆฌก้ชŒ่ฏ | 5 | +| `--validation_prompts` | ็”จไบŽ่ง†้ข‘็”Ÿๆˆ้ชŒ่ฏ็š„ prompt | ่‹ฑๆ–‡ prompt | +| `--validation_paths` | I2V ๆจกๅผ็š„้ชŒ่ฏๅ›พ็‰‡่ทฏๅพ„๏ผˆไป… i2v ๆจกๅผ๏ผ‰ | `"asset/1.png"` | + +**Normal ๆจกๅผ็คบไพ‹**๏ผˆT2V ้ชŒ่ฏ๏ผ‰๏ผš + +```bash + --validation_steps=2000 \ + --validation_epochs=5 \ + --validation_prompts="A dog shaking head. The video is of high quality, and the view is very clear." +``` + +**่ฏดๆ˜Ž**๏ผš +- ้ชŒ่ฏ่ง†้ข‘ไผšไฟๅญ˜ๅˆฐ `output_dir` ็›ฎๅฝ•ใ€‚ +- ๅคš prompt ้ชŒ่ฏๆ ผๅผ๏ผš`--validation_prompts "prompt1" "prompt2" "prompt3"`ใ€‚ +- ่’ธ้ฆๆจกๅž‹้ชŒ่ฏไฝฟ็”จ `denoising_step_indices_list` ๅฎšไน‰็š„ๅฐ‘ๆญฅใ€ๅนถไปฅๆ—  CFG๏ผˆ`guidance_scale=1.0`๏ผ‰่ฟ่กŒ๏ผŒไธŽๆŽจ็†ไฟๆŒไธ€่‡ดใ€‚ + +### 3.5 ไฝฟ็”จ FSDP ่ฎญ็ปƒ + +**่‹ฅไฝฟ็”จ DeepSpeed-Zero-2 ๅคšๅก่ฎญ็ปƒๆ—ถๆ˜พๅญ˜ไธ่ถณ**๏ผŒๅฏๅˆ‡ๆขๅˆฐ FSDP ่ฎญ็ปƒใ€‚`scripts/fastvideo/train_distill_wan2.1.sh` ๆ˜ฏไธ€ไธชๅผ€็ฎฑๅณ็”จ็š„ FSDP launcher๏ผˆๅฎƒไปฅ `FULL_SHARD` ๅŒ…่ฃ… `WanAttentionBlock`๏ผ‰ใ€‚ + +```sh +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +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. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + +accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +### 3.6 ๅ…ถไป–ๅŽ็ซฏ + +#### 3.6.1 ไฝฟ็”จ DeepSpeed-Zero-3 ่ฟ›่กŒ่ฎญ็ปƒ + +็›ฎๅ‰ไธๅผบ็ƒˆๆŽจ่ DeepSpeed Zero-3ใ€‚ๅœจๆœฌไป“ๅบ“ไธญ๏ผŒไฝฟ็”จ FSDP ๆŠฅ้”™ๆ›ดๅฐ‘ใ€ๆ›ด็จณๅฎšใ€‚ + +DeepSpeed Zero-3 ้€‚ๅˆ้ซ˜ๅˆ†่พจ็Ž‡ไธ‹็š„ 14B Wanใ€‚่ฎญ็ปƒๅŽๅฏ็”จไปฅไธ‹ๅ‘ฝไปคๅพ—ๅˆฐๆœ€็ปˆๆจกๅž‹๏ผš +```bash +python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization +``` + +่ฎญ็ปƒ shell ๅ‘ฝไปคๅฆ‚ไธ‹๏ผš +```sh +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +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/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +#### 3.6.2 ไธไฝฟ็”จ DeepSpeed ไธŽ FSDP ่ฎญ็ปƒ + +**ไธๆŽจ่่ฟ™็งๆ–นๅผ๏ผŒๅ› ไธบๆฒกๆœ‰็œๆ˜พๅญ˜ๅŽ็ซฏๆ—ถๆžๆ˜“ๅฏผ่‡ดๆ˜พๅญ˜ไธ่ถณ**ใ€‚่ฟ™้‡Œไป…ไฝœไธบ่ฎญ็ปƒๅ‚่€ƒ shell ๆไพ›ใ€‚ + +```sh +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +export DATASET_NAME="datasets/X-Fun-Videos-Demo/" +export DATASET_META_NAME="datasets/X-Fun-Videos-Demo/metadata_add_width_height.json" +NCCL_DEBUG=INFO + +accelerate launch --mixed_precision="bf16" scripts/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +### 3.7 ๅคšๆœบๅˆ†ๅธƒๅผ่ฎญ็ปƒ + +**้€‚็”จไบŽ**๏ผš่ถ…ๅคงๆ•ฐๆฎ้›†ใ€ๆ›ดๅฟซ็š„่ฎญ็ปƒ้€Ÿๅบฆ + +#### 3.7.1 ็Žฏๅขƒ้…็ฝฎ + +ๅ‡่ฎพๆœ‰ 2 ๅฐๆœบๅ™จ๏ผŒๆฏๅฐ 8 ๅผ  GPU๏ผš + +**ๆœบๅ™จ 0๏ผˆMaster๏ผ‰**๏ผš +```bash +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +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 +export MASTER_PORT=10086 +export WORLD_SIZE=2 # ๆœบๅ™จๆ€ปๆ•ฐ +export NUM_PROCESS=16 # ๆ€ป่ฟ›็จ‹ๆ•ฐ = ๆœบๅ™จๆ•ฐ ร— 8 +export RANK=0 # ๅฝ“ๅ‰ๆœบๅ™จ rank๏ผˆ0 ๆˆ– 1๏ผ‰ +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/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram +``` + +**ๆœบๅ™จ 1๏ผˆWorker๏ผ‰**๏ผš +```bash +export MASTER_ADDR="192.168.1.100" # ไธŽ Master ็›ธๅŒ +export MASTER_PORT=10086 +export WORLD_SIZE=2 +export NUM_PROCESS=16 +export RANK=1 # ๆณจๆ„่ฟ™้‡Œๆ˜ฏ 1 +NCCL_DEBUG=INFO + +# ไฝฟ็”จไธŽๆœบๅ™จ 0 ็›ธๅŒ็š„ accelerate launch ๅ‘ฝไปค +``` + +#### 3.7.2 ๅคšๆœบ่ฎญ็ปƒๆณจๆ„ไบ‹้กน + +- **็ฝ‘็ปœ่ฆๆฑ‚**๏ผš + - ๆŽจ่ RDMA/InfiniBand๏ผˆ้ซ˜ๆ€ง่ƒฝ๏ผ‰ + - ๆฒกๆœ‰ RDMA ๆ—ถ๏ผŒๆทปๅŠ ็Žฏๅขƒๅ˜้‡๏ผš + ```bash + export NCCL_IB_DISABLE=1 + export NCCL_P2P_DISABLE=1 + ``` +- **ๆ•ฐๆฎๅŒๆญฅ**๏ผšๆ‰€ๆœ‰ๆœบๅ™จๅฟ…้กป่ƒฝ่ฎฟ้—ฎ็›ธๅŒ็š„ๆ•ฐๆฎ่ทฏๅพ„๏ผˆNFS/ๅ…ฑไบซๅญ˜ๅ‚จ๏ผ‰ + +--- + +## ๅ››ใ€ๆŽจ็†ๆต‹่ฏ• + +### 4.1 ๆŽจ็†ๅ‚ๆ•ฐ่งฃๆž + +**ๅ…ณ้”ฎๅ‚ๆ•ฐ่ฏดๆ˜Ž**๏ผˆๅ–่‡ช `examples/fastvideo/predict_t2v_wan2.1.py`๏ผ‰๏ผš + +| ๅ‚ๆ•ฐ | ๆ่ฟฐ | ็คบไพ‹ๅ€ผ | +|------|------|-------| +| `GPU_memory_mode` | GPU ๆ˜พๅญ˜ๆจกๅผ๏ผŒ้€‰้กน่งไธ‹่กจ | `model_cpu_offload` | +| `ulysses_degree` | Ulysses ๅนถ่กŒๅบฆ๏ผˆVSA ไป…ๅ•ๅก๏ผŒไฟๆŒ 1๏ผ‰ | 1 | +| `ring_degree` | Ring ๅนถ่กŒๅบฆ๏ผˆVSA ไป…ๅ•ๅก๏ผŒไฟๆŒ 1๏ผ‰ | 1 | +| `config_path` | ้…็ฝฎๆ–‡ไปถ่ทฏๅพ„ | `config/wan2.1/wan_civitai.yaml` | +| `model_name` | ๅŸบๅบงๆจกๅž‹่ทฏๅพ„๏ผˆconfig + dense ไธปๅนฒ + VAE + T5๏ผ‰ | `models/Diffusion_Transformer/Wan2.1-T2V-1.3B` | +| `student_topk` | VSA ไฟ็•™ๆฏ”๏ผ›้กปไธŽ่ฎญ็ปƒๆ—ถๅ–ๅ€ผไธ€่‡ด | `0.20` | +| `sampler_name` | ้‡‡ๆ ทๅ™จ็ฑปๅž‹๏ผš`Flow`ใ€`Flow_Unipc`ใ€`Flow_DPM++` | `Flow_Unipc` | +| `shift` | ๅ™ชๅฃฐ่ฐƒๅบฆ shift๏ผˆๅฎ˜ๆ–น FastWan flow_shift๏ผ‰ | `3.0` | +| `transformer_path` | FastWan ๅญฆ็”Ÿ Transformer ๆƒ้‡่ทฏๅพ„ | ่ฎญ็ปƒ checkpoint / ๅฎ˜ๆ–น FastWan | +| `vae_path` | ่ฎญ็ปƒๅฅฝ็š„ VAE ๆƒ้‡่ทฏๅพ„ | `None` | +| `lora_path` | LoRA ๆƒ้‡่ทฏๅพ„ | `None` | +| `sample_size` | ็”Ÿๆˆ่ง†้ข‘ๅˆ†่พจ็Ž‡ `[height, width]` | `[480, 848]` | +| `video_length` | ็”Ÿๆˆๅธงๆ•ฐ | `81` | +| `fps` | ๆฏ็ง’ๅธงๆ•ฐ | `16` | +| `weight_dtype` | ๆจกๅž‹ๆƒ้‡็ฒพๅบฆ๏ผŒไธๆ”ฏๆŒ bf16 ็š„ GPU ็”จ `torch.float16` | `torch.bfloat16` | +| `prompt` | ๆ่ฟฐ็”Ÿๆˆๅ†…ๅฎน็š„ๆญฃๅ‘ prompt | `"A cinematic aerial shot..."` | +| `negative_prompt` | ็”จไบŽ่ง„้ฟ็‰นๅฎšๅ†…ๅฎน็š„่ดŸๅ‘ prompt | ไธญๆ–‡่ดŸๅ‘ prompt | +| `guidance_scale` | ๅผ•ๅฏผๅผบๅบฆ๏ผˆDMD ๅญฆ็”Ÿๆ—  CFG๏ผ‰ | 1.0 | +| `seed` | ๅค็Žฐ็”จ้šๆœบ็งๅญ | 1024 | +| `num_inference_steps` | ๆŽจ็†ๆญฅๆ•ฐ๏ผˆFastWan ็”จ 3 ๆญฅ๏ผ‰ | 3 | +| `save_path` | ็”Ÿๆˆ่ง†้ข‘ไฟๅญ˜่ทฏๅพ„ | `samples/fastvideo-wan2.1-official` | + +**GPU ๆ˜พๅญ˜ๆจกๅผ่ฏดๆ˜Ž**๏ผš + +| ๆจกๅผ | ๆ่ฟฐ | ๆ˜พๅญ˜ๅ ็”จ | +|------|------|---------| +| `model_full_load` | ๆ•ดไธชๆจกๅž‹ๅŠ ่ฝฝๅˆฐ GPU | ๆœ€้ซ˜ | +| `model_full_load_and_qfloat8` | ๅ…จ้‡ๅŠ ่ฝฝ + FP8 ้‡ๅŒ– | ้ซ˜ | +| `model_cpu_offload` | ็”จๅŽๅฐ†ๆจกๅž‹ๅธ่ฝฝๅˆฐ CPU | ไธญ | +| `model_cpu_offload_and_qfloat8` | CPU ๅธ่ฝฝ + FP8 ้‡ๅŒ– | ไธญไฝŽ | +| `model_group_offload` | ๅฑ‚็ป„ๅœจ CPU/CUDA ้—ดๅˆ‡ๆข | ไฝŽ | +| `sequential_cpu_offload` | ้€ๅฑ‚ๅธ่ฝฝ๏ผˆๆœ€ๆ…ข๏ผ‰ | ๆœ€ไฝŽ | + +### 4.2 ๆ–‡ๆœฌ็”Ÿๆˆ่ง†้ข‘๏ผˆT2V๏ผ‰ๆŽจ็† + +่ฟ่กŒๅ•ๅกๆŽจ็†๏ผš + +```bash +python examples/fastvideo/predict_t2v_wan2.1.py +``` + +ๆ นๆฎ้œ€ๆฑ‚็ผ–่พ‘ `examples/fastvideo/predict_t2v_wan2.1.py`ใ€‚้ฆ–ๆฌกๆŽจ็†้‡็‚นๅ…ณๆณจไปฅไธ‹ๅ…ณ้”ฎๅ‚ๆ•ฐ๏ผŒๅ…ถๅฎƒๅ‚ๆ•ฐๅ‚่งไธŠๆ–‡ๆŽจ็†ๅ‚ๆ•ฐ่ฏดๆ˜Žใ€‚ + +```python +# ๆ นๆฎ GPU ๆ˜พๅญ˜้€‰ๆ‹ฉ +GPU_memory_mode = "model_cpu_offload" +# ๅŸบๅบงๆจกๅž‹่ทฏๅพ„๏ผˆconfig + dense ไธปๅนฒ + VAE + T5๏ผ‰ +model_name = "models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +# VSA ไฟ็•™ๆฏ”๏ผ›้กปไธŽ checkpoint ่ฎญ็ปƒๆ—ถๅ–ๅ€ผไธ€่‡ด +student_topk = 0.20 +# FastWan ๅญฆ็”Ÿๆƒ้‡่ทฏๅพ„๏ผˆไฝ ่ฎญ็ปƒ็š„ checkpoint๏ผŒๆˆ–ๅฎ˜ๆ–น FastWan ๅ‘ๅธƒๆƒ้‡๏ผ‰ +transformer_path = "models/Personalized_Model/FastWan2.1-T2V-1.3B.pth" +# FastWan ้‡‡ๆ ทๅ™จ + shift๏ผˆๆ—  CFG ๅฐ‘ๆญฅๅญฆ็”Ÿ๏ผ‰ +sampler_name = "Flow_Unipc" +shift = 3.0 +guidance_scale = 1.0 +num_inference_steps = 3 +# ๆ นๆฎไฝ ็š„็”Ÿๆˆๅ†…ๅฎน็ผ–ๅ†™ +prompt = "A cinematic aerial shot of a snow-covered mountain village at dawn..." +# ... +``` + +> ๐Ÿ’ก ๅฎ˜ๆ–น `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` checkpoint ้‡‡็”จ diffusers ้”ฎๅ๏ผ›`examples/fastvideo/predict_t2v_wan2.1.py` ไผšๅœจๅŠ ่ฝฝๅ‰ๅฐ†ๅ…ถ remap ๅˆฐ VideoX-Fun civitai ๅธƒๅฑ€๏ผˆๅŒ…ๅซ VSA `to_gate_compress` ้—จๆŽง๏ผ‰ใ€‚ๆœฌ่ฎญ็ปƒๆŒ‡ๅ—ไบงๅ‡บ็š„ checkpoint ๅทฒๆ˜ฏๅŽŸ็”Ÿๅธƒๅฑ€ใ€‚ + +### 4.3 ไป…ๆ”ฏๆŒๅ•ๅก + +VSA ๅ—็จ€็–ๅ†…ๆ ธๆฒกๆœ‰ๅบๅˆ—ๅนถ่กŒ่ทฏๅพ„๏ผŒๅ› ๆญค FastWan ๅญฆ็”Ÿๅช่ƒฝๅœจ**ๅ•ๅก**ไธŠ่ฟ่กŒใ€‚่ฏทไฟๆŒ `ulysses_degree = ring_degree = 1`๏ผ›ๅฆๅˆ™่ฐƒ็”จ `enable_multi_gpus_inference()` ไผš raiseใ€‚่ฆๅœจๅ•ๅกไธŠๅฎน็บณๆ›ดๅคงๅˆ†่พจ็Ž‡๏ผŒ่ฏท้€š่ฟ‡ `GPU_memory_mode`๏ผˆๅฆ‚ `model_cpu_offload` ๆˆ– `sequential_cpu_offload`๏ผ‰้™ไฝŽๆ˜พๅญ˜ๅ ็”จ๏ผŒ่€Œไธๆ˜ฏ่ทจๅกๆ‹†ๅˆ†ใ€‚ + +--- + +## ไบ”ใ€ๆ›ดๅคš่ต„ๆบ + +- **ๅฎ˜ๆ–น GitHub**๏ผšhttps://github.com/aigc-apps/VideoX-Fun +- **FastVideo GitHub**๏ผšhttps://github.com/hao-ai-lab/FastVideo +- **VSA ่ฎบๆ–‡๏ผˆVideo Sparse Attention๏ผ‰**๏ผšhttps://arxiv.org/abs/2505.13389 diff --git a/scripts/fastvideo/train_distill_wan2.1.py b/scripts/fastvideo/train_distill_wan2.1.py new file mode 100644 index 0000000..a7bca5a --- /dev/null +++ b/scripts/fastvideo/train_distill_wan2.1.py @@ -0,0 +1,2411 @@ +"""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 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, deprecate, 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, + FastVideoWanTransformer3DModel, WanT5EncoderModel, + WanTransformer3DModel) +from videox_fun.pipeline import WanI2VPipeline, WanPipeline +from videox_fun.utils.discrete_sampler import DiscreteSampling +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 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": + pipeline = WanI2VPipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d, + scheduler=scheduler, + clip_image_encoder=clip_image_encoder, + ) + else: + pipeline = WanPipeline( + 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": + start_image = Image.open(args.validation_paths[i]) + width, height = start_image.width, start_image.height + width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height) + + video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1 + input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[args.video_sample_size, args.video_sample_size]) + sample = pipeline( + args.validation_prompts[i], + num_frames = video_length, + negative_prompt = "่‰ฒ่ฐƒ่‰ณไธฝ๏ผŒ่ฟ‡ๆ›๏ผŒ้™ๆ€๏ผŒ็ป†่Š‚ๆจก็ณŠไธๆธ…๏ผŒๅญ—ๅน•๏ผŒ้ฃŽๆ ผ๏ผŒไฝœๅ“๏ผŒ็”ปไฝœ๏ผŒ็”ป้ข๏ผŒ้™ๆญข๏ผŒๆ•ดไฝ“ๅ‘็ฐ๏ผŒๆœ€ๅทฎ่ดจ้‡๏ผŒไฝŽ่ดจ้‡๏ผŒJPEGๅŽ‹็ผฉๆฎ‹็•™๏ผŒไธ‘้™‹็š„๏ผŒๆฎ‹็ผบ็š„๏ผŒๅคšไฝ™็š„ๆ‰‹ๆŒ‡๏ผŒ็”ปๅพ—ไธๅฅฝ็š„ๆ‰‹้ƒจ๏ผŒ็”ปๅพ—ไธๅฅฝ็š„่„ธ้ƒจ๏ผŒ็•ธๅฝข็š„๏ผŒๆฏๅฎน็š„๏ผŒๅฝขๆ€็•ธๅฝข็š„่‚ขไฝ“๏ผŒๆ‰‹ๆŒ‡่žๅˆ๏ผŒ้™ๆญขไธๅŠจ็š„็”ป้ข๏ผŒๆ‚ไนฑ็š„่ƒŒๆ™ฏ๏ผŒไธ‰ๆก่…ฟ๏ผŒ่ƒŒๆ™ฏไบบๅพˆๅคš๏ผŒๅ€’็€่ตฐ", + height = args.video_sample_size, + width = args.video_sample_size, + generator = generator, + + video = input_video, + mask_video = input_video_mask, + num_inference_steps = len(args.denoising_step_indices_list), + guidance_scale = 1.0, + ).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" + ) + ) + else: + sample = pipeline( + args.validation_prompts[i], + num_frames = args.video_sample_n_frames, + negative_prompt = "่‰ฒ่ฐƒ่‰ณไธฝ๏ผŒ่ฟ‡ๆ›๏ผŒ้™ๆ€๏ผŒ็ป†่Š‚ๆจก็ณŠไธๆธ…๏ผŒๅญ—ๅน•๏ผŒ้ฃŽๆ ผ๏ผŒไฝœๅ“๏ผŒ็”ปไฝœ๏ผŒ็”ป้ข๏ผŒ้™ๆญข๏ผŒๆ•ดไฝ“ๅ‘็ฐ๏ผŒๆœ€ๅทฎ่ดจ้‡๏ผŒไฝŽ่ดจ้‡๏ผŒJPEGๅŽ‹็ผฉๆฎ‹็•™๏ผŒไธ‘้™‹็š„๏ผŒๆฎ‹็ผบ็š„๏ผŒๅคšไฝ™็š„ๆ‰‹ๆŒ‡๏ผŒ็”ปๅพ—ไธๅฅฝ็š„ๆ‰‹้ƒจ๏ผŒ็”ปๅพ—ไธๅฅฝ็š„่„ธ้ƒจ๏ผŒ็•ธๅฝข็š„๏ผŒๆฏๅฎน็š„๏ผŒๅฝขๆ€็•ธๅฝข็š„่‚ขไฝ“๏ผŒๆ‰‹ๆŒ‡่žๅˆ๏ผŒ้™ๆญขไธๅŠจ็š„็”ป้ข๏ผŒๆ‚ไนฑ็š„่ƒŒๆ™ฏ๏ผŒไธ‰ๆก่…ฟ๏ผŒ่ƒŒๆ™ฏไบบๅพˆๅคš๏ผŒๅ€’็€่ตฐ", + height = args.video_sample_size, + width = args.video_sample_size, + generator = generator, + num_inference_steps = len(args.denoising_step_indices_list), + guidance_scale = 1.0, + ).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( + "--non_ema_revision", + type=str, + default=None, + required=False, + help=( + "Revision of pretrained non-ema model identifier. Must be a branch, tag or git identifier of the local or" + " remote repository specified with --pretrained_model_name_or_path." + ), + ) + 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( + "--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=6.0, + help="The cfg scale for real score.", + ) + parser.add_argument( + '--denoising_step_indices_list', + nargs='+', + default=[1000, 750, 500, 250], + help="The denoising step list.", + ) + parser.add_argument( + "--randomize_step_indices", + action="store_true", + help="whether to use randomize timesteps indices in training.", + ) + + 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 + + # default to using the same revision for the non-ema model if not specified + if args.non_ema_revision is None: + args.non_ema_revision = args.revision + + return args + + +def main(): + args = parse_args() + + 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." + ) + + if args.non_ema_revision is not None: + deprecate( + "non_ema_revision!=None", + "0.15.0", + message=( + "Downloading 'non_ema' weights from revision branches of the Hub is deprecated. Please make sure to" + " use `--variant=non_ema` instead." + ), + ) + 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 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}") + + 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 + + args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] + # 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 + generator_transformer3d = FastVideoWanTransformer3DModel.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, + ).to(weight_dtype) + real_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, + ).to(weight_dtype) + fake_score_transformer3d = FastVideoWanTransformer3DModel.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, + ).to(weight_dtype) + + # 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) + + 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 + + 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.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 + + # `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_stage != 0 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=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: + accelerator.unwrap_model(models[0]).save_pretrained(os.path.join(output_dir, "transformer")) + 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 = accelerator.unwrap_model(models.pop()) + + # load diffusers style into model + load_model = FastVideoWanTransformer3DModel.from_pretrained( + input_dir, subfolder="transformer", + low_cpu_mem_usage=True, + torch_dtype=next(model.parameters()).dtype, + ) + model.set_student_topk(load_model.student_topk) + 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.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 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": + 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": + 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 + + 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) + + # 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) + + 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 + + for epoch in range(first_epoch, args.num_train_epochs): + 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 + 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": + # 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 + + 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, :, :] + 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, :, :] + 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) + 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) + + # 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) + + 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] + + target_shape = ( + len(text), + vae.latent_channels, + int((num_frames - 1) // vae.temporal_compression_ratio + 1), + int(local_sample_size[0] // vae.spatial_compression_ratio), + int(local_sample_size[1] // vae.spatial_compression_ratio), + ) + + 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 + 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) + patch_h, patch_w = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + seq_len = math.ceil((width * height) / (patch_h * patch_w) * num_frames) + + 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: + generator_pred = generator_transformer3d( + x=generator_noise, + 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, + ) + 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 + ) + + indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu() + generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device) + 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) + + 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 + + if args.low_vram: + real_score_transformer3d = real_score_transformer3d.to("cpu") + fake_score_transformer3d = fake_score_transformer3d.to("cpu") + torch.cuda.empty_cache() + + accelerator.backward(dmd_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.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 --- + with torch.no_grad(): + 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] + + patch_h, patch_w = accelerator.unwrap_model(generator_transformer3d).config.patch_size[1:] + seq_len = math.ceil((width * height) / (patch_h * patch_w) * num_frames) + + 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): + fake_score_denoised_pred = generator_transformer3d( + x=fake_score_critic_noise, + 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, + ) + 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) + 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 + ) + + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): + 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, + ) + + 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 + + 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 + + accelerator_fake_score_transformer3d.backward(denoising_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: + + progress_bar.update(1) + global_step += 1 + accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step) + train_dmd_loss = 0.0 + train_gen_log_count = 0 + train_denoising_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(): + 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 args.validation_prompts is not None and epoch % args.validation_epochs == 0: + with progress_bar.paused(): + 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 accelerator.is_main_process: + generator_transformer3d = unwrap_model(generator_transformer3d) + + 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/fastvideo/train_distill_wan2.1.sh b/scripts/fastvideo/train_distill_wan2.1.sh new file mode 100644 index 0000000..bc0af53 --- /dev/null +++ b/scripts/fastvideo/train_distill_wan2.1.sh @@ -0,0 +1,48 @@ +# This shell script is used to distill a regular Wan2.1-1.3B into FastWan2.1-1.3B (VSA sparse student) +export MODEL_NAME="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" +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. +# export NCCL_IB_DISABLE=1 +# export NCCL_P2P_DISABLE=1 +NCCL_DEBUG=INFO + + +accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP \ + --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" \ + --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" \ + --fsdp_cpu_ram_efficient_loading False scripts/fastvideo/train_distill_wan2.1.py \ + --config_path="config/wan2.1/wan_civitai.yaml" \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --train_data_dir=$DATASET_NAME \ + --train_data_meta=$DATASET_META_NAME \ + --image_sample_size=640 \ + --video_sample_size=640 \ + --token_sample_size=640 \ + --video_sample_stride=2 \ + --video_sample_n_frames=81 \ + --train_batch_size=1 \ + --video_repeat=1 \ + --gradient_accumulation_steps=1 \ + --dataloader_num_workers=8 \ + --num_train_epochs=100 \ + --checkpointing_steps=50 \ + --learning_rate=2e-06 \ + --learning_rate_critic=2e-06 \ + --lr_scheduler="constant_with_warmup" \ + --lr_warmup_steps=100 \ + --seed=42 \ + --output_dir="output_dir_fastvideo_wan2.1_1.3b" \ + --gradient_checkpointing \ + --mixed_precision="bf16" \ + --adam_weight_decay=3e-2 \ + --adam_epsilon=1e-10 \ + --vae_mini_batch=1 \ + --max_grad_norm=0.05 \ + --random_hw_adapt \ + --training_with_video_token_length \ + --enable_bucket \ + --uniform_sampling \ + --train_mode="normal" \ + --trainable_modules "." \ + --low_vram diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md b/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md index 92fea91..745ee4e 100644 --- a/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION.md @@ -172,6 +172,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 \ @@ -211,6 +212,7 @@ bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh | `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ar_diffusion` | | `--gradient_checkpointing` | Enable activation checkpointing | - | | `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--fsdp_reduce_fp32` | FP32 gradient reduction under FSDP mixed precision; forward and parameter gathering stay low precision. CLI default is off; enabled in the launcher | enabled | | `--adam_weight_decay` | AdamW weight decay | 3e-2 | | `--adam_epsilon` | AdamW epsilon | 1e-10 | | `--max_grad_norm` | Gradient clipping threshold | 0.05 | diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md b/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md index 08aefe2..3e28894 100644 --- a/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_AR_DIFFUSION_zh-CN.md @@ -172,6 +172,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 \ @@ -211,6 +212,7 @@ bash scripts/wan2.1_causal_forcing/train_ar_diffusion.sh | `--output_dir` | ่พ“ๅ‡บ็›ฎๅฝ• | `output_dir_wan2.1_causal_forcing_ar_diffusion` | | `--gradient_checkpointing` | ๅฏ็”จๆฟ€ๆดป้‡่ฎก็ฎ— | - | | `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--fsdp_reduce_fp32` | FSDP ๆททๅˆ็ฒพๅบฆไธ‹ไฝฟ็”จ FP32 ๆขฏๅบฆๅฝ’็บฆ๏ผŒๅ‰ๅ‘ไธŽๅ‚ๆ•ฐ่šๅˆไปไฟๆŒไฝŽ็ฒพๅบฆใ€‚CLI ้ป˜่ฎคๅ…ณ้—ญ๏ผŒlauncher ๅทฒๅฏ็”จ | ๅฏ็”จ | | `--adam_weight_decay` | AdamW ๆƒ้‡่กฐๅ‡ | 3e-2 | | `--adam_epsilon` | AdamW epsilon | 1e-10 | | `--max_grad_norm` | ๆขฏๅบฆ่ฃๅ‰ช้˜ˆๅ€ผ | 0.05 | diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md index 5333d9f..c3dce4f 100644 --- a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md @@ -183,6 +183,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 \ @@ -228,6 +229,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh | `--output_dir` | Output directory | `output_dir_wan2.1_causal_forcing_ccd` | | `--gradient_checkpointing` | Enable activation checkpointing | - | | `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--fsdp_reduce_fp32` | FP32 gradient reduction under FSDP mixed precision; forward and parameter gathering stay low precision. CLI default is off; enabled in the launcher | enabled | | `--adam_weight_decay` | AdamW weight decay | 0.0 | | `--adam_beta1` | AdamW beta1 (CCD uses 0.0) | 0.0 | | `--adam_beta2` | AdamW beta2 | 0.999 | diff --git a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md index 6d1f36f..9befd85 100644 --- a/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md +++ b/scripts/wan2.1_causal_forcing/README_TRAIN_CAUSAL_CONSISTENCY_DISTILL_zh-CN.md @@ -183,6 +183,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 \ @@ -228,6 +229,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_consistency_distill.sh | `--output_dir` | ่พ“ๅ‡บ็›ฎๅฝ• | `output_dir_wan2.1_causal_forcing_ccd` | | `--gradient_checkpointing` | ๅฏ็”จๆฟ€ๆดป้‡่ฎก็ฎ— | - | | `--mixed_precision` | `fp16` / `bf16` | `bf16` | +| `--fsdp_reduce_fp32` | FSDP ๆททๅˆ็ฒพๅบฆไธ‹ไฝฟ็”จ FP32 ๆขฏๅบฆๅฝ’็บฆ๏ผŒๅ‰ๅ‘ไธŽๅ‚ๆ•ฐ่šๅˆไปไฟๆŒไฝŽ็ฒพๅบฆใ€‚CLI ้ป˜่ฎคๅ…ณ้—ญ๏ผŒlauncher ๅทฒๅฏ็”จ | ๅฏ็”จ | | `--adam_weight_decay` | AdamW ๆƒ้‡่กฐๅ‡ | 0.0 | | `--adam_beta1` | AdamW beta1๏ผˆCCD ไฝฟ็”จ 0.0๏ผ‰ | 0.0 | | `--adam_beta2` | AdamW beta2 | 0.999 | 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 263f2fb..3f935c8 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 @@ -26,6 +26,7 @@ This document provides the complete workflow for **Causal-Forcing Stage 3 โ€” Di - [5.1 Quick Start](#51-quick-start) - [5.2 Key Parameters](#52-key-parameters) - [5.3 DMD-Specific Parameters](#53-dmd-specific-parameters) + - [5.4 Optional One-Forcing GAN and Precision Controls](#54-optional-one-forcing-gan-and-precision-controls) - [6. Use the Trained Checkpoint](#6-use-the-trained-checkpoint) - [7. Additional Resources](#7-additional-resources) @@ -103,7 +104,7 @@ The Stage 2 CCD checkpoint (generator/critic init) is loaded via `--ode_transfor ## 4. Prepare Training Data -DMD uses a **TextDataset** (`train_mode="normal"`) โ€” it only needs prompts, not video data, because the generator creates its own training samples via autoregressive rollout. However, a `metadata.json` with prompts is still required. +With `train_mode="normal"`, teacher forcing disabled, and all GAN/R1/R2 weights set to zero, DMD uses a **TextDataset**: only prompts are required because the generator creates training samples via autoregressive rollout. Enabling any GAN/R1/R2 weight switches to **ImageVideoDataset** and requires captioned real videos; see Section 5.4. ### 4.1 Quick Demo Dataset @@ -118,7 +119,7 @@ Same as Stage 1/2. See [README_TRAIN_CAUSAL_CONSISTENCY_DISTILL.md](./README_TRA ### 4.3 metadata.json Format -DMD only uses the `"text"` field from each entry. The `file_path` and other fields are ignored when `--train_mode="normal"` (TextDataset mode). +In prompt-only TextDataset mode, DMD only uses the `"text"` field from each entry. GAN/R1/R2 or teacher forcing requires video metadata with `file_path`, `type="video"`, and `text`, as in Stage 1/2. ```json [ @@ -131,7 +132,7 @@ DMD only uses the `"text"` field from each entry. The `file_path` and other fiel ] ``` -> **Note**: You can reuse any video metadata.json โ€” only the `text` field matters for DMD. +> **Note**: Existing video metadata can be reused. Only `text` is read in prompt-only mode; GAN training also reads the video files. --- @@ -179,6 +180,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \ --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 \ @@ -212,7 +214,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh | `--real_score_pretrained_model_name_or_path` | 14B real-score teacher for DMD | `models/Diffusion_Transformer/Wan2.1-T2V-14B` | | `--config_path` | Model config YAML | `config/wan2.1/wan_civitai.yaml` | | `--train_data_dir` | Data root directory | `""` | -| `--train_data_meta` | Path to `metadata.json` (prompts only) | `datasets/X-Fun-Videos-Demo/metadata.json` | +| `--train_data_meta` | Metadata path: prompts for plain DMD, captioned videos for GAN/R1/R2 | `datasets/X-Fun-Videos-Demo/metadata.json` | | `--ode_transformer_path` | Stage 2 CCD checkpoint for generator/critic init | `$STAGE2_CKPT` | | `--train_batch_size` | Per-GPU batch size | 1 | | `--gradient_accumulation_steps` | Gradient accumulation steps | 1 | @@ -264,11 +266,48 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh | `--context_noise` | Context noise level for KV cache update | 0 | | `--use_teacher_forcing` | Enable teacher forcing (pass clean_x to transformer) | - | | `--teacher_forcing_prob` | Probability of applying teacher forcing per step (1.0 = always) | 1.0 | -| `--train_mode` | Training mode: `normal` (TextDataset, prompt-only) or `i2v` | `normal` | +| `--train_mode` | `normal` (prompt-only unless teacher forcing or GAN/R1/R2 is enabled) or `i2v` | `normal` | | `--resume_from_checkpoint` | Resume from checkpoint. Use `"latest"` to auto-select | `"latest"` | --- +### 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. + +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`. + +| Parameter | Description | Default | +|-----------|-------------|---------| +| `--gan_g_weight` | Non-saturating generator GAN loss weight | `0.0` | +| `--gan_d_weight` | Discriminator GAN loss weight | `0.0` | +| `--gan_feature_layers` | Fake-score block indices used by the head | `21 29` | +| `--gan_block_ffn_dim` | Head cross-attention FFN dimension | `2048` | +| `--gan_block_num_heads` | Head count; must divide fake-score hidden dimension | `12` | +| `--relativistic_discriminator` | Use paired real/fake logit differences | off | +| `--r1_weight` | APT-style finite-difference penalty on noisy real latents | `0.0` | +| `--r1_sigma` | R1 perturbation standard deviation | `0.01` | +| `--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 | + +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))`. + +R1 follows the source's APT form: `r1_weight * mean((D_real_perturbed - D_real)^2)`, **without dividing by sigma**. R2 follows `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`. These are finite-difference penalties, not input-gradient double-backward regularizers. Real/fake samples share timestep and noise; perturbations use independent noise. + +To enable GAN, pass positive generator/discriminator weights and, optionally, R1/R2 weights to `train_causal_dmd.py`. Requirements: +- Captioned real videos, not a prompt-only JSON. Use the Stage 1/2 video metadata schema and enough frames for `score_num_frames` and the selected causal block geometry. +- `--gradient_accumulation_steps=1`; DeepSpeed is not supported by the GAN branch. +- Valid feature-layer indices for the selected fake-score model. The default `21 29` layers and `12` heads fit Wan2.1-T2V-1.3B; other backbones may need different settings. +- Enabling only `gan_g_weight` does not supply a discriminator classification-learning objective. Choose the discriminator weight deliberately; nonzero values are not a guaranteed quality improvement. + +Trainable fake-score and head master weights stay **FP32**, while forward uses the configured autocast/FSDP precision. Logits, softplus losses, and finite-difference reductions are evaluated in FP32. `--fsdp_reduce_fp32` changes gradient communication precision, not the forward dtype or model architecture. The clamp changes score-training timesteps, not the generator rollout or validation schedule. + +Resume GAN checkpoints with the same GAN/head configuration. The head is saved with fake-score, including FP32 head tensors, and is recreated before loading. It is not part of the exported generator, so inference needs no discriminator or new GAN flags. + +--- + ## 6. Use the Trained Checkpoint The Stage 3 DMD checkpoint is the **final model** in the Causal-Forcing pipeline. Use it for inference: @@ -300,4 +339,5 @@ python examples/wan2.1_causal_forcing/predict_t2v.py ## 7. Additional Resources - **Causal-Forcing Paper**: https://github.com/thu-ml/Causal-Forcing +- **One-Forcing (optional GAN mechanism)**: https://github.com/Aurora-edu/One-Forcing - **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun 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 f3a24ff..3e399d1 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 @@ -26,6 +26,7 @@ - [5.1 ๅฟซ้€Ÿๅผ€ๅง‹](#51-ๅฟซ้€Ÿๅผ€ๅง‹) - [5.2 ๅ…ณ้”ฎๅ‚ๆ•ฐ](#52-ๅ…ณ้”ฎๅ‚ๆ•ฐ) - [5.3 DMD ็‰นๆœ‰ๅ‚ๆ•ฐ](#53-dmd-็‰นๆœ‰ๅ‚ๆ•ฐ) + - [5.4 ๅฏ้€‰ One-Forcing GAN ไธŽ็ฒพๅบฆๆŽงๅˆถ](#54-ๅฏ้€‰-one-forcing-gan-ไธŽ็ฒพๅบฆๆŽงๅˆถ) - [ๅ…ญใ€ไฝฟ็”จ่ฎญ็ปƒๅฅฝ็š„ Checkpoint](#ๅ…ญไฝฟ็”จ่ฎญ็ปƒๅฅฝ็š„-checkpoint) - [ไธƒใ€ๆ›ดๅคš่ต„ๆบ](#ไธƒๆ›ดๅคš่ต„ๆบ) @@ -39,7 +40,7 @@ Stage 3 ้œ€่ฆ๏ผš 2. **Wan2.1-T2V-14B** ๆจกๅž‹๏ผšไฝœไธบ DMD real-score teacherใ€‚ ```bash -# ็คบไพ‹๏ผšStage 2 CCD ่ฎญ็ปƒ็š„ checkpoint +# Example: Stage 2 CCD checkpoint export STAGE2_CKPT="output_dir_wan2.1_causal_forcing_ccd/checkpoint-5000/transformer/diffusion_pytorch_model.safetensors" ``` @@ -70,10 +71,10 @@ 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 ``` @@ -87,13 +88,13 @@ DMD ้œ€่ฆ**ไธคไธช**้ข„่ฎญ็ปƒๆจกๅž‹๏ผš - **Wan2.1-T2V-14B**๏ผš้žๅ› ๆžœ็š„ real-score teacher๏ผŒ็”จไบŽ่ฎก็ฎ—็œŸๅฎžๅˆ†ๅธƒๅพ—ๅˆ†ใ€‚ ```bash -# ๅˆ›ๅปบๆจกๅž‹็›ฎๅฝ• +# Create the model directory mkdir -p models/Diffusion_Transformer -# ไธ‹่ฝฝ Wan2.1 T2V 1.3B๏ผˆstudent ๅŸบ็ก€ๆจกๅž‹๏ผ‰ +# 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 -# ไธ‹่ฝฝ Wan2.1 T2V 14B๏ผˆDMD real-score teacher๏ผ‰ +# 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 ``` @@ -103,12 +104,12 @@ Stage 2 CCD checkpoint๏ผˆ็”Ÿๆˆๅ™จ/ๅˆคๅˆซๅ™จๅˆๅง‹ๅŒ–๏ผ‰้€š่ฟ‡ `--ode_transforme ## ๅ››ใ€ๆ•ฐๆฎๅ‡†ๅค‡ -DMD ไฝฟ็”จ **TextDataset**๏ผˆ`train_mode="normal"`๏ผ‰โ€” ๅช้œ€่ฆๆ็คบ่ฏ๏ผŒไธ้œ€่ฆ่ง†้ข‘ๆ•ฐๆฎ๏ผŒๅ› ไธบ็”Ÿๆˆๅ™จ้€š่ฟ‡่‡ชๅ›žๅฝ’ rollout ่‡ช่กŒๅˆ›ๅปบ่ฎญ็ปƒๆ ทๆœฌใ€‚ไฝ†ไป้œ€ๆไพ›ๅŒ…ๅซๆ็คบ่ฏ็š„ `metadata.json`ใ€‚ +ๅฝ“ `train_mode="normal"`ใ€ๆœชๅฏ็”จ teacher forcing ไธ” GAN/R1/R2 ๆƒ้‡ๅ‡ไธบ้›ถๆ—ถ๏ผŒDMD ไฝฟ็”จ **TextDataset**๏ผšๅช้œ€่ฆๆ็คบ่ฏ๏ผŒ็”Ÿๆˆๅ™จ้€š่ฟ‡่‡ชๅ›žๅฝ’ rollout ่‡ช่กŒๅˆ›ๅปบ่ฎญ็ปƒๆ ทๆœฌใ€‚ไปปไธ€ GAN/R1/R2 ๆƒ้‡ไธบๆญฃๆ—ถไผšๅˆ‡ๆขไธบ **ImageVideoDataset**๏ผŒ้œ€่ฆๅธฆๆ่ฟฐ็š„็œŸๅฎž่ง†้ข‘๏ผŒ่ฏฆ่ง 5.4ใ€‚ ### 4.1 ๅฟซ้€Ÿๆต‹่ฏ•ๆ•ฐๆฎ้›† ```bash -# ไธ‹่ฝฝๅฎ˜ๆ–น็คบไพ‹ๆ•ฐๆฎ้›† +# Download the official demo dataset modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun-Videos-Demo ``` @@ -118,7 +119,7 @@ modelscope download --dataset PAI/X-Fun-Videos-Demo --local_dir ./datasets/X-Fun ### 4.3 metadata.json ๆ ผๅผ -DMD ไป…ไฝฟ็”จๆฏๆก่ฎฐๅฝ•็š„ `"text"` ๅญ—ๆฎตใ€‚ๅฝ“ `--train_mode="normal"`๏ผˆTextDataset ๆจกๅผ๏ผ‰ๆ—ถ๏ผŒ`file_path` ๅ’Œๅ…ถไป–ๅญ—ๆฎตไผš่ขซๅฟฝ็•ฅใ€‚ +ไป…ๆ็คบ่ฏ็š„ TextDataset ๆจกๅผๅช่ฏปๅ–ๆฏๆก่ฎฐๅฝ•็š„ `"text"` ๅญ—ๆฎตใ€‚GAN/R1/R2 ๆˆ– teacher forcing ้œ€่ฆไธŽ Stage 1/2 ็›ธๅŒ็š„่ง†้ข‘ๅ…ƒๆ•ฐๆฎ๏ผŒๅŒ…ๆ‹ฌ `file_path`ใ€`type="video"` ๅ’Œ `text`ใ€‚ ```json [ @@ -131,7 +132,7 @@ DMD ไป…ไฝฟ็”จๆฏๆก่ฎฐๅฝ•็š„ `"text"` ๅญ—ๆฎตใ€‚ๅฝ“ `--train_mode="normal"`๏ผˆTex ] ``` -> **่ฏดๆ˜Ž**๏ผšไฝ ๅฏไปฅๅค็”จไปปไฝ•่ง†้ข‘ metadata.json โ€” DMD ๅชๅ…ณๅฟƒ `text` ๅญ—ๆฎตใ€‚ +> **่ฏดๆ˜Ž**๏ผšๅฏไปฅๅค็”จ็Žฐๆœ‰่ง†้ข‘ๅ…ƒๆ•ฐๆฎใ€‚ไป…ๆ็คบ่ฏๆจกๅผๅช่ฏปๅ– `text`๏ผŒGAN ่ฎญ็ปƒ่ฟ˜ไผš่ฏปๅ–่ง†้ข‘ๆ–‡ไปถใ€‚ --- @@ -179,6 +180,7 @@ accelerate launch --mixed_precision="bf16" --use_fsdp \ --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 \ @@ -212,7 +214,7 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh | `--real_score_pretrained_model_name_or_path` | 14B real-score teacher | `models/Diffusion_Transformer/Wan2.1-T2V-14B` | | `--config_path` | ๆจกๅž‹้…็ฝฎ YAML | `config/wan2.1/wan_civitai.yaml` | | `--train_data_dir` | ๆ•ฐๆฎๆ น็›ฎๅฝ• | `""` | -| `--train_data_meta` | `metadata.json` ่ทฏๅพ„๏ผˆไป…ๅซๆ็คบ่ฏ๏ผ‰ | `datasets/X-Fun-Videos-Demo/metadata.json` | +| `--train_data_meta` | ๅ…ƒๆ•ฐๆฎ่ทฏๅพ„๏ผšๆ™ฎ้€š DMD ไฝฟ็”จๆ็คบ่ฏ๏ผŒGAN/R1/R2 ไฝฟ็”จๅธฆๆ่ฟฐ็š„่ง†้ข‘ | `datasets/X-Fun-Videos-Demo/metadata.json` | | `--ode_transformer_path` | Stage 2 CCD checkpoint๏ผˆ็”Ÿๆˆๅ™จ/ๅˆคๅˆซๅ™จๅˆๅง‹ๅŒ–๏ผ‰ | `$STAGE2_CKPT` | | `--train_batch_size` | ๆฏ GPU batch ๅคงๅฐ | 1 | | `--gradient_accumulation_steps` | ๆขฏๅบฆ็ดฏ็งฏๆญฅๆ•ฐ | 1 | @@ -264,23 +266,60 @@ bash scripts/wan2.1_causal_forcing/train_causal_dmd.sh | `--context_noise` | KV ็ผ“ๅญ˜ๆ›ดๆ–ฐ็š„ไธŠไธ‹ๆ–‡ๅ™ชๅฃฐ็บงๅˆซ | 0 | | `--use_teacher_forcing` | ๅฏ็”จ teacher forcing๏ผˆๅฐ† clean_x ไผ ็ป™ transformer๏ผ‰ | - | | `--teacher_forcing_prob` | ๆฏๆญฅๅบ”็”จ teacher forcing ็š„ๆฆ‚็Ž‡๏ผˆ1.0 = ๅง‹็ปˆ๏ผ‰ | 1.0 | -| `--train_mode` | ่ฎญ็ปƒๆจกๅผ๏ผš`normal`๏ผˆTextDataset๏ผŒไป…ๆ็คบ่ฏ๏ผ‰ๆˆ– `i2v` | `normal` | +| `--train_mode` | `normal`๏ผˆๆœชๅฏ็”จ teacher forcing ๆˆ– GAN/R1/R2 ๆ—ถไป…ๆ็คบ่ฏ๏ผ‰ๆˆ– `i2v` | `normal` | | `--resume_from_checkpoint` | ๆขๅค่ฎญ็ปƒใ€‚ไฝฟ็”จ `"latest"` ่‡ชๅŠจ้€‰ๆ‹ฉๆœ€ๆ–ฐ checkpoint | `"latest"` | --- +### 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 ๆŸๅคฑใ€‚ + +ๅˆคๅˆซๅ™จๅ…ฑไบซ **fake-score backbone**๏ผŒๆฏไธช้€‰ๅฎš็‰นๅพๅฑ‚ๅฏนๅบ”ไธ€ไธชๅฏๅญฆไน  register tokenใ€‚ๅˆคๅˆซๅ™จๅคด็”ฑ critic optimizer ๆ›ดๆ–ฐ๏ผŒไธๅŠ ๅ…ฅ generator optimizerใ€‚็”Ÿๆˆๅ™จๆŸๅคฑไธบ `DMD + ๅŠ ๆƒ GAN`๏ผŒcritic ๆŸๅคฑไธบ `denoising + ๅŠ ๆƒ GAN + R1 + R2`ใ€‚ + +| ๅ‚ๆ•ฐ | ่ฏดๆ˜Ž | ้ป˜่ฎคๅ€ผ | +|------|------|--------| +| `--gan_g_weight` | ้ž้ฅฑๅ’Œ็”Ÿๆˆๅ™จ GAN ๆŸๅคฑๆƒ้‡ | `0.0` | +| `--gan_d_weight` | ๅˆคๅˆซๅ™จ GAN ๆŸๅคฑๆƒ้‡ | `0.0` | +| `--gan_feature_layers` | ๅˆคๅˆซๅ™จๅคดไฝฟ็”จ็š„ fake-score block ็ดขๅผ• | `21 29` | +| `--gan_block_ffn_dim` | ๅˆคๅˆซๅ™จไบคๅ‰ๆณจๆ„ๅŠ› FFN ็ปดๅบฆ | `2048` | +| `--gan_block_num_heads` | ๅˆคๅˆซๅ™จๅคดๆ•ฐ๏ผŒ้กปๆ•ด้™ค fake-score ้š่—็ปดๅบฆ | `12` | +| `--relativistic_discriminator` | ไฝฟ็”จ้…ๅฏน็œŸๅฎž/็”Ÿๆˆๆ ทๆœฌ logit ๅทฎๅ€ผ | ๅ…ณ้—ญ | +| `--r1_weight` | ๅ™ชๅฃฐ็œŸๅฎž latent ไธŠ็š„ APT ๅฝขๅผๆœ‰้™ๅทฎๅˆ†ๆญฃๅˆ™ๆƒ้‡ | `0.0` | +| `--r1_sigma` | R1 ๆ‰ฐๅŠจๆ ‡ๅ‡†ๅทฎ | `0.01` | +| `--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` | ๅ…ณ้—ญ | + +่ฎพ 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))`ใ€‚ + +R1 ไฟๆŒๆบๅฎž็Žฐ็š„ APT ๅฝขๅผ๏ผš`r1_weight * mean((D_real_perturbed - D_real)^2)`๏ผŒ**ไธ้™คไปฅ sigma**ใ€‚R2 ไธบ `0.5 * r2_weight * mean(((D_fake_perturbed - D_fake) / r2_sigma)^2)`ใ€‚ไบŒ่€…้ƒฝๆ˜ฏๆœ‰้™ๅทฎๅˆ†ๆญฃๅˆ™๏ผŒไธๆ˜ฏ้œ€่ฆ double backward ็š„่พ“ๅ…ฅๆขฏๅบฆๆƒฉ็ฝšใ€‚็œŸๅฎž/็”Ÿๆˆๆ ทๆœฌๅ…ฑ็”จๆ—ถ้—ดๆญฅๅ’Œๅ™ชๅฃฐ๏ผŒๆ‰ฐๅŠจๅ™ชๅฃฐ็‹ฌ็ซ‹ๆŠฝๆ ทใ€‚ + +ๅผ€ๅฏ GAN ๆ—ถ๏ผŒๅ‘ `train_causal_dmd.py` ไผ ๅ…ฅๆญฃ็š„็”Ÿๆˆๅ™จ/ๅˆคๅˆซๅ™จๆƒ้‡๏ผŒๅนถๆŒ‰้œ€้…็ฝฎ R1/R2ใ€‚่ฆๆฑ‚๏ผš +- ไฝฟ็”จๅธฆๆ่ฟฐ็š„็œŸๅฎž่ง†้ข‘๏ผŒไธ่ƒฝๅชๆไพ›ๆ็คบ่ฏ JSONใ€‚ๅ…ƒๆ•ฐๆฎ้‡‡็”จ Stage 1/2 ่ง†้ข‘ๆ ผๅผ๏ผŒๅธงๆ•ฐ้กปๆปก่ถณ `score_num_frames` ไธŽๅ› ๆžœๅ—ๅ‡ ไฝ•ใ€‚ +- `--gradient_accumulation_steps=1`๏ผ›GAN ๅˆ†ๆ”ฏไธๆ”ฏๆŒ DeepSpeedใ€‚ +- ็‰นๅพๅฑ‚็ดขๅผ•ๅฟ…้กปๅญ˜ๅœจใ€‚้ป˜่ฎค `21 29` ไธŽ `12` ไธชๅคด้€‚็”จไบŽ Wan2.1-T2V-1.3B๏ผŒๅ…ถไป– backbone ๅฏ่ƒฝ้œ€่ฆ่ฐƒๆ•ดใ€‚ +- ไป…ๅผ€ๅฏ `gan_g_weight` ไธไผšๆไพ›ๅˆคๅˆซๅ™จ็š„ๅˆ†็ฑปๅญฆไน ็›ฎๆ ‡ใ€‚้กปๆ˜Ž็กฎ้€‰ๆ‹ฉๅˆคๅˆซๅ™จๆƒ้‡๏ผŒ้ž้›ถๆƒ้‡ไธไฟ่ฏ่ดจ้‡ๆๅ‡ใ€‚ + +ๅฏ่ฎญ็ปƒ fake-score ไธŽๅˆคๅˆซๅ™จๅคด็š„ master weights ไฟๆŒ **FP32**๏ผŒๅ‰ๅ‘ๆŒ‰้…็ฝฎไฝฟ็”จ autocast/FSDP ๆททๅˆ็ฒพๅบฆใ€‚Logitsใ€softplus ๆŸๅคฑๅŠๆœ‰้™ๅทฎๅˆ†ๅฝ’็บฆไฝฟ็”จ FP32ใ€‚`--fsdp_reduce_fp32` ๅชๆ”นๅ˜ๆขฏๅบฆ้€šไฟก็ฒพๅบฆ๏ผŒไธๆ”นๅ˜ๅ‰ๅ‘ dtype ๆˆ–ๆจกๅž‹็ป“ๆž„ใ€‚ๆ—ถ้—ดๆญฅ clamp ๅชไฝœ็”จไบŽ score ่ฎญ็ปƒ๏ผŒไธๆ”นๅ˜็”Ÿๆˆๅ™จ rollout ๆˆ– validation scheduleใ€‚ + +ๆขๅค GAN checkpoint ๆ—ถๅบ”ไฝฟ็”จ็›ธๅŒ็š„ GAN/ๅˆคๅˆซๅ™จๅคด้…็ฝฎใ€‚ๅˆคๅˆซๅ™จๅคด้š fake-score ไฟๅญ˜๏ผŒๅŒ…ๆ‹ฌ FP32 ๅˆคๅˆซๅ™จๅคดๅผ ้‡๏ผŒๅนถๅœจๅŠ ่ฝฝๅ‰้‡ๅปบใ€‚ๅฏผๅ‡บ็š„ generator ไธๅŒ…ๅซ่ฏฅๅคด๏ผŒๆŽจ็†ๆ— ้œ€ๅˆคๅˆซๅ™จๆˆ–ๆ–ฐๅขž GAN ๅ‚ๆ•ฐใ€‚ + +--- + ## ๅ…ญใ€ไฝฟ็”จ่ฎญ็ปƒๅฅฝ็š„ Checkpoint Stage 3 DMD checkpoint ๆ˜ฏ Causal-Forcing ๆตๆฐด็บฟ็š„**ๆœ€็ปˆๆจกๅž‹**๏ผŒ็›ดๆŽฅ็”จไบŽๆŽจ็†๏ผš ```python -# ๅœจ examples/wan2.1_causal_forcing/predict_t2v.py ไธญ +# In examples/wan2.1_causal_forcing/predict_t2v.py transformer_path = "output_dir_wan2.1_causal_forcing_dmd/checkpoint-{N}/diffusion_pytorch_model.safetensors" -# DMD Stage 3 ๆŽจ็†้…็ฝฎ -guidance_scale = 1.0 # CFG ๅทฒ็ƒ˜็„ฆๅˆฐ่’ธ้ฆๆƒ้‡ไธญ -num_inference_steps = 4 # 4 ๆญฅ DMD +# DMD Stage 3 inference config +guidance_scale = 1.0 # CFG is baked into distilled weights +num_inference_steps = 4 # 4-step DMD stochastic_sampling = True -num_frame_per_block = 3 # ๅˆ†ๅ—็”Ÿๆˆ +num_frame_per_block = 3 # Chunk-wise generation ``` ๆˆ–่ฟ่กŒ๏ผš @@ -299,4 +338,5 @@ python examples/wan2.1_causal_forcing/predict_t2v.py ## ไธƒใ€ๆ›ดๅคš่ต„ๆบ - **Causal-Forcing ่ฎบๆ–‡**๏ผšhttps://github.com/thu-ml/Causal-Forcing +- **One-Forcing๏ผˆๅฏ้€‰ GAN ๆœบๅˆถ๏ผ‰**๏ผšhttps://github.com/Aurora-edu/One-Forcing - **ๅฎ˜ๆ–น GitHub**๏ผšhttps://github.com/aigc-apps/VideoX-Fun diff --git a/scripts/wan2.1_causal_forcing/train_ar_diffusion.py b/scripts/wan2.1_causal_forcing/train_ar_diffusion.py index 965b0c9..b6f2087 100644 --- a/scripts/wan2.1_causal_forcing/train_ar_diffusion.py +++ b/scripts/wan2.1_causal_forcing/train_ar_diffusion.py @@ -18,6 +18,7 @@ scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing # See the License for the specific language governing permissions and import argparse +import dataclasses import gc import logging import math @@ -314,6 +315,10 @@ def parse_args(): 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( + "--fsdp_reduce_fp32", action="store_true", + help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.", + ) parser.add_argument( "--logging_dir", type=str, @@ -596,6 +601,9 @@ def main(): 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: @@ -1254,6 +1262,8 @@ def main(): training_weights_table = training_weights_table.to(accelerator.device) for epoch in range(first_epoch, args.num_train_epochs): + if global_step >= args.max_train_steps: + break train_loss = 0.0 batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) for step, batch in enumerate(train_dataloader): @@ -1487,6 +1497,9 @@ def main(): 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(): log_validation( diff --git a/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py b/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py index d53e74c..6ceb0af 100644 --- a/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py +++ b/scripts/wan2.1_causal_forcing/train_causal_consistency_distill.py @@ -18,6 +18,7 @@ scripts/wan2.1_self_forcing/train_distill.py and Causal-Forcing # See the License for the specific language governing permissions and import argparse +import dataclasses import gc import logging import math @@ -317,6 +318,10 @@ def parse_args(): 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( + "--fsdp_reduce_fp32", action="store_true", + help="Reduce gradients in fp32 under FSDP mixed precision; parameter gathering stays in low precision.", + ) parser.add_argument( "--logging_dir", type=str, @@ -610,6 +615,9 @@ def main(): 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: @@ -769,7 +777,7 @@ def main(): 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=weight_dtype, + torch_dtype=torch.float32, ) ema_transformer3d.num_frame_per_block = args.num_frame_per_block ema_transformer3d.independent_first_frame = args.independent_first_frame @@ -1437,6 +1445,8 @@ def main(): torch.cuda.empty_cache() for epoch in range(first_epoch, args.num_train_epochs): + if global_step >= args.max_train_steps: + break train_loss = 0.0 batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) for step, batch in enumerate(train_dataloader): @@ -1737,6 +1747,9 @@ def main(): 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(): log_validation( diff --git a/scripts/wan2.1_causal_forcing/train_causal_dmd.py b/scripts/wan2.1_causal_forcing/train_causal_dmd.py index 8b80703..e9cdb8b 100644 --- a/scripts/wan2.1_causal_forcing/train_causal_dmd.py +++ b/scripts/wan2.1_causal_forcing/train_causal_dmd.py @@ -17,6 +17,7 @@ import argparse import contextlib +import dataclasses import gc import json import logging @@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, get_random_mask) from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel, WanTransformer3DModel, - WanTransformer3DModel_SelfForcing) + 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, @@ -158,6 +161,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi 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__) @@ -494,6 +506,14 @@ def parse_args(): ), ) parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.") + 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( "--dataloader_num_workers", type=int, @@ -838,6 +858,47 @@ def parse_args(): 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: @@ -848,6 +909,14 @@ def parse_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( @@ -875,6 +944,8 @@ def main(): 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 @@ -896,6 +967,9 @@ def main(): 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: @@ -1111,28 +1185,55 @@ def main(): 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.") - ema_generator_transformer3d = EMAModel( - generator_transformer3d.parameters(), - model_cls=WanTransformer3DModel_SelfForcing, - model_config=generator_transformer3d.config, - ) + 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, + ) + 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_stage != 0 or zero_stage == 3: + 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=weight_dtype) for k, v in accelerate_state_dict.items()} + 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: @@ -1149,10 +1250,14 @@ def main(): # 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: - if args.use_ema and models[0] is generator_transformer3d: - ema_generator_transformer3d.save_pretrained(os.path.join(output_dir, "transformer_ema")) - - models[0].save_pretrained(os.path.join(output_dir, "transformer")) + 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() @@ -1160,34 +1265,23 @@ def main(): pickle.dump([batch_sampler.sampler._pos_start, first_epoch], file) def load_model_hook(models, input_dir): - if args.use_ema and models[0] is generator_transformer3d: - ema_path = os.path.join(input_dir, "transformer_ema") - _, 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, - ) - 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 - for i in range(len(models)): # pop models so that they are not loaded again model = models.pop() - # load diffusers style into model - 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 + # 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): @@ -1201,6 +1295,32 @@ def main(): 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() @@ -1267,8 +1387,8 @@ def main(): fake_trainable_params = list(filter(lambda p: p.requires_grad, fake_score_transformer3d.parameters())) fake_trainable_params_optim = [ - {'params': [], 'lr': args.learning_rate}, - {'params': [], 'lr': args.learning_rate / 2}, + {'params': [], 'lr': args.learning_rate_critic}, + {'params': [], 'lr': args.learning_rate_critic / 2}, ] in_already = [] for name, param in fake_score_transformer3d.named_parameters(): @@ -1281,7 +1401,7 @@ def main(): 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}") + print(f"Set {name} to lr : {args.learning_rate_critic}") break if high_lr_flag: continue @@ -1290,9 +1410,15 @@ def main(): 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 / 2}") + 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, @@ -1334,7 +1460,7 @@ def main(): args.random_hw_adapt = False # Get the dataset - if args.train_mode != "normal" or args.use_teacher_forcing: + 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, @@ -1362,7 +1488,7 @@ def main(): return length_to_frame_num - if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing: + 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( @@ -1672,7 +1798,7 @@ def main(): ) fake_score_lr_scheduler = get_scheduler( args.lr_scheduler, - optimizer=optimizer, + optimizer=critic_optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, ) @@ -1798,6 +1924,8 @@ def main(): 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, @@ -1852,12 +1980,16 @@ def main(): return result 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 @@ -1878,7 +2010,7 @@ def main(): 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: + 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) @@ -1942,8 +2074,9 @@ def main(): temp_n_frames = (temp_n_frames - 1) // sample_n_frames_bucket_interval + 1 pixel_values = pixel_values[:, :temp_n_frames, :, :] - mask_pixel_values = mask_pixel_values[:, :temp_n_frames, :, :] - mask = mask[:, :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: @@ -1966,8 +2099,9 @@ def main(): actual_video_length = (actual_video_length - 1) // sample_n_frames_bucket_interval + 1 pixel_values = pixel_values[:, :actual_video_length, :, :] - mask_pixel_values = mask_pixel_values[:, :actual_video_length, :, :] - mask = mask[:, :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() @@ -1990,7 +2124,7 @@ def main(): 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: + if args.use_teacher_forcing or use_gan: clean_latents = _batch_encode_vae(pixel_values) else: clean_latents = None @@ -2022,7 +2156,7 @@ def main(): clip_context.append(_clip_context) clip_context = torch.cat(clip_context) - if args.use_teacher_forcing: + if args.use_teacher_forcing or use_gan: target_shape = clean_latents.size() else: target_shape = mask_latents.size() @@ -2426,17 +2560,7 @@ def main(): # Final output โ€” slice generated frames (may be < num_frames for variable-length) generator_pred_full = output_pred[:, :, :num_generated_frames] - # Gradient mask: first block gets no gradient when generating > min frames - # (matches Self-Forcing reference: model/base.py L182-L190) - min_num_frames_score = args.score_num_frames - need_gradient_mask = (num_generated_frames != min_num_frames_score) - gradient_mask = None - if need_gradient_mask: - gradient_mask = torch.ones_like(generator_pred_full, dtype=torch.bool) - if args.independent_first_frame: - gradient_mask[:, :, :1] = False - else: - gradient_mask[:, :, :args.num_frame_per_block] = False + 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: @@ -2450,11 +2574,13 @@ def main(): generator_pred_for_score = generator_pred_full score_num_frames = num_generated_frames - # Compute score_mask for DMD loss (matches Self-Forcing: dmd.py L199-204) score_mask = None - if gradient_mask is not None: - mask_offset = num_generated_frames - score_num_frames - score_mask = gradient_mask[:, :, mask_offset:mask_offset + score_num_frames] + 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 @@ -2530,7 +2656,7 @@ def main(): 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 (ไธ€ๆฌกๆ€งๅค„็†ๆ•ดไธช่ง†้ข‘) + # Use block_mask for causal training generator_pred = generator_transformer3d( x=generator_noise_list, context=prompt_embeds, @@ -2563,6 +2689,8 @@ def main(): # 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), @@ -2668,17 +2796,53 @@ def main(): 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") - fake_score_transformer3d = fake_score_transformer3d.to("cpu") + if args.gan_g_weight == 0: + fake_score_transformer3d = fake_score_transformer3d.to("cpu") torch.cuda.empty_cache() - accelerator.backward(dmd_loss) + 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) @@ -2967,6 +3131,8 @@ def main(): 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( @@ -2975,15 +3141,25 @@ def main(): critic_timestep ) - with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): - 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, - ) + # 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() @@ -2997,11 +3173,136 @@ def main(): final_loss = masked_loss.mean() return final_loss - 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 - - accelerator_fake_score_transformer3d.backward(denoising_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() @@ -3019,10 +3320,15 @@ def main(): ema_generator_transformer3d.step(generator_transformer3d.parameters()) progress_bar.update(1) global_step += 1 - accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step) + 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: @@ -3092,6 +3398,9 @@ def main(): 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(): if args.use_ema: @@ -3121,12 +3430,6 @@ def main(): # Create the pipeline using the trained modules and save it. accelerator.wait_for_everyone() - if args.use_ema: - # Copy EMA weights to generator on ALL processes before save_state (collective op). - ema_generator_transformer3d.copy_to(generator_transformer3d.parameters()) - if accelerator.is_main_process: - generator_transformer3d = unwrap_model(generator_transformer3d) - if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process: gc.collect() torch.cuda.empty_cache() @@ -3135,8 +3438,6 @@ def main(): 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) - if args.use_ema and accelerator.is_main_process: - ema_generator_transformer3d.save_pretrained(os.path.join(save_path, "transformer_ema")) logger.info(f"Saved state to {save_path}") accelerator.end_training() diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md index d434b76..fd71d6e 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_DISTILL.md @@ -258,6 +258,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--adam_epsilon` | AdamW epsilon | 1e-10 | | `--vae_mini_batch` | VAE encoding mini-batch size | 1 | | `--max_grad_norm` | Gradient clipping threshold | 0.05 | +| `--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 | - | | `--training_with_video_token_length` | Train based on token length, supports arbitrary resolutions | - | @@ -275,12 +276,17 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | Parameter | Description | Example Value | |-----|------|-------| | `--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 | +| `--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 | | `--negative_prompt` | Negative prompt for distillation | Chinese negative prompt | | `--train_sampling_steps` | Training sampling steps | 1000 | | `--ode_transformer_path` | Path to ODE-trained weights to load into generator transformer3d | `models/Diffusion_Transformer/Self-Forcing/checkpoints/ode_init.pt` | +| `--real_score_model_name_or_path` | Optional: separate model used as real score (defaults to `pretrained_model_name_or_path`) | `models/Diffusion_Transformer/Wan2.1-T2V-14B` | +| `--score_num_frames` | Latent frames in the score-computation window; `fake_score`/`real_score` always receive this many frames | 21 | +| `--min_length_prob_bias` | Probability bias toward sampling the minimum length (`score_num_frames`); `0.0` = uniform sampling | 0.0 | **Self-Forcing-Specific Parameters**: @@ -294,6 +300,23 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--use_teacher_forcing` | Enable teacher forcing training (pass clean_x to transformer) | - | | `--teacher_forcing_prob` | Probability of applying teacher forcing per step | 1.0 | +**GAN / Discriminator Parameters (One-Forcing)**: + +> All default to `0` / off, so a plain DMD run is unaffected. Setting **any** positive `--gan_g_weight` / `--gan_d_weight` / `--r1_weight` / `--r2_weight` attaches a discriminator head to `fake_score` and **requires real videos** in the training data. + +| Parameter | Description | Example Value | +|-----|------|-------| +| `--gan_g_weight` | One-Forcing: weight of the non-saturating GAN loss added to the DMD generator loss (`0` disables) | 0.0 | +| `--gan_d_weight` | One-Forcing: weight of the GAN discriminator loss added to the `fake_score` denoising loss (`0` disables) | 0.0 | +| `--gan_feature_layers` | `fake_score` block indices whose outputs feed the discriminator head (one register token per layer) | `21 29` | +| `--gan_block_ffn_dim` | FFN dim of the discriminator cross-attention blocks | 2048 | +| `--gan_block_num_heads` | Number of heads of the discriminator cross-attention blocks | 12 | +| `--relativistic_discriminator` | Use the relativistic GAN loss (fake vs. real logit difference) | - | +| `--r1_weight` | Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma) | 0.0 | +| `--r1_sigma` | Perturbation std of the R1 penalty | 0.01 | +| `--r2_weight` | Finite-difference R2 penalty on noisy generated latents | 0.0 | +| `--r2_sigma` | Perturbation std of the R2 penalty | 0.01 | + **Sample Size Configuration Guide**: - `video_sample_size` represents the resolution size of videos; when `random_hw_adapt` is True, it represents the minimum value between video and image resolutions. - `image_sample_size` represents the resolution size of images; when `random_hw_adapt` is True, it represents the maximum value between video and image resolutions. 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 305a831..4224fb6 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 @@ -259,6 +259,7 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--adam_epsilon` | AdamW epsilon ๅ€ผ | 1e-10 | | `--vae_mini_batch` | VAE ็ผ–็ ๆ—ถ็š„่ฟทไฝ ๆ‰นๆฌกๅคงๅฐ | 1 | | `--max_grad_norm` | ๆขฏๅบฆ่ฃๅ‰ช้˜ˆๅ€ผ | 0.05 | +| `--fsdp_reduce_fp32` | ๅฏ้€‰๏ผšFSDP ๆทท็ฒพไธ‹ๆขฏๅบฆไปฅ fp32 ๅš่ทจๅก reduce๏ผˆ้ป˜่ฎค bf16๏ผ‰ | - | | `--enable_bucket` | ๅฏ็”จๅˆ†ๆกถ่ฎญ็ปƒ๏ผŒไธ่ฃๅ‰ชๅ›พ็‰‡/่ง†้ข‘๏ผŒๆŒ‰ๅˆ†่พจ็Ž‡ๅˆ†็ป„่ฎญ็ปƒ | - | | `--random_hw_adapt` | ่‡ชๅŠจ็ผฉๆ”พๅ›พ็‰‡/่ง†้ข‘ๅˆฐ `[min_size, max_size]` ่Œƒๅ›ดๅ†…็š„้šๆœบๅฐบๅฏธ | - | | `--training_with_video_token_length` | ๆ นๆฎ token ้•ฟๅบฆ่ฎญ็ปƒ๏ผŒๆ”ฏๆŒไปปๆ„ๅˆ†่พจ็Ž‡ | - | @@ -276,12 +277,17 @@ 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 | +| `--score_timestep_clamp` | ๅฏ้€‰๏ผšๅฐ† DMD/critic ็š„ timestep ๆˆชๆ–ญๅˆฐ `[0.02, 0.98] * num_train_timesteps` | - | | `--fake_guidance_scale` | ็”จไบŽ่ฏ„ๅˆ†็š„่™šๆ‹Ÿ guidance scale | 0.0 | | `--gen_update_interval` | ็”Ÿๆˆๅ™จๆ›ดๆ–ฐ้—ด้š” | 5 | | `--negative_prompt` | ็”จไบŽ่’ธ้ฆ็š„่ดŸๅ‘ๆ็คบ่ฏ | ไธญๆ–‡่ดŸๅ‘ๆ็คบ่ฏ | | `--train_sampling_steps` | ่ฎญ็ปƒ้‡‡ๆ ทๆญฅๆ•ฐ | 1000 | | `--ode_transformer_path` | ODE ่ฎญ็ปƒๆƒ้‡่ทฏๅพ„๏ผŒๅŠ ่ฝฝๅˆฐ generator transformer3d ไธญ | `models/Diffusion_Transformer/Self-Forcing/checkpoints/ode_init.pt` | +| `--real_score_model_name_or_path` | ๅฏ้€‰๏ผšๅ•็‹ฌๆŒ‡ๅฎš real score ๆจกๅž‹๏ผˆ้ป˜่ฎคไธŽ `pretrained_model_name_or_path` ็›ธๅŒ๏ผ‰ | `models/Diffusion_Transformer/Wan2.1-T2V-14B` | +| `--score_num_frames` | ่ฏ„ๅˆ†็ช—ๅฃ็š„ latent ๅธงๆ•ฐ๏ผ›`fake_score`/`real_score` ๅง‹็ปˆๆŽฅๆ”ถ่ฟ™ไนˆๅคšๅธง | 21 | +| `--min_length_prob_bias` | ้‡‡ๆ ทๆœ€ๅฐ้•ฟๅบฆ๏ผˆ`score_num_frames`๏ผ‰็š„ๆฆ‚็Ž‡ๅ็ฝฎ๏ผ›`0.0` ่กจ็คบๅ‡ๅŒ€้‡‡ๆ ท | 0.0 | **Self-Forcing ็‰นๆœ‰ๅ‚ๆ•ฐ**๏ผš @@ -295,6 +301,23 @@ accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_con | `--use_teacher_forcing` | ๅฏ็”จ teacher forcing ่ฎญ็ปƒ๏ผˆๅฐ† clean_x ไผ ็ป™ transformer๏ผ‰ | - | | `--teacher_forcing_prob` | ๆฏๆญฅๅบ”็”จ teacher forcing ็š„ๆฆ‚็Ž‡ | 1.0 | +**GAN / ๅˆคๅˆซๅ™จๅ‚ๆ•ฐ๏ผˆOne-Forcing๏ผ‰**๏ผš + +> ๅ…จ้ƒจ้ป˜่ฎคไธบ `0` / ๅ…ณ้—ญ๏ผŒๅ› ๆญคๆ™ฎ้€š DMD ่ฎญ็ปƒไธๅ—ๅฝฑๅ“ใ€‚ๅช่ฆ `--gan_g_weight` / `--gan_d_weight` / `--r1_weight` / `--r2_weight` ไธญ**ไปปไธ€**ไธบๆญฃๅ€ผ๏ผŒๅฐฑไผšๅœจ `fake_score` ไธŠๆŒ‚ไธ€ไธชๅˆคๅˆซๅ™จๅคด๏ผŒๅนถไธ”่ฎญ็ปƒๆ•ฐๆฎ**ๅฟ…้กปๅŒ…ๅซ็œŸๅฎž่ง†้ข‘**ใ€‚ + +| ๅ‚ๆ•ฐ | ่ฏดๆ˜Ž | ็คบไพ‹ๅ€ผ | +|-----|------|-------| +| `--gan_g_weight` | One-Forcing๏ผšๅŠ ๅˆฐ DMD ็”Ÿๆˆๅ™จๆŸๅคฑไธŠ็š„้ž้ฅฑๅ’Œ GAN ๆŸๅคฑๆƒ้‡๏ผˆ`0` ๅ…ณ้—ญ๏ผ‰ | 0.0 | +| `--gan_d_weight` | One-Forcing๏ผšๅŠ ๅˆฐ `fake_score` ๅŽปๅ™ชๆŸๅคฑไธŠ็š„ GAN ๅˆคๅˆซๅ™จๆŸๅคฑๆƒ้‡๏ผˆ`0` ๅ…ณ้—ญ๏ผ‰ | 0.0 | +| `--gan_feature_layers` | ่พ“ๅ‡บๅ–‚็ป™ๅˆคๅˆซๅ™จๅคด็š„ `fake_score` block ็ดขๅผ•๏ผˆๆฏๅฑ‚ไธ€ไธช register token๏ผ‰ | `21 29` | +| `--gan_block_ffn_dim` | ๅˆคๅˆซๅ™จ cross-attention block ็š„ FFN ็ปดๅบฆ | 2048 | +| `--gan_block_num_heads` | ๅˆคๅˆซๅ™จ cross-attention block ็š„ๅคดๆ•ฐ | 12 | +| `--relativistic_discriminator` | ไฝฟ็”จ relativistic GAN ๆŸๅคฑ๏ผˆfake ไธŽ real logit ไน‹ๅทฎ๏ผ‰ | - | +| `--r1_weight` | ๅฏนๅซๅ™ช็œŸๅฎž latent ็š„ๆœ‰้™ๅทฎๅˆ† R1 ๆƒฉ็ฝš๏ผˆAPT ๅฝขๅผ๏ผŒไธ้™คไปฅ sigma๏ผ‰ | 0.0 | +| `--r1_sigma` | R1 ๆƒฉ็ฝš็š„ๆ‰ฐๅŠจๆ ‡ๅ‡†ๅทฎ | 0.01 | +| `--r2_weight` | ๅฏนๅซๅ™ช็”Ÿๆˆ latent ็š„ๆœ‰้™ๅทฎๅˆ† R2 ๆƒฉ็ฝš | 0.0 | +| `--r2_sigma` | R2 ๆƒฉ็ฝš็š„ๆ‰ฐๅŠจๆ ‡ๅ‡†ๅทฎ | 0.01 | + **Sample Size ้…็ฝฎๆŒ‡ๅ—**๏ผš - `video_sample_size` ่กจ็คบ่ง†้ข‘็š„ๅˆ†่พจ็Ž‡ๅคงๅฐ๏ผ›ๅฝ“ `random_hw_adapt` ไธบ True ๆ—ถ๏ผŒ่กจ็คบ่ง†้ข‘ๅ’Œๅ›พๅƒๅˆ†่พจ็Ž‡็š„ๆœ€ๅฐๅ€ผใ€‚ - `image_sample_size` ่กจ็คบๅ›พๅƒ็š„ๅˆ†่พจ็Ž‡ๅคงๅฐ๏ผ›ๅฝ“ `random_hw_adapt` ไธบ True ๆ—ถ๏ผŒ่กจ็คบ่ง†้ข‘ๅ’Œๅ›พๅƒๅˆ†่พจ็Ž‡็š„ๆœ€ๅคงๅ€ผใ€‚ diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_ODE.md b/scripts/wan2.1_self_forcing/README_TRAIN_ODE.md index 800b6b4..86df740 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_ODE.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_ODE.md @@ -277,6 +277,9 @@ bash scripts/wan2.1_self_forcing/train_ode.sh | `--max_grad_norm` | Gradient clipping threshold | 0.05 | | `--trainable_modules` | Trainable modules (`"."` = all) | `"."` | | `--resume_from_checkpoint` | Resume path or `"latest"` | `latest` | +| `--save_state` | Save the full accelerator state on checkpoint (required to resume training) | - | +| `--use_kv_cache_training` | Run block-by-block KV cache training that matches `pipeline_wan_self_forcing` inference; otherwise fall back to one-shot causal-mask ODE regression (baseline) | - | +| `--prob_full_zero_start` | Per-sample probability of forcing all frames in all blocks to timestep index `0` (pure-noise start) to bridge the train-inference gap; `0.0` disables | 0.0 | **ODE-specific parameters** (must match Step 1 unless you understand the consequences): diff --git a/scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md b/scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md index b1310c0..cfa1c92 100755 --- a/scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md +++ b/scripts/wan2.1_self_forcing/README_TRAIN_ODE_zh-CN.md @@ -277,6 +277,9 @@ bash scripts/wan2.1_self_forcing/train_ode.sh | `--max_grad_norm` | ๆขฏๅบฆ่ฃๅ‰ช้˜ˆๅ€ผ | 0.05 | | `--trainable_modules` | ๅฏ่ฎญ็ปƒๆจกๅ—๏ผˆ`"."` ่กจ็คบๅ…จ้‡๏ผ‰ | `"."` | | `--resume_from_checkpoint` | ๆขๅค่ฎญ็ปƒ่ทฏๅพ„ๆˆ– `"latest"` | `latest` | +| `--save_state` | ๅœจ checkpoint ไธญไฟๅญ˜ๅฎŒๆ•ด็š„ accelerator ็Šถๆ€๏ผˆๆขๅค่ฎญ็ปƒๆ‰€้œ€๏ผ‰ | - | +| `--use_kv_cache_training` | ่ฟ่กŒ้€ๅ— KV ็ผ“ๅญ˜่ฎญ็ปƒ๏ผŒๅฎŒๅ…จๅŒน้… `pipeline_wan_self_forcing` ๆŽจ็†่กŒไธบ๏ผ›ๅฆๅˆ™ๅ›ž้€€ๅˆฐไธ€ๆฌกๆ€ง็š„ causal-mask ODE ๅ›žๅฝ’๏ผˆbaseline๏ผ‰ | - | +| `--prob_full_zero_start` | ๆฏไธชๆ ทๆœฌๅผบๅˆถๆ‰€ๆœ‰ๅ—็š„ๆ‰€ๆœ‰ๅธงไฝฟ็”จ timestep ็ดขๅผ• `0`๏ผˆ็บฏๅ™ชๅฃฐ่ตท็‚น๏ผ‰็š„ๆฆ‚็Ž‡๏ผŒ็”จไบŽๅผฅๅˆ่ฎญ็ปƒ-ๆŽจ็†ๅทฎ่ท๏ผ›`0.0` ่กจ็คบๅ…ณ้—ญ | 0.0 | **ODE ็‰นๆœ‰ๅ‚ๆ•ฐ**๏ผˆ้™ค้žๆธ…ๆฅšๅŽๆžœ๏ผŒๅฆๅˆ™้œ€ไธŽ็ฌฌไธ€ๆญฅไฟๆŒไธ€่‡ด๏ผ‰๏ผš diff --git a/scripts/wan2.1_self_forcing/train_distill.py b/scripts/wan2.1_self_forcing/train_distill.py index 4f148e7..141cfa9 100644 --- a/scripts/wan2.1_self_forcing/train_distill.py +++ b/scripts/wan2.1_self_forcing/train_distill.py @@ -17,6 +17,7 @@ import argparse import contextlib +import dataclasses import gc import json import logging @@ -74,10 +75,12 @@ from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, get_random_mask) from videox_fun.models import (AutoencoderKLWan, CLIPModel, WanT5EncoderModel, WanTransformer3DModel, - WanTransformer3DModel_SelfForcing) + 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, @@ -165,6 +168,15 @@ def slice_for_score(pred, vae, weight_dtype, score_num_frames=21, independent_fi 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__) @@ -492,6 +504,8 @@ def parse_args(): parser.add_argument( "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." ) + parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model for the generator.") + parser.add_argument("--ema_decay", type=float, default=0.9999, help="EMA decay applied per generator update (Causal-Forcing uses 0.99).") parser.add_argument( "--allow_tf32", action="store_true", @@ -523,6 +537,9 @@ def parse_args(): 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( + "--fsdp_reduce_fp32", action="store_true", help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision)." + ) 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( @@ -722,6 +739,12 @@ def parse_args(): default=None, help=("If you want to load the ode-trained weight into generator transformer3d, input its path."), ) + parser.add_argument( + "--real_score_model_name_or_path", + type=str, + default=None, + help=("Path to a separate model (e.g. Wan2.1-T2V-14B) used as real score. Defaults to pretrained_model_name_or_path."), + ) parser.add_argument( "--vae_path", type=str, @@ -782,6 +805,11 @@ def parse_args(): default=4.5, help="The cfg scale for real score.", ) + parser.add_argument( + "--score_timestep_clamp", + action="store_true", + help="Clamp DMD and critic timesteps to [0.02, 0.98] * num_train_timesteps.", + ) parser.add_argument( '--denoising_step_indices_list', nargs='+', @@ -842,6 +870,67 @@ def parse_args(): 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: weight of the non-saturating GAN loss added to the DMD generator loss (0 disables). " + "Any positive GAN/R1/R2 weight attaches a discriminator head to fake_score and requires real videos." + ) + parser.add_argument( + "--gan_d_weight", + type=float, + default=0.0, + help="One-Forcing: weight of the GAN discriminator loss added to the fake_score denoising loss (0 disables)." + ) + parser.add_argument( + "--gan_feature_layers", + nargs="+", + type=int, + default=[21, 29], + help="fake_score block indices whose outputs feed the discriminator head (one register token per layer)." + ) + parser.add_argument( + "--gan_block_ffn_dim", + type=int, + default=2048, + help="FFN dim 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 the relativistic GAN loss (fake vs. real logit difference)." + ) + parser.add_argument( + "--r1_weight", + type=float, + default=0.0, + help="Finite-difference R1 penalty on noisy real latents (APT form, not divided by sigma)." + ) + parser.add_argument( + "--r1_sigma", + type=float, + default=0.01, + help="Perturbation std of the R1 penalty." + ) + parser.add_argument( + "--r2_weight", + type=float, + default=0.0, + help="Finite-difference R2 penalty on noisy generated latents." + ) + parser.add_argument( + "--r2_sigma", + type=float, + default=0.01, + help="Perturbation std of the R2 penalty." + ) args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) @@ -914,6 +1003,9 @@ def main(): 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: @@ -970,6 +1062,9 @@ def main(): args.mixed_precision = accelerator.mixed_precision args.denoising_step_indices_list = [int(i) for i in args.denoising_step_indices_list] + use_gan = any(w > 0 for w in (args.gan_g_weight, args.gan_d_weight, args.r1_weight, args.r2_weight)) + 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.") # Load scheduler, tokenizer and models. noise_scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs'])) @@ -1030,7 +1125,7 @@ def main(): low_cpu_mem_usage=True, ).to(weight_dtype) real_score_transformer3d = WanTransformer3DModel.from_pretrained( - os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), + os.path.join(args.real_score_model_name_or_path or 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, ).to(weight_dtype) @@ -1062,7 +1157,8 @@ def main(): 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) + if args.real_score_model_name_or_path is None: + 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 @@ -1116,6 +1212,33 @@ def main(): if trainable_module_name in name: param.requires_grad = True break + if use_gan: + add_wan_gan_discriminator( + fake_score_transformer3d, args.gan_feature_layers, args.gan_block_ffn_dim, args.gan_block_num_heads, + weight_dtype, 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 only supports 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(), + decay=args.ema_decay, + 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"): @@ -1178,6 +1301,32 @@ def main(): 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() @@ -1269,6 +1418,11 @@ def main(): 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( @@ -1311,7 +1465,7 @@ def main(): args.random_hw_adapt = False # Get the dataset - if args.train_mode != "normal" or args.use_teacher_forcing: + 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, @@ -1339,7 +1493,7 @@ def main(): return length_to_frame_num - if (args.enable_bucket and args.train_mode != "normal") or args.use_teacher_forcing: + 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( @@ -1682,6 +1836,9 @@ def main(): 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: @@ -1772,6 +1929,8 @@ def main(): 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, @@ -1832,6 +1991,8 @@ def main(): # 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 @@ -1852,7 +2013,7 @@ def main(): 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: + 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) @@ -1964,7 +2125,7 @@ def main(): 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: + if args.use_teacher_forcing or use_gan: clean_latents = _batch_encode_vae(pixel_values) else: clean_latents = None @@ -1996,7 +2157,7 @@ def main(): clip_context.append(_clip_context) clip_context = torch.cat(clip_context) - if args.use_teacher_forcing: + if args.use_teacher_forcing or use_gan: target_shape = clean_latents.size() else: target_shape = mask_latents.size() @@ -2504,7 +2665,7 @@ def main(): 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 (ไธ€ๆฌกๆ€งๅค„็†ๆ•ดไธช่ง†้ข‘) + # Use block_mask for causal training (process the whole video at once) generator_pred = generator_transformer3d( x=generator_noise_list, context=prompt_embeds, @@ -2537,6 +2698,8 @@ def main(): # 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), @@ -2642,17 +2805,53 @@ def main(): train_dmd_loss += avg_dmd_loss.item() train_gen_log_count += 1 + # One-Forcing GAN generator loss: gradients flow through fake_score into generator_pred + 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") - fake_score_transformer3d = fake_score_transformer3d.to("cpu") + if args.gan_g_weight == 0: + fake_score_transformer3d = fake_score_transformer3d.to("cpu") torch.cuda.empty_cache() - accelerator.backward(dmd_loss) + 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: + # Drop the fake_score grads produced by the generator GAN loss + critic_optimizer.zero_grad() if args.low_vram: fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device) @@ -2941,6 +3140,8 @@ def main(): 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( @@ -2950,14 +3151,41 @@ def main(): ) with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): - 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, - ) + if use_gan: + # Generated and real latents share critic_noise / critic_timestep. The critic forward is + # split into three passes so the peak activation stays at 2x batch instead of 4x: + # A) fake + real for the denoising and GAN loss, B) real + R1 perturbation, C) fake + R2. + 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) + 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 + ) + 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, + ) def custom_mse_loss(noise_pred, target, weighting=None, threshold=50): noise_pred = noise_pred.float() @@ -2974,8 +3202,53 @@ def main(): 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 - - accelerator_fake_score_transformer3d.backward(denoising_loss) + + 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 + + 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() @@ -2988,13 +3261,22 @@ def main(): # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: + # EMA tracks the generator only on rounds where it actually ran optimizer.step(); + # with gen_update_interval > 1 the critic-only rounds leave the generator unchanged. + if args.use_ema and generator_update: + ema_generator_transformer3d.step(generator_transformer3d.parameters()) progress_bar.update(1) global_step += 1 - accelerator.log({"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)}, step=global_step) + train_log = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / max(train_gen_log_count, 1)} + if use_gan: + train_log.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(train_log, 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: @@ -3038,6 +3320,10 @@ def main(): if args.validation_prompts is not None and global_step % args.validation_steps == 0: with progress_bar.paused(): + if args.use_ema: + # Store the generator parameters temporarily and load the EMA parameters to perform inference. + ema_generator_transformer3d.store(generator_transformer3d.parameters()) + ema_generator_transformer3d.copy_to(generator_transformer3d.parameters()) log_validation( vae, text_encoder, @@ -3050,8 +3336,15 @@ def main(): weight_dtype, global_step, ) + if args.use_ema: + # Switch back to the original generator parameters. + ema_generator_transformer3d.restore(generator_transformer3d.parameters()) logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + if use_gan: + logs["gan_d_loss"] = gan_d_loss.detach().item() + if args.gan_g_weight > 0: + logs["gan_g_loss"] = gan_g_loss.detach().item() progress_bar.set_postfix(**logs) if global_step >= args.max_train_steps: @@ -3059,6 +3352,10 @@ def main(): if args.validation_prompts is not None and epoch % args.validation_epochs == 0: with progress_bar.paused(): + if args.use_ema: + # Store the generator parameters temporarily and load the EMA parameters to perform inference. + ema_generator_transformer3d.store(generator_transformer3d.parameters()) + ema_generator_transformer3d.copy_to(generator_transformer3d.parameters()) log_validation( vae, text_encoder, @@ -3071,6 +3368,9 @@ def main(): weight_dtype, global_step, ) + if args.use_ema: + # Switch back to the original generator parameters. + ema_generator_transformer3d.restore(generator_transformer3d.parameters()) # 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 diff --git a/scripts/wan2.1_self_forcing/train_ode.py b/scripts/wan2.1_self_forcing/train_ode.py index 72c9050..2acf5ba 100644 --- a/scripts/wan2.1_self_forcing/train_ode.py +++ b/scripts/wan2.1_self_forcing/train_ode.py @@ -15,6 +15,7 @@ # See the License for the specific language governing permissions and import argparse +import dataclasses import gc import logging import math @@ -362,6 +363,10 @@ def parse_args(): parser.add_argument("--adam_weight_decay", type=float, default=3e-2, help="Weight decay to use.") parser.add_argument("--adam_epsilon", type=float, default=1e-10, 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( + "--fsdp_reduce_fp32", action="store_true", + help="Reduce gradients in fp32 under FSDP mixed precision (params are still gathered in low precision).", + ) parser.add_argument( "--mixed_precision", type=str, @@ -581,6 +586,11 @@ def main(): 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: diff --git a/videox_fun/dist/fsdp.py b/videox_fun/dist/fsdp.py index b8d3a87..8197cd4 100755 --- a/videox_fun/dist/fsdp.py +++ b/videox_fun/dist/fsdp.py @@ -146,6 +146,13 @@ def shard_model( and p.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2)): p.data = p.data.to(param_dtype) + # `device_id` only moves the managed params/buffers, but `sync_module_states` checks every buffer of the root + # module, including those of the ignored modules (e.g. a computed rope `inv_freq`). Place the ignored modules + # on the device before wrapping. They are small fp32 heads kept replicated (not sharded), so FSDP's CPUOffload + # does not cover them: they stay GPU-resident even in offload mode, which is the intended numerical-precision pin. + for m in ignored_modules: + m.to(device_id) + # With CPU offload, bound the prefetch depth so at most one extra all-gather is in flight; otherwise FSDP # would keep the gathered (unsharded) params of upcoming blocks resident and eat the savings. cpu_offload = CPUOffload(offload_params=True) if offload_to_cpu else None @@ -164,12 +171,6 @@ def shard_model( limit_all_gathers=offload_to_cpu, ignored_states=ignored_modules if ignored_modules else None) - # `device_id`/`sync_module_states` only manage the sharded params; the ignored modules must be placed on - # the device explicitly. They are small fp32 heads kept replicated (not sharded), so FSDP's CPUOffload does - # not cover them: they stay GPU-resident even in offload mode, which is the intended numerical-precision pin. - for m in ignored_modules: - m.to(device_id) - return model diff --git a/videox_fun/models/__init__.py b/videox_fun/models/__init__.py index c6540de..5acd291 100755 --- a/videox_fun/models/__init__.py +++ b/videox_fun/models/__init__.py @@ -94,7 +94,12 @@ from .spark_transformer3d import (SparkWan2_2Transformer3DModel, SparkWanSelfAttention, SparkWanTransformer3DModel) from .turbowan_transformer3d import TurboWanTransformer3DModel +from .fastvideo_wan_transformer3d import (FastVideoWanSelfAttention, + FastVideoWanTransformer3DModel) +from .fastvideo_minimax_h3_transformer3d import \ + FastVideoMiniMaxH3Transformer3DModel from .wan_audio_encoder import WanAudioEncoder +from .wan_gan_discriminator import add_wan_gan_discriminator from .wan_image_encoder import CLIPModel from .wan_latent_upsampler import WanLatentUpsamplerModel from .wan_tae import AutoencoderTinyWan diff --git a/videox_fun/models/attention_kernel.py b/videox_fun/models/attention_kernel.py index 449f3c4..de720a2 100644 --- a/videox_fun/models/attention_kernel.py +++ b/videox_fun/models/attention_kernel.py @@ -14,10 +14,575 @@ Citation (please cite if you use this code): } """ +import math + import torch import triton import triton.language as tl +# ====================================================================================== +# This module holds TWO sparse-attention kernel families, kept in contiguous sections. +# +# SECTION A -- VSA, variable-length block-sparse. This is the path the repo actually runs +# for MiniMax-H3 and Wan sparse attention (attention_utils.VSASparseAttention.forward_h3 +# -> sparse_linear_attention_varlen). Everything that path needs lives in this section, +# in call order: the two vendored upstream blocks (the official tile-64 forward kernel, +# then the on-device mask -> key-list compaction), our own varlen backward kernels, the +# autograd Function, and its wrapper. +# +# SECTION B -- the fixed-length SLA original this file was modified from (see the licence +# header above). Still live, NOT dead code: attention_utils' flat sparse path calls +# get_block_map + _sparse_linear_attention.apply, and its tile pooling calls mean_pool. +# +# The two families share no kernels. A symbol suffixed `_varlen` belongs to section A. +# ====================================================================================== + +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> VENDORED FROM UPSTREAM (Apache-2.0) -- BLOCK 1/2 BEGIN >>>>>>>>>>>>>>>>>>>> +# Source: hao-ai-lab/fastvideo `fastvideo_kernel` 0.3.5, file +# python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py, lines 11-166 and 703-752. +# It replaces the hand-written `_attn_fwd_varlen` that used to live further down this module. +# +# Why vendored: upstream's compiled sm90 CUDA kernel (`fastvideo_kernel._C`) ships only cp312 / +# torch==2.12.0 wheels, so it is not installable here (we run torch 2.7.0 / py3.10), whereas its +# pure-Triton tile-64 block-sparse forward runs as-is with no build step. A kernel-level micro-bench +# on the Wan VSA workload (seq 33390, 672 tiles, sparsity 0.8, bf16, H800) measured this forward at +# 6.31 ms/call vs `_attn_fwd_varlen` at 13.48 ms/call (2.14x), with identical output (rel_l2 +# 6.2e-5). The wins are @triton.autotune over num_stages/num_warps, tl.make_block_ptr + tl.advance +# loads, a fused tl.dot(p, v, acc) accumulate, and a per-query-tile variable kv-block count (vs a +# uniform-MAXK loop over a dummy-padded tile). +# +# This is deliberately NOT the fastest forward available. That same sm90 CUDA kernel (ThunderKittens) +# IS buildable from source on this box -- only its prebuilt wheels are pinned -- and measures +# 4.66 ms/call on the workload above, i.e. 1.35x over this Triton kernel (step-level VSA-vs-dense +# 1.019x -> 1.115x at 480P, 1.233x -> 1.381x at 720P). It stays out of the production path because +# it adds a build dependency and fails SILENTLY three ways: on a non-sm90a device pass it compiles +# to an empty stub returning uninitialised memory; it reinterpret_casts `variable_block_sizes` to +# int32 with no dtype check, so an int64 tensor yields all-NaN; and its lse is [B,H,L,1] in the +# natural-log domain scaled by -1/qk_scale rather than our [B,H,L] base-2. The build recipe, the +# runtime gate it would need, and the measurement set are archived in +# `_fastvideo_sm90_kernel/README.md`. +# +# Deviations from upstream, listed so a diff against the sdist still audits this block: +# 1. its `import torch / triton / triton.language`, and its SPARSE-ADDITION `import math`, are +# hoisted to this module's top import block -- same names, no behavioural change, and the +# two `SPARSE ADDITION BEGIN/END` banner comments upstream wrapped around that import are +# dropped with it (the one around `@triton.autotune` below is kept); +# 2. the inner dot's `p.to(tl.bfloat16)` is `p.to(v.dtype)` here, because our varlen path and +# its fp32 oracle tests must not be forced to bf16 (output identical when v IS bf16). +# Everything else between the BEGIN/END markers is byte-identical to upstream. Our own backward +# (`_attn_bwd_*_varlen`) is untouched by this block: it supplies the forward only. + +# BLOCK_M / BLOCK_N are fixed at 64 because they are structural, not tunable: +# the kernel indexes the top-k list per BLOCK_M q-tile and addresses keys as +# kv_idx * BLOCK_N, so both must match the granularity q2k_index and +# variable_block_sizes were built at. +# +# num_stages / num_warps ARE free, and the previous {3, 4, 7} was inherited from +# the upstream tutorial rather than tuned here. It skips 5 and 6; on Blackwell +# (sm_121) the optimum is num_stages=5, so the search could not reach it. Both +# block paths independently select 5 once it is available. Autotune still picks +# per architecture, so other GPUs re-tune rather than inheriting this choice. +configs = [ + triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \ + for BM in [64]\ + for BN in [64]\ + for s in [2, 3, 4, 5, 6, 7]\ + for w in [4, 8]\ +] + + +# โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ SPARSE ADDITION BEGIN โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +@triton.autotune(configs, key=["N_CTX_Q", "HEAD_DIM"]) +@triton.jit +def _attn_fwd_sparse( + Q, + K, + V, + sm_scale, # + q2k_index, + q2k_num, + max_kv_blks, # + variable_block_sizes, + M, + Out, # + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vk, + stride_vn, + stride_oz, + stride_oh, + stride_om, + stride_on, + Z, + H, + N_CTX_Q, # + N_CTX_KV, # + HEAD_DIM: tl.constexpr, # + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + STAGE: tl.constexpr): + """ + 64ร—64 **block-sparse** forward kernel. Back-prop kernels remain dense + (32ร—64 and 64ร—32) โ€“ memory footprint unchanged. + """ + + # ----- program-id mapping ----- + q_blk = tl.program_id(0) # Q-tile index + off_hz = tl.program_id(1) # fused (batch, head) + b = off_hz // H + h = off_hz % H + q_tiles = N_CTX_Q // BLOCK_M + meta_base = ((b * H + h) * q_tiles + q_blk) + + kv_blocks = tl.load(q2k_num + meta_base) # int32 + kv_ptr = q2k_index + meta_base * max_kv_blks # ptr to list + + # ----- base pointers ----- + # Note: when q and kv have different sequence lengths, their per-(batch,head) + # strides differ, so we must compute separate base offsets. + q_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh) + k_off = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh) + v_off = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh) + o_off = (b.to(tl.int64) * stride_oz + h.to(tl.int64) * stride_oh) + + Q_ptr = tl.make_block_ptr(base=Q + q_off, + shape=(N_CTX_Q, HEAD_DIM), + strides=(stride_qm, stride_qk), + offsets=(q_blk * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0)) + + K_base = tl.make_block_ptr(base=K + k_off, + shape=(HEAD_DIM, N_CTX_KV), + strides=(stride_kk, stride_kn), + offsets=(0, 0), + block_shape=(HEAD_DIM, BLOCK_N), + order=(0, 1)) + + v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0) + V_base = tl.make_block_ptr(base=V + v_off, + shape=(N_CTX_KV, HEAD_DIM), + strides=(stride_vk, stride_vn), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=v_order) + + O_ptr = tl.make_block_ptr(base=Out + o_off, + shape=(N_CTX_Q, HEAD_DIM), + strides=(stride_om, stride_on), + offsets=(q_blk * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0)) + + # ----- accumulators ----- + offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M) + m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 + acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + qk_scale = sm_scale * 1.44269504 # 1/ln2 + q = tl.load(Q_ptr) + + # ----- sparse loop over valid K/V tiles ----- + for i in range(0, kv_blocks): + kv_idx = tl.load(kv_ptr + i).to(tl.int32) + block_size = tl.load(variable_block_sizes + kv_idx) + K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N)) + V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0)) + + k = tl.load(K_ptr) + qk = tl.dot(q, k) + # mask out invalid columns + mask = tl.arange(0, BLOCK_N) < block_size + qk = tl.where(mask[None, :], qk, -float("inf")) + + m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) + p = tl.math.exp2(qk * qk_scale - m_ij[:, None]) + l_ij = tl.sum(p, 1) + + alpha = tl.math.exp2(m_i - m_ij) + l_i = l_i * alpha + l_ij + acc = acc * alpha[:, None] + + v = tl.load(V_ptr) + acc = tl.dot(p.to(v.dtype), v, acc) # dtype-agnostic: official hardcodes bf16, + # but our varlen path/tests also run fp32; p.to(v.dtype) keeps bf16 identical. + m_i = m_ij + + # ----- epilogue ----- + m_i += tl.math.log2(l_i) + acc = acc / l_i[:, None] + tl.store(M + off_hz * N_CTX_Q + offs_m, m_i) + tl.store(O_ptr, acc.to(Out.type.element_ty)) + +def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes): + B, H, Tq, D = q.shape + Tkv = k.shape[2] + sm_scale = 1.0 / math.sqrt(D) + max_kv_blks = q2k_index.shape[-1] + assert Tq % 64 == 0, f"q length must be a multiple of 64, but got {Tq}" + assert Tkv % 64 == 0, f"kv length must be a multiple of 64, but got {Tkv}" + assert q2k_num.shape[ + -1] == Tq // 64, f"shape mismatch, Tq // 64 = {Tq // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}" + assert variable_block_sizes.numel() == Tkv // 64, ( + f"shape mismatch, variable_block_sizes must have length {Tkv // 64}, " + f"got {variable_block_sizes.numel()}") + o = torch.empty_like(q) + M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device) + + grid = lambda _: (triton.cdiv(Tq, 64), B * H, 1) + _attn_fwd_sparse[grid](q, + k, + v, + sm_scale, + q2k_index, + q2k_num, + max_kv_blks, + variable_block_sizes, + M, + o, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + o.stride(0), + o.stride(1), + o.stride(2), + o.stride(3), + B, + H, + Tq, + Tkv, + HEAD_DIM=D, + STAGE=3) + + return o, M +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> VENDORED FROM UPSTREAM (Apache-2.0) -- BLOCK 1/2 END >>>>>>>>>>>>>>>>>>>> + +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> VENDORED FROM UPSTREAM (Apache-2.0) -- BLOCK 2/2 BEGIN >>>>>>>>>>>>>>>>>>>> +# Source: hao-ai-lab/fastvideo `fastvideo_kernel` 0.3.5, file +# python/fastvideo_kernel/triton_kernels/index.py, lines 32-61 and 106-144. +# +# Why vendored: same install blocker as block 1/2, and this helper is pure Triton with no +# dependency on the compiled `_C` extension. +# +# What it replaces: our mask -> key-list conversion used `mask.float().topk(int(counts.max()))`, +# whose `int(...)` forces a GPU->CPU sync on every layer of every step and materializes an fp32 +# copy of the [B,H,n,n] mask. `map_to_index` compacts each row on-device +# (grid = (bs, h, num_q_blocks)), pads with -1, and returns the per-row counts as an int32 tensor +# -- no host sync, no fp32 temp. +# +# Semantics relied on by the caller: for each query block, the first `index_num` entries of +# `index` are its selected kv-block ids in ASCENDING order; the rest are -1 and are never read +# (both the official forward kernel above and our backward loop bound their iteration by it). +# +# Deviation from upstream: its `import torch / triton / triton.language` are hoisted to this +# module's top, and its unused `from typing import Tuple` (kept upstream only for parity) is +# dropped. Between the BEGIN/END markers the block is otherwise byte-identical. + +@triton.jit +def map_to_index_kernel( + map_ptr, + index_ptr, + index_num_ptr, + map_bs_stride, + map_h_stride, + map_q_stride, + map_kv_stride, + index_bs_stride, + index_h_stride, + index_q_stride, + index_kv_stride, + index_num_bs_stride, + index_num_h_stride, + index_num_q_stride, + num_kv_blocks, +): + b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2) + index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride + map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride + + num = 0 + for i in tl.range(num_kv_blocks): + map_entry = tl.load(map_ptr_base + i * map_kv_stride) + if map_entry: + tl.store(index_ptr_base + num * index_kv_stride, i) + num += 1 + + tl.store(index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + q * index_num_q_stride, num) + + +def map_to_index(block_map: torch.Tensor): + """ + Convert a block map to indices and counts. + + Args: + block_map: [bs, h, num_q_blocks, num_kv_blocks] + The block map tensor. + + Returns: + index: [bs, h, num_q_blocks, num_kv_blocks] + The indices of the blocks. + index_num: [bs, h, num_q_blocks] + The number of blocks for each q block. + """ + bs, h, num_q_blocks, num_kv_blocks = block_map.shape + + index = torch.full((block_map.shape), -1, dtype=torch.int32, device=block_map.device) + index_num = torch.empty((bs, h, num_q_blocks), dtype=torch.int32, device=block_map.device) + + grid = (bs, h, num_q_blocks) + map_to_index_kernel[grid]( + block_map, + index, + index_num, + block_map.stride(0), + block_map.stride(1), + block_map.stride(2), + block_map.stride(3), + index.stride(0), + index.stride(1), + index.stride(2), + index.stride(3), + index_num.stride(0), + index_num.stride(1), + index_num.stride(2), + num_kv_blocks=num_kv_blocks, + ) + + return index, index_num +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> VENDORED FROM UPSTREAM (Apache-2.0) -- BLOCK 2/2 END >>>>>>>>>>>>>>>>>>>> + + +@triton.jit +def _attn_bwd_preprocess_varlen( + O, DO, DELTAS, QSIZE, + L_PAD, + D: tl.constexpr, + BLOCK_M: tl.constexpr, +): + # delta[row] = sum_d o[row,d] * do[row,d], per query tile, over the tile-order buffer. + idx_m = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + qkv_offset = idx_bh * L_PAD * D + offs_m = tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, D) + qsize = tl.load(QSIZE + idx_m) + q_mask = offs_m < qsize + ptrs = qkv_offset + idx_m * BLOCK_M * D + offs_m[:, None] * D + offs_d[None, :] + o = tl.load(O + ptrs, mask=q_mask[:, None], other=0.0) + do = tl.load(DO + ptrs, mask=q_mask[:, None], other=0.0) + delta = tl.sum(o.to(tl.float32) * do.to(tl.float32), axis=1) + tl.store(DELTAS + idx_bh * L_PAD + idx_m * BLOCK_M + offs_m, delta, mask=q_mask) + + +@triton.jit +def _attn_bwd_dq_varlen( + Q, K, V, LSE, DELTAS, DO, DQ, LUT, LUTN, + QSIZE, KSIZE, + qk_scale, + MAXK, + N_TILES, + L_PAD, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # dQ for one query tile: loop over the LUTN[row] key tiles that tile actually selected (ids in + # LUT, ascending, -1 beyond the count) and recompute p from LSE. The bound is a runtime per-row + # value, so a row that selected few tiles does no extra work and never touches the -1 pads. + idx_m = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + qkv_offset = idx_bh * L_PAD * D + lse_offset = idx_bh * L_PAD + offs_m = tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + + qsize = tl.load(QSIZE + idx_m) + q_mask = offs_m < qsize + m_ptrs = qkv_offset + idx_m * BLOCK_M * D + offs_m[:, None] * D + offs_d[None, :] + q = tl.load(Q + m_ptrs, mask=q_mask[:, None], other=0.0) + do = tl.load(DO + m_ptrs, mask=q_mask[:, None], other=0.0) + delta = tl.load(DELTAS + lse_offset + idx_m * BLOCK_M + offs_m, mask=q_mask, other=0.0) + lse = tl.load(LSE + lse_offset + idx_m * BLOCK_M + offs_m, mask=q_mask, other=float('inf')) + + dq = tl.zeros([BLOCK_M, D], dtype=tl.float32) + lut_base = LUT + (idx_bh * N_TILES + idx_m) * MAXK + kv_blocks = tl.load(LUTN + idx_bh * N_TILES + idx_m) + for i in tl.range(0, kv_blocks): + idx_n = tl.load(lut_base + i) + ksize = tl.load(KSIZE + idx_n) + n_mask = offs_n < ksize + n_ptrs = qkv_offset + idx_n * BLOCK_N * D + offs_n[:, None] * D + offs_d[None, :] + k = tl.load(K + n_ptrs, mask=n_mask[:, None], other=0.0) # [BLOCK_N, D] + v = tl.load(V + n_ptrs, mask=n_mask[:, None], other=0.0) # [BLOCK_N, D] + qk = tl.dot(q, k.T) * (qk_scale * 1.4426950408889634) # [BLOCK_M, BLOCK_N] + p = tl.math.exp2(qk - lse[:, None]) + p = tl.where(n_mask[None, :], p, 0.0) + dp = tl.dot(do, v.T).to(tl.float32) # [BLOCK_M, BLOCK_N] + ds = p * (dp - delta[:, None]) + dq += tl.dot(ds.to(k.dtype), k) # [BLOCK_M, D] + tl.store(DQ + m_ptrs, (dq * qk_scale).to(DQ.type.element_ty), mask=q_mask[:, None]) + + +@triton.jit +def _attn_bwd_dkdv_varlen( + Q, K, V, DO, DK, DV, KBID, LSE, DELTAS, + QSIZE, KSIZE, + qk_scale, + N_TILES, + L_PAD, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + # dK/dV for one key tile: loop over ALL query tiles, keep those whose KBID incidence is 1. + idx_n = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + qkv_offset = idx_bh * L_PAD * D + lse_offset = idx_bh * L_PAD + offs_m = tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + + ksize = tl.load(KSIZE + idx_n) + n_mask = offs_n < ksize + n_ptrs = qkv_offset + idx_n * BLOCK_N * D + offs_n[:, None] * D + offs_d[None, :] + k = tl.load(K + n_ptrs, mask=n_mask[:, None], other=0.0) # [BLOCK_N, D] + v = tl.load(V + n_ptrs, mask=n_mask[:, None], other=0.0) # [BLOCK_N, D] + + dk = tl.zeros([BLOCK_N, D], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, D], dtype=tl.float32) + kbid_base = KBID + idx_bh * N_TILES * N_TILES + idx_n + for idx_m in tl.range(0, N_TILES): + kbid = tl.load(kbid_base + idx_m * N_TILES) + if kbid == 1: + qsize = tl.load(QSIZE + idx_m) + m_mask = offs_m < qsize + m_ptrs = qkv_offset + idx_m * BLOCK_M * D + offs_m[:, None] * D + offs_d[None, :] + q = tl.load(Q + m_ptrs, mask=m_mask[:, None], other=0.0) # [BLOCK_M, D] + do = tl.load(DO + m_ptrs, mask=m_mask[:, None], other=0.0) + lse = tl.load(LSE + lse_offset + idx_m * BLOCK_M + offs_m, mask=m_mask, other=float('inf')) + delta = tl.load(DELTAS + lse_offset + idx_m * BLOCK_M + offs_m, mask=m_mask, other=0.0) + qkT = tl.dot(k, q.T) * (qk_scale * 1.4426950408889634) # [BLOCK_N, BLOCK_M] + pT = tl.math.exp2(qkT - lse[None, :]) + pT = tl.where(m_mask[None, :], pT, 0.0) + dv += tl.dot(pT.to(do.dtype), do) # [BLOCK_N, D] + dpT = tl.dot(v, do.T).to(tl.float32) # [BLOCK_N, BLOCK_M] + dsT = pT * (dpT - delta[None, :]) + dk += tl.dot(dsT.to(q.dtype), q) # [BLOCK_N, D] + tl.store(DK + n_ptrs, (dk * qk_scale).to(DK.type.element_ty), mask=n_mask[:, None]) + tl.store(DV + n_ptrs, dv.to(DV.type.element_ty), mask=n_mask[:, None]) + + +class _sparse_linear_attention_varlen(torch.autograd.Function): + @staticmethod + def forward(ctx, q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes, kb_id, qk_scale): + assert q_tile.is_contiguous() and k_tile.is_contiguous() and v_tile.is_contiguous() + B, H, L_BUF, D = q_tile.shape + # Forward runs the vendored official Triton block-sparse kernel, which consumes the + # official (q2k_index, q2k_num) pair map_to_index produced: the first q2k_num[row] entries + # of each row are that query tile's kv-tile ids and the -1 pads are never read. Its M is + # [B,H,L_BUF] base-2 lse (m_i + log2 l_i), byte-compatible with what our backward expects, + # so the save_for_backward contract is unchanged. + out, lse = triton_block_sparse_attn_forward( + q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes) + ctx.save_for_backward(q_tile, k_tile, v_tile, q2k_idx, q2k_num, + block_sizes, kb_id, lse, out) + ctx.qk_scale = qk_scale + return out + + @staticmethod + def backward(ctx, do): + (q_tile, k_tile, v_tile, q2k_idx, q2k_num, + block_sizes, kb_id, lse, out) = ctx.saved_tensors + do = do.contiguous() + B, H, L_BUF, D = q_tile.shape + num_tiles = L_BUF // 64 + MAXK = q2k_idx.shape[-1] + dq = torch.zeros_like(q_tile) + dk = torch.zeros_like(k_tile) + dv = torch.zeros_like(v_tile) + delta = torch.empty((B, H, L_BUF), device=q_tile.device, dtype=torch.float32) + grid = (num_tiles, B * H) + _attn_bwd_preprocess_varlen[grid]( + out, do, delta, block_sizes, L_BUF, D, 64, + num_warps=8, num_stages=3, + ) + _attn_bwd_dq_varlen[grid]( + q_tile, k_tile, v_tile, lse, delta, do, dq, q2k_idx, q2k_num, + block_sizes, block_sizes, ctx.qk_scale, + MAXK, num_tiles, L_BUF, D, 64, 64, + num_warps=8, num_stages=3, + ) + _attn_bwd_dkdv_varlen[grid]( + q_tile, k_tile, v_tile, do, dk, dv, kb_id, lse, delta, + block_sizes, block_sizes, ctx.qk_scale, + num_tiles, L_BUF, D, 64, 64, + num_warps=8, num_stages=3, + ) + return dq, dk, dv, None, None, None, None, None + + +def sparse_linear_attention_varlen(q_tile, k_tile, v_tile, block_mask, block_sizes, + qk_scale=None): + """Tile-order block-sparse attention for the H3/Wan VSA path (forward + backward). + + q_tile/k_tile/v_tile: [B, H, L_BUF, D] contiguous, rows in tile order, pad slots zero. + block_mask: [B, H, num_tiles, num_tiles] bool query->key incidence at the PHYSICAL 64-token + block granularity the kernels run at (a logical 128/256-tile selection is expanded down to + it by ``_h3_expand_mask_and_sizes_to_64``). It drives both stages: map_to_index compacts it + into the per-query key list used by the forward and by dQ, and it is the dK/dV reverse-loop + incidence. block_sizes: [num_tiles] int valid token count per 64-token block; 0 is legal + (a partial logical tile's trailing sub-blocks) and makes that block contribute nothing. + num_tiles is derived as L_BUF // 64, so the buffer carries no separate dummy tile. + Returns a NEW [B, H, L_BUF, D] output buffer (pad rows are garbage; the caller untiles only the + real rows). Differentiable -- gradients flow to q_tile/k_tile/v_tile. + """ + assert q_tile.is_contiguous() and k_tile.is_contiguous() and v_tile.is_contiguous() + B, H, L_BUF, D = q_tile.shape + assert L_BUF > 0 and L_BUF % 64 == 0, "VSA tile ็ผ“ๅ†ฒ้•ฟๅบฆๅฟ…้กปๆ˜ฏๆญฃ็š„ 64 ๅ€ๆ•ฐใ€‚" + assert k_tile.shape == q_tile.shape and v_tile.shape == q_tile.shape + num_tiles = L_BUF // 64 + assert tuple(block_mask.shape) == (B, H, num_tiles, num_tiles), \ + f"block_mask {tuple(block_mask.shape)} vs expected {(B, H, num_tiles, num_tiles)}" + assert block_sizes.numel() == num_tiles, \ + f"block_sizes has {block_sizes.numel()} entries, expected {num_tiles}" + default_scale = D ** -0.5 + if qk_scale is not None and not math.isclose(qk_scale, default_scale, rel_tol=1e-7): + raise ValueError("VSA ๅ‰ๅ‘ๆ ธไป…ๆ”ฏๆŒ head_dim**-0.5๏ผ›ไธๆŽฅๅ—่‡ชๅฎšไน‰ qk_scaleใ€‚") + qk_scale = default_scale + # On-device compaction by the vendored official Triton kernel: no host sync (the topk over + # int(counts.max()) it replaces synced once per layer per step). `q2k_idx` keeps the full + # num_tiles width and pads with -1; both kernels bound their loops by q2k_num. + q2k_idx, q2k_num = map_to_index(block_mask.contiguous()) + block_sizes = block_sizes.to(torch.int32).contiguous() + kb_id = block_mask.to(torch.int8).contiguous() + return _sparse_linear_attention_varlen.apply( + q_tile, k_tile, v_tile, q2k_idx, q2k_num, block_sizes, kb_id, qk_scale) + + +# ====================================================================================== +# SECTION B -- fixed-length SLA original (NOT varlen): mean_pool / get_block_map / _attn_fwd / +# _attn_bwd_{preprocess,dq,dkdv} / class _sparse_linear_attention. Do not confuse these with +# section A's `*_varlen` kernels; nothing in section A calls into this section. +# ====================================================================================== @triton.jit def compress_kernel( diff --git a/videox_fun/models/attention_utils.py b/videox_fun/models/attention_utils.py index b8c9c45..969b3a9 100644 --- a/videox_fun/models/attention_utils.py +++ b/videox_fun/models/attention_utils.py @@ -1,3 +1,5 @@ +import functools +import math import os import warnings @@ -42,7 +44,8 @@ except Exception: SAGE_ATTENTION_AVAILABLE = False -from .attention_kernel import _sparse_linear_attention, get_block_map +from .attention_kernel import (_sparse_linear_attention, get_block_map, + mean_pool, sparse_linear_attention_varlen) def convert_qkv_dtype(q, k, v): @@ -550,4 +553,498 @@ class RoLASparseAttention(nn.Module): [B, H, L, 1] gate in (0, 1). ''' gb = self.gate_bias if head_slice is None else self.gate_bias[:, head_slice] - return torch.sigmoid(self.gate_proj(x_bhld) + gb) \ No newline at end of file + return torch.sigmoid(self.gate_proj(x_bhld) + gb) + + +# Logical tile geometries keyed by tile volume -- the official VSA_H3_TILE_SHAPES. 64 is what the +# official defaults actually select for the checkpoints this repo runs: the Wan backend's +# VSA_TILE_SIZE is (4,4,4), and both released FastH3 inference contracts pin "vsa_tile_size": 64. +# (256/(4,8,8) is the H3 backend's *code* default, aimed at the Blackwell CuTe fastpath.) 128 and +# 256 stay selectable and reach the same 64-block kernel through the route-A expansion below. +_H3_TILE_SHAPES = {64: (4, 4, 4), 128: (4, 4, 8), 256: (4, 8, 8)} +_H3_TILE_ELEMS = 64 +# Physical block size the Triton kernels are instantiated at: a logical tile of N elems is N // 64 +# CONTIGUOUS 64-blocks in the padded tile buffer. +_H3_KERNEL_BLOCK = 64 + + +def _h3_compute_topk(sparsity, num_blocks): + """Key tiles to KEEP for a sparsity level, clamped to [1, num_blocks].""" + return max(1, min(math.ceil((1.0 - sparsity) * num_blocks), num_blocks)) + + +@functools.lru_cache(maxsize=32) +def _h3_tile_partition_indices(grid, tile_shape, device): + """Region tile-order -> original flat region index (port of get_tile_partition_indices).""" + T, H, W = grid + ts, hs, ws = tile_shape + indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W) + parts = [] + for t in range(math.ceil(T / ts)): + for h in range(math.ceil(H / hs)): + for w in range(math.ceil(W / ws)): + parts.append(indices[t * ts:min(t * ts + ts, T), + h * hs:min(h * hs + hs, H), + w * ws:min(w * ws + ws, W)].flatten()) + return torch.cat(parts, dim=0) + + +@functools.lru_cache(maxsize=32) +def _h3_variable_block_sizes(grid, num_tiles, device, tile_shape): + """Valid (non-pad) token count inside every 3-D tile (port of construct_variable_block_sizes).""" + t, h, w = grid + ts_t, ts_h, ts_w = tile_shape + n_t, n_h, n_w = num_tiles + + def _sizes(dim_len, tile, n): + sizes = torch.full((n,), tile, dtype=torch.long, device=device) + rem = dim_len - (n - 1) * tile + sizes[-1] = rem if rem > 0 else tile + return sizes + + t_s = _sizes(t, ts_t, n_t) + h_s = _sizes(h, ts_h, n_h) + w_s = _sizes(w, ts_w, n_w) + return (t_s[:, None, None] * h_s[None, :, None] * w_s[None, None, :]).reshape(-1) + + +def _h3_non_pad_index(variable_block_sizes, tile_elems): + """Padded-buffer slot index of every real row, in tile order (port of get_non_pad_index).""" + n_win = variable_block_sizes.shape[0] + device = variable_block_sizes.device + starts = torch.arange(n_win, device=device) * tile_elems + index_pad = starts[:, None] + torch.arange(tile_elems, device=device)[None, :] + index_mask = torch.arange(tile_elems, device=device)[None, :] < variable_block_sizes[:, None] + return index_pad[index_mask] + + +def _h3_validate_segment_geometry(segments, variable_block_sizes, untile_combined_index, tile_elems): + R'''Fail synchronously on out-of-bounds tile geometry (port of the official guard). + + Invariants the block-sparse kernels trust without checking: every tile's valid size is in + (0, tile_elems]; the sizes sum to the packed sequence length; and ``untile_combined_index`` + maps each packed row to exactly one non-pad slot of the padded tile buffer. A violation would + otherwise surface only as an async device fault at some later kernel or collective (e.g. an + FSDP all-gather), which is unattributable -- so raise here, once per cached geometry. + ''' + total = sum(math.prod(seg) if isinstance(seg, tuple) else seg for seg in segments) + n_pad = variable_block_sizes.numel() * tile_elems + sizes_min = int(variable_block_sizes.min()) + sizes_max = int(variable_block_sizes.max()) + sizes_sum = int(variable_block_sizes.sum()) + if sizes_min < 1 or sizes_max > tile_elems or sizes_sum != total: + raise ValueError( + f"VSA-H3 tile sizes out of bounds for segments={segments}, tile_elems={tile_elems}: " + f"min={sizes_min}, max={sizes_max}, sum={sizes_sum}, expected sum={total}.") + if untile_combined_index.numel() != total: + raise ValueError( + f"VSA-H3 untile index has {untile_combined_index.numel()} entries for a packed " + f"sequence of {total} rows (segments={segments}).") + idx_min = int(untile_combined_index.min()) + idx_max = int(untile_combined_index.max()) + if idx_min < 0 or idx_max >= n_pad: + # Range first: the pad-slot gather below would itself index out of bounds (the very async + # fault this guard exists to preempt). + raise ValueError( + f"VSA-H3 untile index is not an injective map into non-pad slots: range " + f"[{idx_min}, {idx_max}] vs padded length {n_pad} (segments={segments}).") + in_tile_offset = untile_combined_index % tile_elems + maps_into_pad = bool( + (in_tile_offset >= variable_block_sizes[untile_combined_index // tile_elems]).any()) + if maps_into_pad or int(torch.unique(untile_combined_index).numel()) != total: + raise ValueError( + f"VSA-H3 untile index is not an injective map into non-pad slots: " + f"pad-slot hit={maps_into_pad} (segments={segments}).") + + +@functools.lru_cache(maxsize=32) +def build_h3_tile_geometry(segments, device, tile_shape=_H3_TILE_SHAPES[_H3_TILE_ELEMS]): + R'''Tile a packed sequence of interleaved dense segments and video regions. + + Port of the official ``_h3_segment_tile_geometry``. ``segments`` lists, in packed + order, dense segment row counts (``int``) and video-region token grids (``(t,h,w)`` + triples). Returns ``(tile_partition_indices, variable_block_sizes, + untile_combined_index, num_prefix_tiles, num_video_tiles, video_tile_spans, + tile_elems)`` where ``video_tile_spans`` holds one ``[start, end)`` tile range per + video region (prefix tiles included in the numbering). + ''' + tile_elems = math.prod(tile_shape) + ts_t, ts_h, ts_w = tile_shape + prefix_sizes, prefix_rows, video_entries = [], [], [] + cursor = 0 + for seg in segments: + if isinstance(seg, tuple): + video_entries.append((cursor, seg)) + cursor += math.prod(seg) + continue + if seg <= 0: + continue + full, rem = divmod(seg, tile_elems) + prefix_sizes.extend([tile_elems] * full) + if rem: + prefix_sizes.append(rem) + prefix_rows.append(torch.arange(cursor, cursor + seg, device=device, dtype=torch.long)) + cursor += seg + num_prefix_tiles = len(prefix_sizes) + + video_index_parts, video_size_parts, video_tile_spans = [], [], [] + tile_cursor = num_prefix_tiles + for start, grid in video_entries: + t, h, w = grid + num_tiles = (math.ceil(t / ts_t), math.ceil(h / ts_h), math.ceil(w / ts_w)) + region_sizes = _h3_variable_block_sizes(grid, num_tiles, device, tile_shape) + video_index_parts.append(_h3_tile_partition_indices(grid, tile_shape, device) + start) + video_size_parts.append(region_sizes) + video_tile_spans.append((tile_cursor, tile_cursor + int(region_sizes.numel()))) + tile_cursor += int(region_sizes.numel()) + num_video_tiles = tile_cursor - num_prefix_tiles + + tile_partition_indices = torch.cat(prefix_rows + video_index_parts) + variable_block_sizes = torch.cat( + [torch.tensor(prefix_sizes, dtype=torch.long, device=device)] + video_size_parts) + non_pad_index = _h3_non_pad_index(variable_block_sizes, tile_elems) + untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)] + # One-time (lru-cached) synchronous bounds check; see _h3_validate_segment_geometry. + _h3_validate_segment_geometry(segments, variable_block_sizes, untile_combined_index, tile_elems) + return (tile_partition_indices, variable_block_sizes, untile_combined_index, + num_prefix_tiles, num_video_tiles, tuple(video_tile_spans), tile_elems) + + +# Per-process reusable padded tile scratch for forward_h3's no-grad (inference) path. +# Mirrors the official FastVideo builder-owned tile_buf_holder: one buffer per +# (role, shape, dtype, device) shared across all layers and steps, so the per-layer +# torch.zeros((B,H,L_buf,D)) allocation+zeroing churn is paid once instead of every call. +# Keyed by role because q/k/v/gate are all live simultaneously within one forward_h3. +_H3_TILE_BUFS = {} + + +def _h3_tile_into_buf(t, dtype, idx, shape, device, role, reuse): + R'''Scatter packed [B,L,H,D] ``t`` into a [B,H,L_buf,D] tile buffer (pad/dummy stay zero). + + Under ``reuse`` (no-grad inference) the buffer is cached in ``_H3_TILE_BUFS`` and only + re-zeroed when the tile geometry (``idx`` identity) changes, exactly like the official + holder: pad slots are never written and non-pad slots are fully overwritten per call, so + a reused buffer stays valid. Callers consume the result before the next layer overwrites + it. Under grad the buffer is fresh -- an in-place write into a buffer that outlives the + step would anchor that step's autograd graph (the official tile-cache OOM). + ''' + B, H, L_buf, D = shape + if not reuse: + buf = torch.zeros(shape, device=device, dtype=dtype) + buf[:, :, idx, :] = t.transpose(1, 2).to(dtype) + return buf + key = (role, B, H, L_buf, D, dtype, str(device)) + ent = _H3_TILE_BUFS.get(key) + if ent is None or ent[0].shape != shape or ent[0].dtype != dtype: + ent = [torch.zeros(shape, device=device, dtype=dtype), None] + _H3_TILE_BUFS[key] = ent + if ent[1] is not idx: # geometry identity changed -> clear stale valid rows once + ent[0].zero_() + ent[1] = idx + ent[0][:, :, idx, :] = t.transpose(1, 2).to(dtype) + return ent[0] + + +def _h3_pool_tiles(x, variable_block_sizes, tile_elems, n_tiles): + R'''fp32 masked mean over each tile. ``x``: [B,H,L_buf,D] (pad slots zero) -> [B,H,n_tiles,D]. + + Pad slots are guaranteed zero (zeros-init buffer, only non-pad slots written), so a + plain fp32 sum divided by the true tile size is exactly the masked mean. + ''' + B, H, _, D = x.shape + xt = x[:, :, :n_tiles * tile_elems].reshape(B, H, n_tiles, tile_elems, D) + pooled = xt.sum(dim=3, dtype=torch.float32) + return pooled / variable_block_sizes.view(1, 1, -1, 1).to(torch.float32) + + +def _h3_build_block_mask(scores, num_prefix_tiles, sparsity, exempt, + video_tile_spans, span_sparsities): + R'''scores [B,H,n_tiles,n_tiles] -> bool mask (port of official ``_build_block_mask``). + + Prefix rows/cols: with ``exempt`` every query keeps every prefix column, and every + query keeps its own top-k of EACH video region under that region's sparsity. + ``sparsity <= 0`` keeps every tile (dense). + ''' + n_tiles = scores.shape[-1] + if sparsity <= 0.0: + span_sparsities = tuple(sparsity for _ in video_tile_spans) + span_topk = [_h3_compute_topk(s, end - start) + for (start, end), s in zip(video_tile_spans, span_sparsities)] + if all(kk == end - start for kk, (start, end) in zip(span_topk, video_tile_spans)): + return torch.ones_like(scores, dtype=torch.bool) + mask = torch.zeros_like(scores, dtype=torch.bool) + if exempt or num_prefix_tiles == 0: + for (start, end), k_span in zip(video_tile_spans, span_topk): + if k_span == end - start: + mask[..., start:end] = True + continue + idx = scores[..., start:end].topk(k_span, dim=-1).indices + start + mask.scatter_(-1, idx, True) + mask[..., :num_prefix_tiles] = True + else: + k_total = min(span_topk[0] + num_prefix_tiles, n_tiles) + idx = scores.topk(k_total, dim=-1).indices + mask.scatter_(-1, idx, True) + mask[:, :, :num_prefix_tiles, :] = True + return mask + + +def _h3_expand_mask_and_sizes_to_64(block_mask, variable_block_sizes, tile_elems): + R'''Logical tile selection -> the physical 64-block map the Triton kernels run at. + + Port of the official ``_expand_mask_and_sizes_256_to_64`` / ``_128_to_64`` (their "route A"), + generalized over ``r = tile_elems // 64``. A logical tile occupies ``r`` CONTIGUOUS 64-blocks in + the padded buffer, so selecting it selects all r sub-blocks on both the query and the key axis + (``repeat_interleave``), and its valid-token count chops into 64-token strides. The trailing + sub-blocks of a partial logical tile legitimately get size 0, which every kernel treats as + "contributes nothing" (the forward masks those columns to -inf, dQ skips them, dK/dV never + stores them). ``tile_elems == 64`` is the identity. + ''' + r = tile_elems // _H3_KERNEL_BLOCK + if r == 1: + return block_mask, variable_block_sizes + mask_64 = block_mask.repeat_interleave(r, dim=2).repeat_interleave(r, dim=3) + sizes = variable_block_sizes.to(torch.int32) + offsets = torch.arange(r, device=sizes.device, dtype=torch.int32) * _H3_KERNEL_BLOCK + sizes_64 = torch.clamp(sizes[:, None] - offsets[None, :], + min=0, max=_H3_KERNEL_BLOCK).reshape(-1) + return mask_64, sizes_64 + + +class VSASparseAttention(nn.Module): + R''' + Video Sparse Attention (VSA) -- the FastVideo/hao-ai-lab coarse-to-fine block-sparse + attention (arXiv:2505.13389) used by the official FastWan2.1 sparse-distilled students. + + Two branches, mirroring the released FastWan ``to_gate_compress`` design: + - sparse branch : mean-pool q/k into blocks -> coarse block scores -> keep the top-k + key blocks per query block -> block-sparse EXACT softmax over the + selected blocks. + - compressed branch: a global mean-pooled ("compressed") attention over ALL key blocks, + broadcast back to token resolution and scaled per token/head by the + learned gate ``to_gate_compress(x)`` (x = attention-input hidden + states). This is FastVideo's learned compensation for the context + lost to sparsification. ``to_gate_compress`` is zero-init, so a + checkpoint without those weights reduces to sparse-only:: + + out = o_sparse + to_gate_compress(x) * o_compressed + + New learnable params per layer (all under ``attn_op.*``): ``to_gate_compress`` weight+bias. + Absent from a dense checkpoint (loaded with strict=False); the official FastWan release + provides them. + + NOTE on fidelity: the block-sparse forward IS the official FastVideo Triton kernel (vendored), + the tile geometry / exempt mask / top-k are ports of the official H3 backend, and the tile + volume is selectable over the official {64, 128, 256} -- 128/256 reach the same 64-block kernel + through the official route-A expansion. + + ``forward`` accepts BSND (``[B, L, H, D]``) q/k/v like the sibling attention modules; + internal helpers use BHLD (``[B, H, L, *]``). + ''' + def __init__(self, head_dim, num_heads, dim, topk=0.2, + BLKQ=64, BLKK=64, use_bf16=True, tile_elems=_H3_TILE_ELEMS): + R''' + Args: + head_dim: dimension of each head. + num_heads: number of attention heads. + dim: hidden size of the attention input (sizes ``to_gate_compress``). + topk: keep-ratio of key blocks for the sparse branch (0.2 -> 80% sparse, + matching the official ``--VSA-sparsity 0.8``). + BLKQ/BLKK: query/key block sizes for the flat (non-segment) sparse kernel. + use_bf16: compute the sparse kernel in bfloat16 (default) or float16. + tile_elems: logical tile volume for the segment-aware path -- one of the official + ``VSA_H3_TILE_SHAPES`` keys (64 / 128 / 256). Must match the volume the + checkpoint was distilled at; the released FastWan and FastH3 both use 64. + ''' + super().__init__() + self.head_dim = head_dim + self.num_heads = num_heads + self.dim = dim + self.topk = topk + self.BLKQ = BLKQ + self.BLKK = BLKK + self.dtype = torch.bfloat16 if use_bf16 else torch.float16 + if tile_elems not in _H3_TILE_SHAPES: + raise ValueError( + f"Unsupported VSA tile volume {tile_elems}; supported volumes are " + f"{sorted(_H3_TILE_SHAPES)} (the official VSA_H3_TILE_SHAPES).") + self.tile_elems = tile_elems + self.tile_shape = _H3_TILE_SHAPES[tile_elems] + + inner = num_heads * head_dim + # Learned gate for the compressed global branch (FastVideo ``to_gate_compress``). + self.to_gate_compress = nn.Linear(dim, inner, bias=True) + self.init_weights_() + + def init_weights_(self): + with torch.no_grad(): + nn.init.zeros_(self.to_gate_compress.weight) + nn.init.zeros_(self.to_gate_compress.bias) + + def sparse(self, q, k, v): + R'''Block-sparse exact softmax (VSA fine stage). + + Args: + q, k: [B, H, L, D] POST-RoPE queries/keys. + v: [B, H, L, D] values. + Returns: + [B, H, L, D], or ``None`` when the sequence is shorter than one block. + ''' + L = q.shape[2] + if L < max(self.BLKQ, self.BLKK): + return None + orig_dtype = q.dtype + with torch.no_grad(): + pq = self._pool_flat(q, self.BLKQ) + pk = self._pool_flat(k, self.BLKK) + scores = pq @ pk.transpose(-1, -2) + real_topk = max(1, min(scores.shape[-1], math.ceil(self.topk * scores.shape[-1]))) + lut = scores.topk(real_topk, dim=-1, sorted=False).indices.contiguous() + sparse_map = torch.zeros_like(scores, dtype=torch.bool) + sparse_map.scatter_(-1, lut, True) + if self.BLKQ % 64 or self.BLKK % 64: + raise ValueError("VSA ็š„ flat block ๅคงๅฐๅฟ…้กปๆ˜ฏ 64 ็š„ๆญฃๆ•ดๆ•ฐๅ€ใ€‚") + B, H, L, D = q.shape + n_tiles = math.ceil(L / 64) + padding = n_tiles * 64 - L + q, k, v = (F.pad(t.to(self.dtype), (0, 0, 0, padding)).contiguous() for t in (q, k, v)) + sizes = (L - torch.arange(n_tiles, device=q.device) * 64).clamp(max=64).to(torch.int32) + mask = sparse_map.repeat_interleave(self.BLKQ // 64, dim=-2).repeat_interleave( + self.BLKK // 64, dim=-1)[..., :n_tiles, :n_tiles].contiguous() + o_s = sparse_linear_attention_varlen(q, k, v, mask, sizes) + return o_s[:, :, :L].to(orig_dtype) + + @staticmethod + def _pool_flat(x, block_size): + B, H, L, D = x.shape + n_blocks = math.ceil(L / block_size) + padded = F.pad(x.float(), (0, 0, 0, n_blocks * block_size - L)) + counts = (L - torch.arange(n_blocks, device=x.device) * block_size).clamp(max=block_size) + return padded.reshape(B, H, n_blocks, block_size, D).sum(3) / counts.view(1, 1, -1, 1) + + def compressed(self, q, k, v): + R'''Mean-pooled global ("compressed") attention, broadcast to token resolution. + + Args: + q, k, v: [B, H, L, D] POST-RoPE queries/keys/values. + Returns: + [B, H, L, D] compressed global context per token. + ''' + B, H, L, D = q.shape + pq = self._pool_flat(q, self.BLKQ) + pk = self._pool_flat(k, self.BLKK) + pv = self._pool_flat(v, self.BLKK) + scale = D ** -0.5 + scores = (pq.float() @ pk.float().transpose(-1, -2)) * scale + attn = torch.softmax(scores, dim=-1).to(pv.dtype) + oc_blocks = attn @ pv # [B,H,Lq,D] + # Broadcast each query block's context to its BLKQ tokens, then trim to L. + oc = oc_blocks.repeat_interleave(self.BLKQ, dim=2)[:, :, :L] + return oc + + def forward_h3(self, q, k, v, x, segments, return_sparsity=False): + R'''Segment-aware H3 VSA -- the official FastVideo segment/tile mask semantics. + + Args: + q, k, v: [B, L, H, D] POST-RoPE queries/keys/values in PACKED order. + x: [B, L, C] attention-input hidden states (drives the compress gate). + segments: packed layout -- dense row counts (``int``) and video region + ``(t, h, w)`` token grid, in packed order. Built by the model forward. + Returns: + [B, L, H, D] attention output (before the output projection). + + Tiles the packed sequence into [prefix tiles][video cubic tiles] of ``self.tile_shape``, + builds the exempt block mask on pooled tile scores, expands that logical selection down to + the 64-block granularity the kernels run at, runs the block-sparse Triton kernel for the + fine branch, adds the gated compressed branch on the same logical geometry, then untiles + back to packed order. + ''' + B, L, H, D = q.shape + device = q.device + (tile_partition_indices, variable_block_sizes, untile_combined_index, + num_prefix_tiles, num_video_tiles, video_tile_spans, tile_elems) = \ + build_h3_tile_geometry(tuple(segments), device, self.tile_shape) + n_tiles = num_prefix_tiles + num_video_tiles + n_regions = len(video_tile_spans) + # Logical tiles are laid end to end with no spare: the kernels take a per-64-block valid + # count, so a partial tile needs no zero dummy partner tile any more. + L_buf = n_tiles * tile_elems + sparsity = 1.0 - self.topk # self.topk is the KEEP ratio + vbs = variable_block_sizes.to(device) + + reuse = not torch.is_grad_enabled() + buf_shape = (B, H, L_buf, D) + + def _tile(t, dtype, role): + # [B,L,H,D] packed -> [B,H,L_buf,D] tile order (pad slots stay zero); under no-grad the + # buffer is shared across layers/steps (see _h3_tile_into_buf). + return _h3_tile_into_buf(t, dtype, untile_combined_index, buf_shape, + device, role, reuse) + + qb = _tile(q, self.dtype, 'q') + kb = _tile(k, self.dtype, 'k') + vb = _tile(v, self.dtype, 'v') + gate = self.to_gate_compress(x).view(B, L, H, D) + gb = _tile(gate, gate.dtype, 'g') + + pq = _h3_pool_tiles(qb, vbs, tile_elems, n_tiles) # [B,H,n_tiles,D] fp32 + pk = _h3_pool_tiles(kb, vbs, tile_elems, n_tiles) + pv = _h3_pool_tiles(vb, vbs, tile_elems, n_tiles) + + # Block selection is a non-differentiable constant (exactly like the flat path's + # get_block_map): build the mask under no_grad on detached pooled scores so no graph is + # retained for the top-k. + with torch.no_grad(): + scores = torch.matmul(pq.detach(), pk.detach().transpose(-2, -1)) / (D ** 0.5) + mask = _h3_build_block_mask(scores, num_prefix_tiles, sparsity, True, + video_tile_spans, (sparsity,) * n_regions) + # Logical (tile_elems) selection -> the physical 64-block map + per-block counts. + # Identity when tile_elems == 64. + block_mask, block_sizes = _h3_expand_mask_and_sizes_to_64(mask, vbs, tile_elems) + o_s = sparse_linear_attention_varlen( + qb, kb, vb, block_mask, block_sizes, qk_scale=D ** -0.5) + + # gated compressed branch on the same logical tile geometry + cs = torch.matmul(pq, pk.transpose(-2, -1)) * (D ** -0.5) + attn = torch.softmax(cs, dim=-1).to(pv.dtype) + oc = (attn @ pv).repeat_interleave(tile_elems, dim=2) # [B,H,L_buf,D] + out_buf = o_s.to(gb.dtype) + gb * oc.to(gb.dtype) + + out = out_buf[:, :, untile_combined_index].transpose(1, 2) # [B,L,H,D] packed + out = out.to(q.dtype) + if return_sparsity: + return out, self.topk + return out + + def forward(self, q, k, v, x, return_sparsity=False, segments=None): + R''' + Args: + q, k, v: [B, L, H, D] POST-RoPE queries/keys/values. + x: [B, L, C] attention-input hidden states (drives the compress gate). + return_sparsity: also return the sparse-branch keep ratio. + Returns: + [B, L, H, D] attention output (before the output projection). + ''' + if segments is not None: + # Segment-aware tile geometry with a differentiable variable-length block-sparse + # kernel (forward + backward), so training and inference share one selection + # semantics. Falls back to the flat 1-D block map only when the caller cannot + # supply a clean segment layout (segments is None). + return self.forward_h3(q, k, v, x, segments, return_sparsity=return_sparsity) + B, L, H, D = q.shape + qb = q.transpose(1, 2).contiguous() # [B,H,L,D] + kb = k.transpose(1, 2).contiguous() + vb = v.transpose(1, 2).contiguous() + + o_s = self.sparse(qb, kb, vb) + if o_s is None: + # Shorter than one block: exact dense attention (sparsity gives no benefit). + o_s = F.scaled_dot_product_attention( + qb.float(), kb.float(), vb.float()).to(qb.dtype) + gate = self.to_gate_compress(x).view(B, L, H, D).transpose(1, 2) # [B,H,L,D] + o_c = self.compressed(qb, kb, vb) + out = o_s + gate.to(o_s.dtype) * o_c.to(o_s.dtype) + out = out.transpose(1, 2) # [B,L,H,D] + if return_sparsity: + return out, self.topk + return out \ No newline at end of file diff --git a/videox_fun/models/fastvideo_minimax_h3_transformer3d.py b/videox_fun/models/fastvideo_minimax_h3_transformer3d.py new file mode 100644 index 0000000..ffc573f --- /dev/null +++ b/videox_fun/models/fastvideo_minimax_h3_transformer3d.py @@ -0,0 +1,256 @@ +# FastVideo VSA (Video Sparse Attention) student backbone for MiniMax-H3. Mirrors the Wan-side +# integration in ``fastvideo_wan_transformer3d.py``: it swaps the dense softmax self-attention of +# ``MiniMaxH3Transformer3DModel`` for the coarse-to-fine block-sparse VSA kernel (reusing the +# in-repo SLA Triton kernel) plus FastVideo's learned ``to_gate_compress`` compressed-global +# compensation branch. +# +# MiniMax-H3 runs full self-attention over ONE packed multi-modal sequence (video + text + audio +# rows) through a swappable attention *processor*, so the VSA op is wired in at the processor +# level (``FastVideoMiniMaxH3AttnProcessor``) and the trainable ``attn_op`` gate is attached to +# each main block's ``MiniMaxH3Attention`` module. The token refiner keeps the dense processor. +# +# VSA reference: Zhang et al., "VSA: Faster Video Diffusion with Trainable Sparse Attention", +# arXiv:2505.13389 (FastVideo / hao-ai-lab). +import torch +import torch.nn.functional as F + +from .attention_utils import VSASparseAttention, attention +from .minimax_h3_transformer3d import (MiniMaxH3Transformer3DModel, + apply_minimax_h3_rotary_emb) + + +class FastVideoMiniMaxH3AttnProcessor: + r"""MiniMax-H3 attention processor whose dense softmax kernel is swapped for VSA block-sparse + attention (sparse fine branch + gated mean-pooled compressed branch). + + Mirrors ``MiniMaxH3AttnProcessor.__call__`` exactly, except the padless single-document path + ``attention(q, k, v, causal=False)`` becomes the VSA op ``attn.attn_op(q, k, v, x)``. Two + cases keep the original DENSE attention: + - ``attention_mask is not None``: the packed sequence carries padding rows kept in their own + attention document; the VSA kernel does not accept an arbitrary mask, so fall back to the + masked SDPA path (identical to the base processor). + - ``attn.topk >= 1.0``: sparsity disabled -> the base dense full-softmax path. + """ + + def __call__( + self, + attn, + hidden_states: torch.Tensor, + rotary_emb=None, + attention_mask=None, + valid_length=None, + vsa_segments=None, + ) -> torch.Tensor: + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + if rotary_emb is not None: + query = apply_minimax_h3_rotary_emb(query, *rotary_emb) + key = apply_minimax_h3_rotary_emb(key, *rotary_emb) + + # VSA on the padless single-document path only. query/key/value are [B, L, H, D] POST-RoPE + # and hidden_states is the [B, L, C] attention input that drives the compress gate; the op + # returns [B, L, H, D], matching the dense paths below so the common tail is unchanged. + if attention_mask is None and getattr(attn, "topk", 1.0) < 1.0: + hidden_states = attn.attn_op( + query, key, value, hidden_states, segments=vsa_segments) + elif attention_mask is None: + hidden_states = attention(query, key, value, causal=False) + else: + hidden_states = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=attention_mask[None, None], + ).transpose(1, 2) + + hidden_states = hidden_states.flatten(2, 3).type_as(query) + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class FastVideoMiniMaxH3Transformer3DModel(MiniMaxH3Transformer3DModel): + r"""MiniMax-H3 transformer (joint video + audio) with a VSA sparse self-attention on the main + block stack. Block type stays ``MiniMaxH3TransformerBlock``, so the existing + ``--fsdp_transformer_layer_cls_to_wrap=MiniMaxH3TransformerBlock`` policy keeps working. + + """ + + # NOTE: diffusers from_pretrained filters config.json keys by __init__ signature, so every + # parent parameter must be declared explicitly (a *args/**kwargs passthrough would silently + # drop hidden_size/num_layers/... and build a wrongly-sized model). + def __init__( + self, + num_attention_heads: int = 56, + attention_head_dim: int = 128, + hidden_size: int = 5376, + num_layers: int = 50, + num_refiner_layers: int = 2, + ffn_dim: int = 14336, + in_channels: int = 24, + audio_in_channels: int = 32, + patch_size=(1, 2, 2), + text_dim: int = 5120, + freq_dim: int = 256, + time_embed_hidden_dim: int = 5376, + time_embed_dim: int = 2688, + rope_freq_dim: int = 16, + rope_theta: float = 10000.0, + norm_eps: float = 1e-5, + qk_norm_eps: float = 1e-5, + final_norm_eps: float = 1e-5, + student_topk=0.2, + ): + super().__init__( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + hidden_size=hidden_size, + num_layers=num_layers, + num_refiner_layers=num_refiner_layers, + ffn_dim=ffn_dim, + in_channels=in_channels, + audio_in_channels=audio_in_channels, + patch_size=patch_size, + text_dim=text_dim, + freq_dim=freq_dim, + time_embed_hidden_dim=time_embed_hidden_dim, + time_embed_dim=time_embed_dim, + rope_freq_dim=rope_freq_dim, + rope_theta=rope_theta, + norm_eps=norm_eps, + qk_norm_eps=qk_norm_eps, + final_norm_eps=final_norm_eps, + ) + self.student_topk = student_topk + self.register_to_config(student_topk=student_topk) + self._apply_sparse_self_attention(student_topk) + + def _apply_sparse_self_attention(self, topk): + # Sparsify only the main block stack (the heavy video+audio attention); the text token + # refiner keeps the dense processor. The VSA op is attached to the existing attention + # module (an nn.Module) so its trainable gate lands in the model's parameter tree, and the + # processor is swapped to route the padless path through it. + for block in self.transformer_blocks: + attn = block.attn + attn.topk = topk + attn.set_processor(FastVideoMiniMaxH3AttnProcessor()) + if topk >= 1.0: + # Dense fallback: drop any residual VSA submodule (e.g. one left on meta by an + # earlier low_cpu_mem_usage __init__) so no stale meta params reach the offloader. + if hasattr(attn, "attn_op"): + del attn.attn_op + continue + # ``dim`` sizes to_gate_compress's input, which is the attention-input hidden states + # (hidden_size = to_q.in_features), NOT inner_dim (heads * head_dim). + attn_op = VSASparseAttention( + attn.head_dim, attn.heads, attn.to_q.in_features, topk=topk) + # Under low_cpu_mem_usage from_pretrained the backbone params are meta tensors and the + # freshly built gate is materialized by a later set_student_topk (called after + # from_pretrained) or the checkpoint load; only align device/dtype when concrete, and + # match the backbone dtype so bf16 weights and an fp32 gate never mix under FSDP. + p0 = attn.to_q.weight + if not p0.is_meta: + attn_op = attn_op.to(device=p0.device, dtype=p0.dtype) + attn.attn_op = attn_op + + def set_student_topk(self, topk): + self.student_topk = topk + self.register_to_config(student_topk=topk) + if topk < 1.0 and all(hasattr(block.attn, "attn_op") and + not any(p.is_meta for p in block.attn.attn_op.parameters()) + for block in self.transformer_blocks): + for block in self.transformer_blocks: + block.attn.topk = topk + block.attn.attn_op.topk = topk + else: + self._apply_sparse_self_attention(topk) + + def forward(self, hidden_states, audio_hidden_states, encoder_hidden_states, + timestep, timestep_indices, token_tags, position_ids, + video_indices, audio_indices, text_indices, + attention_kwargs=None, return_dict=True): + attention_kwargs = dict(attention_kwargs or {}) + attention_kwargs["vsa_segments"] = self._compute_h3_segments(token_tags, position_ids) + return super().forward( + hidden_states, audio_hidden_states, encoder_hidden_states, timestep, + timestep_indices, token_tags, position_ids, video_indices, audio_indices, + text_indices, attention_kwargs=attention_kwargs, return_dict=return_dict) + + @staticmethod + def _video_grid_from_positions(pos): + """Recover the (T, H, W) token grid of one contiguous video run from its rotary + ``position_ids`` [n, 3]. The video rows are frame-major then (h, w) raster with w + fastest: frames carry a distinct temporal coord, and within a frame the leading run + of an identical height coord spans exactly one width row. Returns None on mismatch.""" + n = pos.shape[0] + if n == 0: + return None + t_col = pos[:, 0] + rows_per_frame = int((t_col == t_col[0]).sum()) # frame 0's temporal coord is unique + if rows_per_frame <= 0 or n % rows_per_frame != 0: + return None + h_col = pos[:rows_per_frame, 1] + W = int((h_col == h_col[0]).to(torch.int32).cumprod(0).sum()) # one height row + if W <= 0 or rows_per_frame % W != 0: + return None + H = rows_per_frame // W + T = n // rows_per_frame + if T * H * W != n: + return None + return (T, H, W) + + def _compute_h3_segments(self, token_tags, position_ids): + """Packed-order segment list for ``build_h3_tile_geometry``: dense run lengths + (int) plus one (t, h, w) triple per video region, in packed order. Supports + multiple video regions (Ref2VA / PDD reference videos sit as extra tag-0 runs and + each becomes its own sparse region, exempt top-k per region). Returns None (flat + fallback) when the sequence is padded or any video run's grid cannot be recovered.""" + if token_tags.numel() == 0 or bool((token_tags < 0).any()): + return None + tags = token_tags.detach().to(torch.int64).cpu() + change = torch.nonzero(tags[1:] != tags[:-1]).flatten() + 1 + bounds = torch.cat([torch.zeros(1, dtype=change.dtype), change, + torch.tensor([tags.numel()], dtype=change.dtype)]) + bounds = [int(b) for b in bounds.tolist()] + segments = [] + video_runs = 0 + for i in range(len(bounds) - 1): + s, e = bounds[i], bounds[i + 1] + if e <= s: + continue + if int(tags[s]) == 0: # video run + video_runs += 1 + grid = self._video_grid_from_positions(position_ids[s:e]) + if grid is None: + return None + segments.append(grid) + else: # dense (text / audio) run + segments.append(e - s) + if video_runs < 1: + return None + return tuple(segments) + + def enable_multi_gpus_inference(self,): + """VSA has no sequence-parallel path in this repo. + + The in-repo VSA kernel only implements the single-GPU coarse-to-fine block-sparse forward; + there is no ulysses head-parallel or ring-attention binding for it (unlike the base + ``MiniMaxH3MultiGPUsAttnProcessor``). Fail loudly rather than silently falling back to a + DENSE multi-GPU processor (which would bypass VSA and drop the gate params). + """ + raise NotImplementedError( + "FastVideoMiniMaxH3Transformer3DModel (VSA) supports single-GPU inference only; the " + "VSA kernel has no ulysses/ring sequence-parallel path. Set ulysses_degree=ring_degree=1.") + + def sparse_trainable_params(self): + """Names of the extra VSA-branch parameters, for logging / freezing.""" + return [name for name, _ in self.named_parameters() if "attn_op." in name] diff --git a/videox_fun/models/fastvideo_wan_transformer3d.py b/videox_fun/models/fastvideo_wan_transformer3d.py new file mode 100644 index 0000000..8a69e83 --- /dev/null +++ b/videox_fun/models/fastvideo_wan_transformer3d.py @@ -0,0 +1,204 @@ +# FastVideo VSA (Video Sparse Attention) student backbone for the Wan2.1 single-expert +# models. Mirrors the SparkDiffusion/RoLa integration in ``spark_transformer3d.py`` and the +# TurboDiffusion/SLA integration in ``turbowan_transformer3d.py``: it swaps the dense softmax +# self-attention of ``WanTransformer3DModel`` for the coarse-to-fine block-sparse VSA kernel +# (reusing the in-repo SLA Triton kernel) plus FastVideo's learned ``to_gate_compress`` +# compressed-global compensation branch. +# +# VSA reference: Zhang et al., "VSA: Faster Video Diffusion with Trainable Sparse Attention", +# arXiv:2505.13389 (FastVideo / hao-ai-lab). Official weights: FastVideo/FastWan2.1-T2V-1.3B. +import torch + +from .attention_utils import VSASparseAttention, attention +from .wan_transformer3d import (WanSelfAttention, WanTransformer3DModel, + rope_apply_qk) + + +class FastVideoWanSelfAttention(WanSelfAttention): + """Wan self-attention whose dense softmax kernel is swapped for VSA block-sparse attention + (sparse fine branch + gated mean-pooled compressed branch).""" + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6, + topk=0.2): + super().__init__(dim, num_heads, window_size, qk_norm, eps) + # topk: keep-ratio of key blocks per query (0.2 -> 80% sparse, the official FastWan + # ``--VSA-sparsity 0.8``). >= 1.0 falls back to the original dense attention. + self.topk = topk + if topk < 1.0: + self.attn_op = VSASparseAttention( + self.head_dim, num_heads, dim, topk=topk) + + def forward(self, x, seq_lens, grid_sizes, freqs, dtype=torch.bfloat16, t=0): + r""" + Args: + x(Tensor): Shape [B, L, num_heads * head_dim] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d) + k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d) + v = self.v(x.to(dtype)).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + # Dense fallback (topk >= 1.0): original full-softmax attention. + if self.topk >= 1.0: + q, k = rope_apply_qk(q, k, grid_sizes, freqs) + x = attention( + q.to(dtype), + k.to(dtype), + v=v.to(dtype), + k_lens=seq_lens, + window_size=self.window_size + ) + x = x.to(dtype).flatten(2) + return self.o(x) + + q, k = rope_apply_qk(q, k, grid_sizes, freqs) + + g0 = grid_sizes[0] + if (bool((grid_sizes == g0).all()) and s == int(g0.prod()) + and bool((seq_lens == s).all())): + segments = ((int(g0[0]), int(g0[1]), int(g0[2])),) + out = self.attn_op(q.to(dtype), k.to(dtype), v.to(dtype), x.to(dtype), + segments=segments) + else: + outputs = [] + for i, (grid, length) in enumerate(zip(grid_sizes.tolist(), seq_lens.tolist())): + grid = tuple(int(axis) for axis in grid) + length = int(length) + if min(grid) <= 0 or length != grid[0] * grid[1] * grid[2] or not 0 < length <= s: + raise ValueError("Wan VSA ็š„ seq_lens ๅฟ…้กป็ญ‰ไบŽๆœ‰ๆ•ˆ็ฝ‘ๆ ผ token ๆ•ฐไธ”ไธ่ถ…่ฟ‡ๅบๅˆ—้•ฟๅบฆใ€‚") + oi = self.attn_op(q[i:i + 1, :length].to(dtype), k[i:i + 1, :length].to(dtype), + v[i:i + 1, :length].to(dtype), x[i:i + 1, :length].to(dtype), + segments=(grid,)) + outputs.append(torch.nn.functional.pad(oi, (0, 0, 0, 0, 0, s - length))) + out = torch.cat(outputs, dim=0) + x = out.to(dtype).flatten(2) + return self.o(x) + + +class FastVideoWanTransformer3DModel(WanTransformer3DModel): + """Wan2.1 transformer (T2V / I2V, single expert) with a VSA sparse self-attention. Block + type stays WanAttentionBlock, so the existing + `--fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock` policy keeps working.""" + + # NOTE: diffusers from_pretrained filters config.json keys by __init__ signature, so + # every parent parameter must be declared explicitly (a *args/**kwargs passthrough would + # silently drop dim/ffn_dim/... and build a wrongly-sized model). + def __init__( + self, + model_type='t2v', + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + in_channels=16, + hidden_size=2048, + add_control_adapter=False, + in_dim_control_adapter=24, + downscale_factor_control_adapter=8, + add_ref_conv=False, + in_dim_ref_conv=16, + cross_attn_type=None, + student_topk=0.2, + ): + super().__init__( + model_type=model_type, + patch_size=patch_size, + text_len=text_len, + in_dim=in_dim, + dim=dim, + ffn_dim=ffn_dim, + freq_dim=freq_dim, + text_dim=text_dim, + out_dim=out_dim, + num_heads=num_heads, + num_layers=num_layers, + window_size=window_size, + qk_norm=qk_norm, + cross_attn_norm=cross_attn_norm, + eps=eps, + in_channels=in_channels, + hidden_size=hidden_size, + add_control_adapter=add_control_adapter, + in_dim_control_adapter=in_dim_control_adapter, + downscale_factor_control_adapter=downscale_factor_control_adapter, + add_ref_conv=add_ref_conv, + in_dim_ref_conv=in_dim_ref_conv, + cross_attn_type=cross_attn_type, + ) + self.student_topk = student_topk + self.register_to_config(student_topk=student_topk) + self._apply_sparse_self_attention(student_topk) + + def _apply_sparse_self_attention(self, topk): + for block in self.blocks: + dense_attn = block.self_attn + sparse_attn = FastVideoWanSelfAttention( + dense_attn.dim, dense_attn.num_heads, dense_attn.window_size, + dense_attn.qk_norm, dense_attn.eps, topk=topk, + ) + # Under low_cpu_mem_usage from_pretrained the replaced params are meta tensors, so + # init_weights_' zeros carry no data: wan_transformer3d.py's loader re-materialises + # every key the checkpoint lacks, and `to_gate_compress` is on its zero-init + # whitelist, so a cold start from a dense base still gets a zero gate. Only upgrade + # in place (copy the real tensors over) when the weights are concrete. + if not next(dense_attn.parameters()).is_meta: + old_sd = {k: v for k, v in dense_attn.state_dict().items() if not v.is_meta} + # Match the live params' device AND dtype: freshly built VSA params are fp32, + # and mixing them with bf16 backbone weights breaks FSDP flat-param flattening. + p0 = next(dense_attn.parameters()) + sparse_attn = sparse_attn.to(device=p0.device, dtype=p0.dtype) + sparse_attn.load_state_dict(old_sd, strict=False) + block.self_attn = sparse_attn + + def set_student_topk(self, topk): + self.student_topk = topk + self.register_to_config(student_topk=topk) + if topk < 1.0 and all(hasattr(block.self_attn, "attn_op") and + not any(p.is_meta for p in block.self_attn.attn_op.parameters()) + for block in self.blocks): + for block in self.blocks: + block.self_attn.topk = topk + block.self_attn.attn_op.topk = topk + else: + self._apply_sparse_self_attention(topk) + + def enable_multi_gpus_inference(self,): + """VSA has no sequence-parallel path in this repo. + + The in-repo VSA kernel only implements the single-GPU coarse-to-fine block-sparse + forward; there is no ulysses head-parallel or ring-attention binding for it (unlike + the RoLa/SLA ``usp_attn_*`` forwards). The 1.3B student fits comfortably on one GPU, + so run inference single-device. Fail loudly rather than silently falling back to the + parent's DENSE ``usp_attn_forward`` (which would bypass VSA and drop the gate params). + """ + raise NotImplementedError( + "FastVideoWanTransformer3DModel (VSA) supports single-GPU inference only; the VSA " + "kernel has no ulysses/ring sequence-parallel path. Set ulysses_degree=ring_degree=1.") + + def sparse_trainable_params(self): + """Names of the extra VSA-branch parameters, for logging / freezing.""" + return [name for name, _ in self.named_parameters() if "attn_op." in name] diff --git a/videox_fun/models/wan_gan_discriminator.py b/videox_fun/models/wan_gan_discriminator.py new file mode 100644 index 0000000..c73a06f --- /dev/null +++ b/videox_fun/models/wan_gan_discriminator.py @@ -0,0 +1,108 @@ +# Modified from https://github.com/Aurora-edu/One-Forcing/blob/main/wan/modules/model.py +r"""One-Forcing GAN discriminator head for the Wan fake score. + +The discriminator shares the fake-score backbone: one learnable register token per selected block cross-attends to +that block's hidden states, and the concatenated tokens go through a small classification branch to a single logit. +[`add_wan_gan_discriminator`] attaches the head as the `gan_head` submodule of a [`WanTransformer3DModel`] and wraps its +forward, so calling the model with `gan_classify=True` returns `(flow_pred, logits)`. The head runs inside the forward, +so FSDP shards and trains it together with the backbone. +""" + +import math +import types + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .wan_transformer3d import WanLayerNorm, WanRMSNorm + + +class GanAttentionBlock(nn.Module): + """Register-token cross-attention block of the One-Forcing discriminator head.""" + + def __init__(self, dim, ffn_dim, num_heads, eps=1e-6): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) + self.norm_k = WanRMSNorm(dim, eps=eps) + self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)) + + def forward(self, x, token): + b, n, d = x.size(0), self.num_heads, self.head_dim + x = self.norm3(x) + q = self.norm_q(self.q(token)).view(b, -1, n, d).transpose(1, 2) + k = self.norm_k(self.k(x)).view(b, -1, n, d).transpose(1, 2) + v = self.v(x).view(b, -1, n, d).transpose(1, 2) + token = token + self.o(F.scaled_dot_product_attention(q, k, v).transpose(1, 2).flatten(2)) + return token + self.ffn(self.norm2(token)) + + +class GanDiscriminatorHead(nn.Module): + """One-Forcing discriminator head: one register token per fake_score feature layer.""" + + def __init__(self, dim, num_layers, ffn_dim, num_heads): + super().__init__() + self.register_tokens = nn.Parameter(torch.randn(num_layers, dim) * 0.02) + self.register_norm = WanRMSNorm(dim, eps=1e-6) + self.blocks = nn.ModuleList([GanAttentionBlock(dim, ffn_dim, num_heads) for _ in range(num_layers)]) + self.cls_pred_branch = nn.Sequential( + nn.LayerNorm(dim * num_layers), nn.Linear(dim * num_layers, dim), nn.SiLU(), nn.Linear(dim, 1) + ) + + def forward(self, features): + registers = self.register_norm(self.register_tokens) + tokens = [] + for i, (block, feature) in enumerate(zip(self.blocks, features)): + token = registers[i:i + 1].unsqueeze(0).expand(feature.size(0), -1, -1) + # Checkpoint each cross-attention block so its q/k/v activations are recomputed during + # backward instead of being retained; one register token keeps the recompute cost tiny. + if torch.is_grad_enabled() and self.training: + token = torch.utils.checkpoint.checkpoint(block, feature, token, use_reentrant=False) + else: + token = block(feature, token) + tokens.append(token) + return self.cls_pred_branch(torch.cat(tokens, dim=1).flatten(1)) + + +def add_wan_gan_discriminator(transformer, feature_layers, ffn_dim, num_heads, dtype, seed=0): + """ + Attach the One-Forcing discriminator head to fake_score before accelerator.prepare. + Calling the model with gan_classify=True returns (flow_pred, logits); the head runs + inside the wrapped forward so it is sharded and trained together with fake_score. + """ + feature_layers = sorted(set(feature_layers)) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + transformer.gan_head = GanDiscriminatorHead(transformer.dim, len(feature_layers), ffn_dim, num_heads) + transformer.gan_head.to(device=next(transformer.parameters()).device, dtype=dtype) + transformer.gan_head.requires_grad_(True) + transformer._gan_features = None + + def capture_feature(module, inputs, output): + if transformer._gan_features is not None: + transformer._gan_features.append(output) + + for layer in feature_layers: + transformer.blocks[layer].register_forward_hook(capture_feature) + + def forward(self, *args, gan_classify=False, **kwargs): + if not gan_classify: + return type(self).forward(self, *args, **kwargs) + self._gan_features = [] + try: + flow = type(self).forward(self, *args, **kwargs) + features = self._gan_features + finally: + self._gan_features = None + num_tokens = math.prod(flow.shape[2:]) // math.prod(self.patch_size) + return flow, self.gan_head([feature[:, :num_tokens] for feature in features]) + + transformer.forward = types.MethodType(forward, transformer) diff --git a/videox_fun/models/wan_tae.py b/videox_fun/models/wan_tae.py index 6f8a9bd..2465197 100644 --- a/videox_fun/models/wan_tae.py +++ b/videox_fun/models/wan_tae.py @@ -13,8 +13,7 @@ import torch.nn.functional as F from collections import namedtuple from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.loaders.single_file_model import FromOriginalModelMixin -from diffusers.models.autoencoders.vae import (DecoderOutput, - DiagonalGaussianDistribution) +from diffusers.models.autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution from diffusers.models.modeling_outputs import AutoencoderKLOutput from diffusers.models.modeling_utils import ModelMixin from diffusers.utils.accelerate_utils import apply_forward_hook @@ -39,7 +38,9 @@ class Clamp(nn.Module): class MemBlock(nn.Module): def __init__(self, n_in, n_out): super().__init__() - self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out)) + self.conv = nn.Sequential( + conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out) + ) self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() self.act = nn.ReLU(inplace=True) @@ -49,11 +50,13 @@ class MemBlock(nn.Module): class SuperMemBlock(nn.Module): """MemBlock variant used by the Super decoder (ConvNeXt-style: 7x7 depthwise conv + inverted bottleneck).""" + def __init__(self, n_f): super().__init__() self.conv = nn.Sequential( nn.Conv2d(n_f * 2, n_f * 2, 7, padding=3, groups=n_f * 2, bias=False), - nn.Conv2d(n_f * 2, n_f * 4, 1), nn.ReLU(inplace=True), + nn.Conv2d(n_f * 2, n_f * 4, 1), + nn.ReLU(inplace=True), nn.Conv2d(n_f * 4, n_f, 1, bias=False), ) @@ -208,7 +211,15 @@ class TAEWan(nn.Module): (~Gaussian) latents, matching the reference taehv implementation. """ - def __init__(self, encoder_time_downscale=(True, True, False), decoder_time_upscale=(False, True, True), decoder_space_upscale=(True, True, True), patch_size=1, latent_channels=16, arch_variant=None): + def __init__( + self, + encoder_time_downscale=(True, True, False), + decoder_time_upscale=(False, True, True), + decoder_space_upscale=(True, True, True), + patch_size=1, + latent_channels=16, + arch_variant=None, + ): """ Args: encoder_time_downscale: whether temporal downsampling is enabled for each block. @@ -226,29 +237,78 @@ class TAEWan(nn.Module): decoder_time_upscale = (False, *decoder_time_upscale) assert arch_variant in (None, "super"), f"unrecognized arch_variant {arch_variant!r}" self.encoder = nn.Sequential( - conv(self.image_channels * self.patch_size ** 2, 64), nn.ReLU(inplace=True), - TPool(64, 2 if encoder_time_downscale[0] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), - TPool(64, 2 if encoder_time_downscale[1] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), - TPool(64, 2 if encoder_time_downscale[2] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + conv(self.image_channels * self.patch_size ** 2, 64), + nn.ReLU(inplace=True), + TPool(64, 2 if encoder_time_downscale[0] else 1), + conv(64, 64, stride=2, bias=False), + MemBlock(64, 64), + MemBlock(64, 64), + MemBlock(64, 64), + TPool(64, 2 if encoder_time_downscale[1] else 1), + conv(64, 64, stride=2, bias=False), + MemBlock(64, 64), + MemBlock(64, 64), + MemBlock(64, 64), + TPool(64, 2 if encoder_time_downscale[2] else 1), + conv(64, 64, stride=2, bias=False), + MemBlock(64, 64), + MemBlock(64, 64), + MemBlock(64, 64), conv(64, self.latent_channels), ) if arch_variant == "super": n_f = [512, 256, 128, 64] self.decoder = nn.Sequential( nn.Conv2d(self.latent_channels, n_f[0], 1, bias=False), - SuperMemBlock(n_f[0]), SuperMemBlock(n_f[0]), SuperMemBlock(n_f[0]), conv(n_f[0], n_f[1] * (2 if decoder_space_upscale[0] else 1) ** 2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[0] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), - SuperMemBlock(n_f[1]), SuperMemBlock(n_f[1]), SuperMemBlock(n_f[1]), conv(n_f[1], n_f[2] * (2 if decoder_space_upscale[1] else 1) ** 2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[1] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), - SuperMemBlock(n_f[2]), SuperMemBlock(n_f[2]), SuperMemBlock(n_f[2]), conv(n_f[2], n_f[3] * (2 if decoder_space_upscale[2] else 1) ** 2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[2] else 1), TGrow(n_f[3], 2 if decoder_time_upscale[2] else 1), + SuperMemBlock(n_f[0]), + SuperMemBlock(n_f[0]), + SuperMemBlock(n_f[0]), + conv(n_f[0], n_f[1] * (2 if decoder_space_upscale[0] else 1) ** 2), + nn.ReLU(inplace=True), + nn.PixelShuffle(2 if decoder_space_upscale[0] else 1), + TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), + SuperMemBlock(n_f[1]), + SuperMemBlock(n_f[1]), + SuperMemBlock(n_f[1]), + conv(n_f[1], n_f[2] * (2 if decoder_space_upscale[1] else 1) ** 2), + nn.ReLU(inplace=True), + nn.PixelShuffle(2 if decoder_space_upscale[1] else 1), + TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), + SuperMemBlock(n_f[2]), + SuperMemBlock(n_f[2]), + SuperMemBlock(n_f[2]), + conv(n_f[2], n_f[3] * (2 if decoder_space_upscale[2] else 1) ** 2), + nn.ReLU(inplace=True), + nn.PixelShuffle(2 if decoder_space_upscale[2] else 1), + TGrow(n_f[3], 2 if decoder_time_upscale[2] else 1), conv(n_f[3], self.image_channels * self.patch_size ** 2), ) else: n_f = [256, 128, 64, 64] self.decoder = nn.Sequential( - Clamp(), conv(self.latent_channels, n_f[0]), nn.ReLU(inplace=True), - MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 2 if decoder_time_upscale[0] else 1), conv(n_f[0], n_f[1], bias=False), - MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[1] else 1), conv(n_f[1], n_f[2], bias=False), - MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[2] else 1), conv(n_f[2], n_f[3], bias=False), - nn.ReLU(inplace=True), conv(n_f[3], self.image_channels * self.patch_size ** 2), + Clamp(), + conv(self.latent_channels, n_f[0]), + nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), + MemBlock(n_f[0], n_f[0]), + MemBlock(n_f[0], n_f[0]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), + TGrow(n_f[0], 2 if decoder_time_upscale[0] else 1), + conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), + MemBlock(n_f[1], n_f[1]), + MemBlock(n_f[1], n_f[1]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), + TGrow(n_f[1], 2 if decoder_time_upscale[1] else 1), + conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), + MemBlock(n_f[2], n_f[2]), + MemBlock(n_f[2], n_f[2]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), + TGrow(n_f[2], 2 if decoder_time_upscale[2] else 1), + conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), + conv(n_f[3], self.image_channels * self.patch_size ** 2), ) # computed properties self.t_downscale = 2 ** sum(t.stride == 2 for t in self.encoder if isinstance(t, TPool)) @@ -268,7 +328,7 @@ class TAEWan(nn.Module): key = f"decoder.{i}.conv.weight" if key in sd and sd[key].shape[0] > new_sd[key].shape[0]: # take the last-timestep output channels - sd[key] = sd[key][-new_sd[key].shape[0]:] + sd[key] = sd[key][-new_sd[key].shape[0] :] return sd def preprocess_input_frames(self, x): @@ -293,6 +353,8 @@ class TAEWan(nn.Module): if False, frames will be processed sequentially. Returns NTCHW latent tensor with ~Gaussian values. """ + if isinstance(self.encoder, nn.Identity): + raise RuntimeError("ไป… Decoder ็š„ TAE ๆƒ้‡ไธๆ”ฏๆŒ encode๏ผ›่ฏทๅŠ ่ฝฝๅŒ…ๅซ Encoder ็š„ๅฎŒๆ•ดๆƒ้‡ใ€‚") x = self.preprocess_input_frames(x) if x.shape[1] % self.t_downscale != 0: # pad at end to multiple of self.t_downscale @@ -313,7 +375,7 @@ class TAEWan(nn.Module): """ x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar) x = self.postprocess_output_frames(x) - return x[:, self.frames_to_trim:] + return x[:, self.frames_to_trim :] class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): @@ -330,24 +392,142 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ # per-channel statistics of the full-size Wan VAE latent spaces - WAN2_1_MEAN = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, - 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921] - WAN2_1_STD = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, - 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160] + WAN2_1_MEAN = [ + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ] + WAN2_1_STD = [ + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ] WAN2_2_MEAN = [ - -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, - -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, - -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, - -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, - -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, - 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667] + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ] WAN2_2_STD = [ - 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, - 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, - 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, - 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, - 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, - 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744] + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ] @register_to_config def __init__( @@ -361,6 +541,7 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): temporal_compression_ratio=4, spatial_compression_ratio=None, parallel=True, + decoder_only=False, ): super().__init__() if latent_channels == 16: @@ -368,7 +549,9 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): elif latent_channels == 48: mean, std = self.WAN2_2_MEAN, self.WAN2_2_STD else: - raise ValueError(f"AutoencoderTinyWan supports 16ch (taew2_1) or 48ch (taew2_2) latents, got {latent_channels}") + raise ValueError( + f"AutoencoderTinyWan supports 16ch (taew2_1) or 48ch (taew2_2) latents, got {latent_channels}" + ) assert len(mean) == latent_channels and len(std) == latent_channels self.mean = torch.tensor(mean, dtype=torch.float32) self.std = torch.tensor(std, dtype=torch.float32) @@ -387,6 +570,8 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): latent_channels=latent_channels, arch_variant=arch_variant, ) + if decoder_only: + self.model.encoder = nn.Identity() self.parallel = parallel self.gradient_checkpointing = False @@ -406,12 +591,12 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): # latents_mean=zeros/latents_std=ones and passes pipeline latents # through unchanged), so no mean/std scaling is applied here: the # encoder output is already interchangeable with full-VAE latents. - x = (x.add(1).div_(2)).clamp_(0, 1) # [-1, 1] -> [0, 1] - x = x.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW + x = (x.add(1).div_(2)).clamp_(0, 1) # [-1, 1] -> [0, 1] + x = x.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW zs = [] for u in x: z = self.model.encode_video(u.unsqueeze(0), parallel=self.parallel).squeeze(0) - zs.append(z.permute(1, 0, 2, 3)) # TCHW -> CTHW + zs.append(z.permute(1, 0, 2, 3)) # TCHW -> CTHW return torch.stack(zs) @apply_forward_hook @@ -432,7 +617,7 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): # zs: BCLHW normalized latents -> BCTHW in [-1, 1]. # NOTE: no denormalization here, for the same reason as in _encode: # the official TAE decoder consumes normalized latents directly. - zs = zs.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW + zs = zs.permute(0, 2, 1, 3, 4) # BCTHW -> NTCHW dec = [] for u in zs: d = self.model.decode_video(u.unsqueeze(0), parallel=self.parallel).squeeze(0) @@ -450,48 +635,99 @@ class AutoencoderTinyWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): return DecoderOutput(sample=decoded) @classmethod - def from_pretrained(cls, pretrained_model_path, additional_kwargs={}, **kwargs): + def from_pretrained(cls, pretrained_model_path, additional_kwargs=None, **kwargs): 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 - # standard diffusers directory kwargs (e.g. subfolder= from load_model_hook) - subfolder = kwargs.get("subfolder", None) + pretrained_model_path = os.fspath(pretrained_model_path) + subfolder = kwargs.pop("subfolder", None) if subfolder: pretrained_model_path = os.path.join(pretrained_model_path, subfolder) - - kwargs = filter_kwargs(cls, additional_kwargs) - # support diffusers directory checkpoints (save_pretrained output) + model_kwargs = filter_kwargs(cls, additional_kwargs or {}) + model_kwargs.update(filter_kwargs(cls, kwargs)) if os.path.isdir(pretrained_model_path): if os.path.exists(os.path.join(pretrained_model_path, "config.json")): - return super().from_pretrained(pretrained_model_path, **kwargs) - candidates = sorted(f for f in os.listdir(pretrained_model_path) - if f.endswith((".pth", ".safetensors"))) - if not candidates: - raise ValueError(f"No TAE weights found in {pretrained_model_path}") + loader_kwargs = {k: v for k, v in kwargs.items() if k not in model_kwargs} + return_loading_info = loader_kwargs.pop("output_loading_info", False) + model, info = super().from_pretrained( + pretrained_model_path, **model_kwargs, **loader_kwargs, output_loading_info=True + ) + if any(info.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")): + raise RuntimeError(f"TAE ้…็ฝฎไธŽๆƒ้‡ไธๅŒน้…ๆˆ–ๆƒ้‡ไธๅฎŒๆ•ด๏ผš{info}") + return (model, info) if return_loading_info else model + candidates = sorted( + f for f in os.listdir(pretrained_model_path) if f.endswith((".pt", ".pth", ".bin", ".safetensors")) + ) + if len(candidates) != 1: + raise ValueError(f"่ฏทๆŒ‡ๅฎšๅ”ฏไธ€็š„ TAE ๆƒ้‡ๆ–‡ไปถ๏ผŒ็›ฎๅฝ•ๅ†…ๆ‰พๅˆฐ๏ผš{candidates}") pretrained_model_path = os.path.join(pretrained_model_path, candidates[0]) - # infer model flavour from the checkpoint filename when not given explicitly - checkpoint_name = pretrained_model_path.lower() - if "latent_channels" not in kwargs: - if "taew2_2" in checkpoint_name: - kwargs["latent_channels"], kwargs["patch_size"] = 48, 2 - else: - kwargs["latent_channels"], kwargs["patch_size"] = 16, 1 - if kwargs.get("arch_variant") is None and "_super" in checkpoint_name: - kwargs["arch_variant"] = "super" - model = cls(**kwargs) if pretrained_model_path.endswith(".safetensors"): from safetensors.torch import load_file + state_dict = load_file(pretrained_model_path) else: state_dict = torch.load(pretrained_model_path, map_location="cpu", weights_only=True) - # TAE checkpoints already use "encoder.*" / "decoder.*" keys, no prefix needed - state_dict = model.model.patch_tgrow_layers(state_dict) - m, u = model.model.load_state_dict(state_dict, strict=False) - print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};") - print(m, u) + if isinstance(state_dict, dict): + if isinstance(state_dict.get("state_dict"), dict): + state_dict = state_dict["state_dict"] + elif isinstance(state_dict.get("decoder"), dict): + state_dict = state_dict["decoder"] + if ( + not isinstance(state_dict, dict) + or not state_dict + or not all(isinstance(k, str) and isinstance(v, torch.Tensor) for k, v in state_dict.items()) + ): + raise ValueError("ๆ— ๆ•ˆ็š„ TAE ๆƒ้‡๏ผš้œ€่ฆ tensor state_dict ๆˆ– checkpoint['decoder']ใ€‚") + if all(k.startswith("model.") for k in state_dict): + state_dict = {k[len("model.") :]: v for k, v in state_dict.items()} + + 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) + elif "decoder.1.weight" in state_dict and "decoder.0.weight" not in state_dict: + arch_variant, first_key, block_key = None, "decoder.1.weight", "decoder.3.conv.0.weight" + first_channels, first_kernel, block_shape = 256, (3, 3), (256, 512, 3, 3) + else: + raise ValueError("ๆ— ๆณ•่ฏ†ๅˆซ TAE ๆžถๆž„๏ผšDecoder ้ฆ–ๅฑ‚ๆƒ้‡็ผบๅคฑๆˆ–ๅญ˜ๅœจๅ†ฒ็ชใ€‚") + first = state_dict[first_key] + block = state_dict.get(block_key) + last = state_dict.get("decoder.22.weight") + if ( + first.ndim != 4 + or first.shape[0] != first_channels + or tuple(first.shape[2:]) != first_kernel + or block is None + or tuple(block.shape) != block_shape + or last is None + or last.ndim != 4 + ): + raise ValueError("ๆ— ๆณ•่ฏ†ๅˆซ TAE ๆžถๆž„๏ผš้ฆ–ๅฑ‚ใ€่ฎฐๅฟ†ๅ—ๆˆ–่พ“ๅ‡บๅฑ‚็š„ๅฝข็Šถไธ็ฌฆๅˆๆ™ฎ้€š/Super TAEใ€‚") + latent_channels = first.shape[1] + patch_size = {16: 1, 48: 2}.get(latent_channels) + if patch_size is None or tuple(last.shape) != (3 * patch_size ** 2, 64, 3, 3): + raise ValueError("TAE latent ้€š้“ๆˆ–่พ“ๅ‡บ patch_size ไธ็ฌฆๅˆ Wan 2.1/2.2 ๆžถๆž„ใ€‚") + inferred = { + "arch_variant": arch_variant, + "latent_channels": latent_channels, + "patch_size": patch_size, + "decoder_only": not any(k.startswith("encoder.") for k in state_dict), + } + for key, value in inferred.items(): + if key in model_kwargs and model_kwargs[key] != value: + raise ValueError(f"TAE ๅ‚ๆ•ฐ {key}={model_kwargs[key]!r} ไธŽๆƒ้‡่ฏ†ๅˆซๅ€ผ {value!r} ๅ†ฒ็ชใ€‚") + model_kwargs[key] = value + model = cls(**model_kwargs) + state_dict = model.model.patch_tgrow_layers(dict(state_dict)) + model.model.load_state_dict(state_dict, strict=True) + torch_dtype = kwargs.get("torch_dtype") + if torch_dtype == "auto": + torch_dtype = first.dtype + if torch_dtype is not None: + model.to(dtype=torch_dtype) return model diff --git a/videox_fun/models/wan_transformer3d.py b/videox_fun/models/wan_transformer3d.py index acf7192..fa7bffa 100755 --- a/videox_fun/models/wan_transformer3d.py +++ b/videox_fun/models/wan_transformer3d.py @@ -1325,10 +1325,14 @@ class WanTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): elif ( ('audio_injector' in key and '.o.' in key) or ('injector_adain_layers' in key and '.linear.' in key) - or any(name in key for name in ('trainable_cond_mask', 'cond_encoder', 'zip_motion_out.1', 'proj_l')) + or any(name in key for name in ('trainable_cond_mask', 'cond_encoder', 'zip_motion_out.1', 'proj_l', 'to_gate_compress')) ): - # S2V zero_init_weights() / turbowan SparseLinearAttention.init_weights_() - # zero these modules in __init__, which is lost under init_empty_weights. + # S2V zero_init_weights() / turbowan SparseLinearAttention.init_weights_() / + # VSASparseAttention.init_weights_() zero these modules in __init__, which is + # lost under init_empty_weights. A VSA student cold-started from a dense base + # has no attn_op.to_gate_compress keys in the checkpoint, so without this + # branch the gate falls through to xavier_uniform_ below and the compressed + # global branch starts at full random magnitude instead of as a no-op. initialized_dict[key] = torch.zeros(param_shape, dtype=param_dtype) print(f"Initializing missing parameter '{key}' with zero.") elif 'weight' in key: