# Based on https://github.com/guandeh17/Self-Forcing #!/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 gc import logging import math import os import pickle import shutil import sys import time import accelerate import diffusers import numpy as np import torch import torch.nn.functional as F import torchvision.transforms.functional as TF import transformers from accelerate import Accelerator from accelerate.logging import get_logger 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.utils.data import RandomSampler from torch.utils.tensorboard import SummaryWriter from torchvision import transforms from tqdm.auto import tqdm from transformers import AutoTokenizer 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 ImageVideoSafetensorsDataset, RandomSampler from videox_fun.models import (AutoencoderKLWan, WanT5EncoderModel, WanTransformer3DModel_FlexForcing, WanTransformer3DModel_SelfForcing) from videox_fun.pipeline import WanFlexForcingPipeline, WanSelfForcingPipeline from videox_fun.utils.flex_chunking import chunk_boundaries, sample_flex_level0 from videox_fun.utils.tqdm_bar import PauseAwareTqdm from videox_fun.utils.utils import save_videos_grid check_min_version("0.18.0.dev0") logger = get_logger(__name__, log_level="INFO") # ============================================================================ # Utilities # ============================================================================ def filter_kwargs(cls, kwargs): import inspect sig = inspect.signature(cls.__init__) valid_params = set(sig.parameters.keys()) - {'self', 'cls'} return {k: v for k, v in kwargs.items() if k in valid_params} def log_validation(transformer3d, args, config, accelerator, weight_dtype, global_step, fsdp_stage=0, zero_stage=0): """Run validation. Loads tokenizer, text_encoder and vae lazily (only when called) and frees them after.""" text_encoder = None vae = None try: with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): logger.info("Running validation... ") # --- Lazy load tokenizer, text_encoder and vae just for this validation pass --- tokenizer = AutoTokenizer.from_pretrained( os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), ) logger.info("Loading text encoder for validation...") if zero_stage == 3: ctx = accelerate.state.AcceleratorState().deepspeed_plugin.zero3_init_context_manager(enable=False) else: from contextlib import nullcontext ctx = nullcontext() with ctx: 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() text_encoder.requires_grad_(False) logger.info("Loading VAE for validation...") 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']), ).to(weight_dtype) vae.eval() vae.requires_grad_(False) # Apply FSDP sharding to text_encoder if FSDP/DeepSpeed is enabled if fsdp_stage != 0 or zero_stage != 0: from functools import partial from videox_fun.dist import shard_model shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype) text_encoder = shard_fn(text_encoder) scheduler_kwargs = OmegaConf.to_container(config['scheduler_kwargs']) scheduler_kwargs['shift'] = args.shift scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, scheduler_kwargs) ) pipeline_cls = WanFlexForcingPipeline if args.flex_forcing else WanSelfForcingPipeline pipeline = pipeline_cls( 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) for i in range(len(args.validation_prompts)): if args.fix_sample_size is not None: height, width = args.fix_sample_size else: height, width = args.video_sample_size, args.video_sample_size sample = pipeline( args.validation_prompts[i], num_frames=args.video_sample_n_frames, negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", height=height, width=width, generator=generator, guidance_scale=1.0, num_inference_steps=len(args.denoising_step_indices_list), shift=args.shift, num_frame_per_block=args.num_frame_per_block, independent_first_frame=args.independent_first_frame, context_noise=args.context_noise, stochastic_sampling=True, # Validation reuses the inference pipeline. A *fixed* uniform # partition keeps samples comparable across checkpoints, # unlike the random draw used for training. **({'chunk_spec': args.num_frame_per_block} if args.flex_forcing else {}), ).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}-{i}.mp4") ) # --- Free text_encoder, vae and pipeline to release GPU memory --- del pipeline del text_encoder del vae text_encoder = None vae = None gc.collect() torch.cuda.empty_cache() except Exception as e: if text_encoder is not None: del text_encoder if vae is not None: del vae gc.collect() torch.cuda.empty_cache() print(f"Eval error on rank {accelerator.process_index} with info {e}") def get_timestep_for_ode( min_timestep, max_timestep, batch_size, num_frames, num_frame_per_block, independent_first_frame, device, generator=None, ): """ Generate random timestep indices per frame/block. Same timestep within each block (matches Self-Forcing ODE regression). Returns: [batch_size, num_frames] tensor of indices. """ timestep = torch.randint( min_timestep, max_timestep, [batch_size, num_frames], device=device, dtype=torch.long, generator=generator, ) if independent_first_frame: timestep_from_second = timestep[:, 1:] timestep_from_second = timestep_from_second.reshape( timestep_from_second.shape[0], -1, num_frame_per_block) timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1] timestep_from_second = timestep_from_second.reshape( timestep_from_second.shape[0], -1) timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1) else: timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block) timestep[:, :, 1:] = timestep[:, :, 0:1] timestep = timestep.reshape(timestep.shape[0], -1) return timestep def get_timestep_for_ode_flexible(min_timestep, max_timestep, batch_size, chunk_sizes, device, generator=None, clean_index=None): """`get_timestep_for_ode` for a Flex-Forcing partition (arXiv 2607.03509 3.1). Same contract - every frame of a chunk shares one timestep index - but the chunk widths come from the partition instead of a single scalar, so the reshape trick of the uniform version does not apply. ``clean_index`` is the index whose entry in ``train_denoising_step_list`` is t=0, or ``None`` when the schedule has no clean entry. When given, each sample is guaranteed at least one chunk that did *not* draw it. An all-clean draw makes the loss mask ``(timestep != 0)`` empty, and the caller's fallback ``mse(pred_x0, target)`` is then *exactly* zero rather than merely small: those frames gather the clean endpoint itself and sigma is pinned to 0, so ``pred_x0 == target`` bitwise. The iteration contributes no gradient at all. The probability is ``len(train_denoising_step_list) ** -K`` in the number of chunks K, so it is dominated by the coarsest layouts: the whole-clip ``[F]`` partition is a single chunk and on its own carries ``1 / len(train_denoising_step_list)`` of iterations. A single chunk also has no later block to serve as clean context *for*, which is the only reason the clean index exists, so redrawing it loses nothing. """ num_frames = sum(int(c) for c in chunk_sizes) per_chunk = torch.randint( min_timestep, max_timestep, [batch_size, len(chunk_sizes)], device=device, dtype=torch.long, generator=generator, ) if clean_index is not None and per_chunk.shape[1] > 0 \ and max_timestep - min_timestep > 1: all_clean = (per_chunk == clean_index).all(dim=1) if bool(all_clean.any()): rows = all_clean.nonzero(as_tuple=True)[0] # One chunk per affected sample, redrawn uniformly over every index # *except* the clean one: draw from the shrunken range, then shift # past the hole. Works for any clean_index, not just the last one. pick = torch.randint( 0, per_chunk.shape[1], [rows.numel()], device=device, dtype=torch.long, generator=generator) redraw = torch.randint( min_timestep, max_timestep - 1, [rows.numel()], device=device, dtype=torch.long, generator=generator) per_chunk[rows, pick] = redraw + (redraw >= clean_index).long() timestep = torch.empty( batch_size, num_frames, device=device, dtype=per_chunk.dtype) for chunk_idx, (start_frame, end_frame) in enumerate(chunk_boundaries(chunk_sizes)): timestep[:, start_frame:end_frame] = per_chunk[:, chunk_idx:chunk_idx + 1] return timestep def sample_flex_partition(args, num_frames, torch_rng, device): """Draw this iteration's Flex-Forcing partition (§3.1 / §3.3). Returns ``None`` when Flex-Forcing is off, so every caller keeps using the uniform ``num_frame_per_block`` layout. The draw is :func:`videox_fun.utils.flex_chunking.sample_flex_level0` - the *same* four-arm mixture stage 2 uses for its level 0, random arm at ``FLEX_RANDOM_PROB`` (10%) and the three deterministic arms at 30% each. Sharing it is the whole point: this stage's checkpoint is stage 2's ``--ode_transformer_path`` initialisation, so every layout stage 2 spends iterations on has to be one stage 1 actually regressed. The two-band sampler this replaced drew 10% uniform / 90% random - coincidentally the mirror image of today's split, and beside the point, since what made it unusable was its chunk support being exactly ``{--flex_chunk_min .. --flex_chunk_max}``: the whole-clip ``[F]`` layout that 60% of stage 2's level-0 iterations start from (at ``--flex_pyramid_levels=4``) had probability *zero* there (``max_chunk`` below the frame count makes it unreachable, and at most 1/19 of random draws even with ``max_chunk`` raised to it), and so did every 1-frame chunk - ``[1] * F`` is ``1/F!`` under a random draw, i.e. unreachable at any setting of these two knobs. §3.2's pyramid is a *multi-step* schedule and this stage regresses x0 in a single forward per chunk, so there is no ladder to build here and only level 0 is used. ``--flex_pyramid_levels`` is still read, and must match stage 2's, because it decides whether the coarse arm exists and therefore how the other arms split the 90% the random one leaves behind. ``sample_flex_level0`` broadcasts rank 0's draw, so every rank regresses the same layout; otherwise the FlexAttention mask and the per-chunk timesteps would disagree across the group. """ if not args.flex_forcing: return None return sample_flex_level0( num_frames, num_frame_per_block=args.num_frame_per_block, chunk_min=args.flex_chunk_min, chunk_max=args.flex_chunk_max, pyramid=args.flex_pyramid_levels > 1, generator=torch_rng, device=device, independent_first_frame=args.independent_first_frame)["sizes"] def initialize_kv_cache_for_training(batch_size, num_frames, frame_seq_length, num_layers, num_heads, head_dim, dtype, device): """Initialize KV cache for block-by-block training (mirrors train_distill).""" kv_cache_size = num_frames * frame_seq_length kv_cache = [] for _ in range(num_layers): kv_cache.append({ "k": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), "v": torch.zeros([batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device), "global_end_index": torch.tensor([0], dtype=torch.long, device=device), "local_end_index": torch.tensor([0], dtype=torch.long, device=device), }) return kv_cache def initialize_crossattn_cache_for_training(batch_size, text_len, num_layers, num_heads, head_dim, dtype, device): """Initialize cross-attention cache for block-by-block training.""" crossattn_cache = [] for _ in range(num_layers): crossattn_cache.append({ "k": torch.zeros([batch_size, text_len, num_heads, head_dim], dtype=dtype, device=device), "v": torch.zeros([batch_size, text_len, num_heads, head_dim], dtype=dtype, device=device), "is_init": False, }) return crossattn_cache # ============================================================================ # Args # ============================================================================ def parse_args(): parser = argparse.ArgumentParser(description="ODE Regression Training for Self-Forcing") parser.add_argument( "--pretrained_model_name_or_path", type=str, required=True, help="Path to pretrained model or model identifier from huggingface.co/models.", ) parser.add_argument( "--config_path", type=str, required=True, help="The config of the model in training.", ) 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, required=True, help="Path to JSON annotation file listing ODE trajectory safetensors files.", ) parser.add_argument( "--output_dir", type=str, default="output_ode_regression", help="The output directory where the model predictions and checkpoints will be written.", ) parser.add_argument( "--logging_dir", type=str, default="logs", help="TensorBoard log directory. Will default to *output_dir/logs*.", ) parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") parser.add_argument( "--train_batch_size", type=int, default=1, help="Batch size (per device) for the training dataloader." ) 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=2e-6, 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_with_warmup", 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=100, 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( "--use_came", action="store_true", help="Whether or not to use CAME optimizer." ) 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("--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=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( "--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." ), ) 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"`.' ), ) parser.add_argument( "--dataloader_num_workers", type=int, default=8, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.", ) 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( "--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( "--validation_prompts", type=str, default=None, nargs="+", help="A set of prompts evaluated every `--validation_steps` and logged to `--report_to`.", ) parser.add_argument( "--validation_steps", type=int, default=500, help="Run validation every X steps.", ) parser.add_argument( "--validation_epochs", type=int, default=500, help="Run validation every X epochs.", ) parser.add_argument( "--video_sample_size", type=int, default=640, help="Sample size of the video.", ) parser.add_argument( "--video_sample_n_frames", type=int, default=81, help="Num frame of video.", ) 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( "--transformer_path", type=str, default=None, help="If you want to load the weight from other transformers, input its path.", ) parser.add_argument( "--train_sampling_steps", type=int, default=1000, help="Total number of scheduler timesteps for sampling.", ) parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank.") # Self-Forcing specific parser.add_argument( '--denoising_step_indices_list', nargs='+', type=int, default=[1000, 750, 500, 250], help="The denoising step list.", ) parser.add_argument( "--num_frame_per_block", type=int, default=3, help="Number of frames per block for Self-Forcing causal training.", ) parser.add_argument( "--flex_forcing", action="store_true", help="Enable Flex-Forcing (arXiv 2607.03509): replace the scalar " "--num_frame_per_block with a partition of the frame axis that is " "re-drawn every iteration. Instantiates " "WanTransformer3DModel_FlexForcing and validates with " "WanFlexForcingPipeline. Keep this on for the ODE stage too, or the " "model enters distillation having only ever seen uniform blocks." ) parser.add_argument( "--flex_chunk_min", type=int, default=2, help="Smallest chunk drawn by the *random* arm of the level-0 mixture " "(paper §5 uses 2). Unlike stage 2, nothing here refines a coarse " "chunk down to 1-frame leaves, so 1 is a real setting: it lets the " "random arm emit single-frame chunks itself. The fully causal " "[1] * F layout is a separate arm and is drawn either way." ) parser.add_argument( "--flex_chunk_max", type=int, default=10, help="Largest chunk drawn by the *random* arm of the level-0 mixture " "(paper §5 uses 10; setting it equal to --flex_chunk_min pins that " "arm to one fixed layout). The whole-clip layout is a separate arm, " "so this knob does not decide whether stage 1 ever regresses the " "clip as one fully bidirectional chunk. Raising it to the latent " "frame count only spends random-arm draws on a layout the coarse arm " "already covers - at 21 latent frames it makes [21] 1/19 of that arm " "and thins every other one, cutting P([3]*7) from 1/183708 to " "1/1106560." ) parser.add_argument( "--flex_pyramid_levels", type=int, default=1, help="Must match stage 2's --flex_pyramid_levels. Stage 1 builds no ladder " "(it regresses x0 in one forward per chunk, so only level 0 is used), " "but the level-0 mixture is shared with stage 2 and this flag decides " "whether its coarse arm exists: >1 gives 30%% whole-clip / 30%% " "uniform / 30%% full_then_blocks / 10%% random, while =1 divides the " "coarse arm's share between the two survivors - random stays at 10%% - " "for 45/45/10. Leaving it at 1 while stage 2 " "runs 4 would under-train exactly the [F] layout stage 2 starts most " "of its iterations from." ) parser.add_argument( "--independent_first_frame", action="store_true", help="Whether first frame is independent ([1, N, N, ...] pattern).", ) parser.add_argument( "--context_noise", type=int, default=0, help="Context noise level for KV cache update (matches training config).", ) parser.add_argument( '--trainable_modules', nargs='+', default=['.'], 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( "--save_state", action="store_true", help="Whether to save full accelerator state on checkpoint." ) parser.add_argument( "--shift", type=float, default=8.0, help="Shift value for FlowMatchEulerDiscreteScheduler. Default: 8.0 (matches ODE data generation).", ) parser.add_argument( "--use_kv_cache_training", action="store_true", help=( "If set, run block-by-block KV cache training that fully matches the " "pipeline_wan_self_forcing inference behavior. Otherwise fall back to " "the default one-shot causal-mask ODE regression (kept as baseline)." ), ) parser.add_argument( "--prob_full_zero_start", type=float, default=0.0, help=( "Probability (per-sample) of forcing ALL frames in ALL blocks to use " "timestep index=0 (pure-noise start). Bridges the train-inference gap " "so the model also sees the real autoregressive rollout where every " "block starts from fresh noise. 0.0 disables (default)." ), ) args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) if env_local_rank != -1 and env_local_rank != args.local_rank: args.local_rank = env_local_rank return args def main(): args = parse_args() if args.flex_forcing and (args.flex_chunk_min < 1 or args.flex_chunk_max < args.flex_chunk_min): raise ValueError( f"--flex_chunk_min/--flex_chunk_max must satisfy 1 <= min <= max, " f"got {args.flex_chunk_min}/{args.flex_chunk_max}.") 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, ) 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 # Load scheduler. scheduler_kwargs = OmegaConf.to_container(config['scheduler_kwargs']) scheduler_kwargs['shift'] = args.shift noise_scheduler = FlowMatchEulerDiscreteScheduler( **filter_kwargs(FlowMatchEulerDiscreteScheduler, scheduler_kwargs) ) # NOTE: tokenizer, text_encoder and vae are NOT loaded here. They are lazily loaded inside # `log_validation` only when validation actually runs, then freed afterwards. # Get causal Transformer (generator) # Flex-Forcing (arXiv 2607.03509) swaps in a generator that accepts a # per-iteration partition of the frame axis and carries the K-Projection of # 3.3; the ODE regression itself is unchanged. transformer_cls = (WanTransformer3DModel_FlexForcing if args.flex_forcing else WanTransformer3DModel_SelfForcing) # 3.3 is not a launcher knob: which projection the backbone gets comes from # the model config (`transformer_additional_kwargs.flex_kproj_mode`, whose # default is the paper's 'diag_rank1'). It is identity-initialised, so # carrying it from the ODE stage costs nothing and makes the distillation # checkpoint load exactly. transformer_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs']) transformer3d = transformer_cls.from_pretrained( os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')), transformer_additional_kwargs=transformer_kwargs, low_cpu_mem_usage=True, ).to(weight_dtype) # Set transformer3d to non-trainable initially (trainable_modules will toggle below) transformer3d.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 state_dict = load_file(args.transformer_path) else: state_dict = torch.load(args.transformer_path, map_location="cpu") state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict state_dict = state_dict["generator_ema"] if "generator_ema" in state_dict else state_dict state_dict = state_dict["generator"] if "generator" in state_dict else state_dict if any(k.startswith("model.") for k in state_dict.keys()): state_dict = {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in state_dict.items()} m, u = transformer3d.load_state_dict(state_dict, strict=False) print(f"Loaded transformer_path. missing keys: {len(m)}, unexpected keys: {len(u)}") assert len(u) == 0 # Set trainable modules # A good trainable modules is showed below now. # For full finetune: trainable_modules = ['.'] transformer3d.train() for name, param in transformer3d.named_parameters(): for trainable_module_name in args.trainable_modules + args.trainable_modules_low_learning_rate: if trainable_module_name in name: param.requires_grad = True break if args.flex_forcing and any("flex_kproj" in n for n, _ in transformer3d.named_parameters()): # Whatever projection the config asked for has to end up trainable: the # default --trainable_modules ['.'] already covers it, but a narrowed # list would silently freeze Pi at its identity initialisation. for name, param in transformer3d.named_parameters(): if "flex_kproj" in name: param.requires_grad = True if accelerator.is_main_process: accelerator.print( "Flex-Forcing K-Projection " f"('{getattr(transformer3d, 'flex_kproj_mode', 'none')}') " "trainable." ) # `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: 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 = models.pop() # load diffusers style into model # Reopen with the class that actually wrote this directory (the # save hook above calls `models[0].save_pretrained()`); the base # `WanTransformer3DModel` is not even imported here. load_model = transformer_cls.from_pretrained( input_dir, subfolder="transformer" ) 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) if args.gradient_checkpointing: 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, transformer3d.parameters())) trainable_params_optim = [ {'params': [], 'lr': args.learning_rate}, {'params': [], 'lr': args.learning_rate / 2}, ] in_already = [] for name, param in transformer3d.named_parameters(): if not param.requires_grad: continue if name in in_already: continue # 3.3's K-Projection is grouped like every other layer: the default # --trainable_modules ['.'] matches `flex_kproj.*`, so Pi trains at # --learning_rate together with the rest of the generator. high_lr = False for m in args.trainable_modules: if m in name: in_already.append(name) high_lr = True trainable_params_optim[0]['params'].append(param) break if high_lr: continue for m in args.trainable_modules_low_learning_rate: if m in name: in_already.append(name) trainable_params_optim[1]['params'].append(param) 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) ) 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, ) # Get the training dataset train_dataset = ImageVideoSafetensorsDataset( args.train_data_meta, data_root=args.train_data_dir ) # DataLoaders creation: from torch.utils.data import BatchSampler, Dataset, Sampler 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 ) def ode_safetensors_collate_fn(examples): """Collate safetensors-loaded ODE samples into a batch. Each sample is a dict with keys: - 'latents': [S, C, F, H, W] - 'prompt_embeds': [L, D] - 'prompt_attention_mask': [L] The default torch collate fails when, across samples, the same key has slightly different dtypes/lengths (e.g. attention_mask saved as bool/int vs long, or prompt_embeds with different seq lengths). This custom collate normalizes dtypes and pads variable-length text fields so `torch.stack` always succeeds. """ out = {} # ---- latents: assume identical shape across samples (fixed by pipeline) ---- latents = [ex["latents"] for ex in examples] target_latent_dtype = latents[0].dtype latents = [t.to(target_latent_dtype) for t in latents] out["latents"] = torch.stack(latents, dim=0) # ---- prompt_embeds: pad along seq dim, unify dtype ---- embeds = [ex["prompt_embeds"] for ex in examples] embed_dtype = embeds[0].dtype max_len = max(e.shape[0] for e in embeds) padded_embeds = [] for e in embeds: e = e.to(embed_dtype) if e.shape[0] < max_len: pad = torch.zeros( max_len - e.shape[0], *e.shape[1:], dtype=embed_dtype ) e = torch.cat([e, pad], dim=0) padded_embeds.append(e) out["prompt_embeds"] = torch.stack(padded_embeds, dim=0) # ---- prompt_attention_mask: pad along seq dim, force long dtype ---- masks = [ex["prompt_attention_mask"].long() for ex in examples] max_len = max(m.shape[0] for m in masks) padded_masks = [] for m in masks: if m.shape[0] < max_len: pad = torch.zeros(max_len - m.shape[0], dtype=torch.long) m = torch.cat([m, pad], dim=0) padded_masks.append(m) out["prompt_attention_mask"] = torch.stack(padded_masks, dim=0) return out 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, collate_fn=ode_safetensors_collate_fn, ) # 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, ) # Prepare everything with our `accelerator`. transformer3d, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( transformer3d, optimizer, train_dataloader, lr_scheduler ) # Compute denoising_step_list from the scheduler's own grid. # Deliberately *not* calling `set_timesteps(...)`: it rebuilds the grid with # a different `sigma_min`, so indexing it lands away from the t the pipeline # actually denoises at. The `__init__` grid maps index k to exactly # t = 1 - k / num_train_timesteps, which is what # `stochastic_sampling_timesteps` hardcodes, so training, validation and # inference read the same numbers. `train_distill.py` relies on the same # grid. The explicit `.to()` keeps these on the accelerator, which the # per-frame `timestep` below requires. denoising_step_list = noise_scheduler.timesteps.to(accelerator.device)[ args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list) ] # Training denoising step list: append 0 (clean) for train-inference context alignment. # index=4 frames use clean latent as input but are excluded from loss via mask=(timestep!=0). # They serve as clean context for later blocks via causal attention. train_denoising_step_list = denoising_step_list if 0 not in denoising_step_list.tolist(): train_denoising_step_list = torch.cat([ denoising_step_list, torch.tensor([0], device=denoising_step_list.device) ]) num_denoising_steps = len(train_denoising_step_list) # Index of the clean (t=0) entry, derived from the list itself rather than # assumed to be the last one. Passed to `get_timestep_for_ode_flexible`, # which keeps every iteration from drawing it for *all* chunks - an # all-clean draw zeroes the loss mask and makes the fallback MSE exactly 0, # i.e. a zero-gradient iteration. See that function for the probability. _clean_hits = (train_denoising_step_list == 0).nonzero(as_tuple=True)[0] clean_step_index = int(_clean_hits[-1]) if _clean_hits.numel() > 0 else None if accelerator.is_main_process: print(f"Denoising step list (inference): {denoising_step_list.tolist()}") print(f"Denoising step list (training): {train_denoising_step_list.tolist()}") print(f"num_denoising_steps (includes clean): {num_denoising_steps}") print(f"Dataset size: {len(train_dataset)}") # 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}") accelerator.load_state(os.path.join(args.output_dir, 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, ) # ==================================================================== # Training loop # ==================================================================== for epoch in range(first_epoch, args.num_train_epochs): train_loss = 0.0 batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) for step, batch in enumerate(train_dataloader): with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): # --- Read preprocessed prompt embeddings --- prompt_embeds_raw = batch["prompt_embeds"].to(device=accelerator.device, dtype=weight_dtype) encoder_attention_mask = batch["prompt_attention_mask"].to(device=accelerator.device) bsz = prompt_embeds_raw.shape[0] # Truncate to actual token length (same as preprocess training) seq_lens = encoder_attention_mask.gt(0).sum(dim=1).long() prompt_embeds = [prompt_embeds_raw[i, :seq_lens[i].item()] for i in range(bsz)] # --- Prepare ODE regression inputs --- # ode_latent: [B, num_denoising_steps, C, F, H, W] ode_latent = batch["latents"].to(device=accelerator.device, dtype=weight_dtype) # Target: clean endpoint (last timestep) target_latent = ode_latent[:, -1] # [B, C, F, H, W] num_frames = target_latent.shape[2] C_dim, F_dim, H_dim, W_dim = ( ode_latent.shape[2], ode_latent.shape[3], ode_latent.shape[4], ode_latent.shape[5], ) patch_h, patch_w = accelerator.unwrap_model(transformer3d).config.patch_size[1:] frame_seqlen = (H_dim * W_dim) // (patch_h * patch_w) seq_len = frame_seqlen * num_frames # Flex-Forcing (arXiv 2607.03509 3.1/3.3): one random partition of the # frame axis per iteration, drawn outside the accumulate context so # both training branches below cut the clip the same way. # None = Flex-Forcing off. flex_partition = sample_flex_partition( args, num_frames, torch_rng, accelerator.device) with accelerator.accumulate(transformer3d): with torch.cuda.amp.autocast(dtype=weight_dtype): if args.use_kv_cache_training: # ============================================================ # Block-by-block KV cache training (autoregressive, single-step x0) # Starting timestep is randomly sampled per block — same as the # non-KV-cache (baseline) branch. Each block performs ONE forward # to predict x0; KV cache is then refreshed with pred_block + # context_noise to keep the autoregressive rollout intact. # ============================================================ # 1) Block split (mirrors pipeline_wan_self_forcing) if flex_partition is not None: # A Flex-Forcing partition covers `num_frames` exactly by # construction, so the divisibility asserts of the # uniform layout do not apply. all_num_frames = flex_partition num_blocks_split = len(all_num_frames) elif not args.independent_first_frame: assert num_frames % args.num_frame_per_block == 0 num_blocks_split = num_frames // args.num_frame_per_block all_num_frames = [args.num_frame_per_block] * num_blocks_split else: assert (num_frames - 1) % args.num_frame_per_block == 0 num_blocks_split = (num_frames - 1) // args.num_frame_per_block all_num_frames = [1] + [args.num_frame_per_block] * num_blocks_split # 2) Random timestep index per frame/block (same as baseline branch) if flex_partition is not None: index = get_timestep_for_ode_flexible( 0, num_denoising_steps, bsz, all_num_frames, accelerator.device, generator=torch_rng, clean_index=clean_step_index, ) # [B, F] else: index = get_timestep_for_ode( 0, num_denoising_steps, bsz, num_frames, args.num_frame_per_block, args.independent_first_frame, accelerator.device, generator=torch_rng, ) # [B, F] # Optional: force per-sample full-zero start to cover the # real inference rollout (all blocks starting from pure noise). if args.prob_full_zero_start > 0.0: zero_mask = ( torch.rand(bsz, device=accelerator.device, generator=torch_rng) < args.prob_full_zero_start ) if zero_mask.any(): index[zero_mask] = 0 gather_index = index.reshape(bsz, 1, 1, num_frames, 1, 1).expand( -1, -1, C_dim, -1, H_dim, W_dim ) noisy_input_full = torch.gather(ode_latent, dim=1, index=gather_index).squeeze(1) timestep_full = train_denoising_step_list[index] # [B, F] # 3) Initialize KV / cross-attention cache cfg = accelerator.unwrap_model(transformer3d).config num_layers_t = cfg.num_layers num_heads_t = cfg.num_heads head_dim_t = cfg.dim // num_heads_t text_len = 512 # T5 sequence length kv_cache = initialize_kv_cache_for_training( batch_size=bsz, num_frames=num_frames, frame_seq_length=frame_seqlen, num_layers=num_layers_t, num_heads=num_heads_t, head_dim=head_dim_t, dtype=weight_dtype, device=accelerator.device, ) crossattn_cache = initialize_crossattn_cache_for_training( batch_size=bsz, text_len=text_len, num_layers=num_layers_t, num_heads=num_heads_t, head_dim=head_dim_t, dtype=weight_dtype, device=accelerator.device, ) # 4) Sigma / timestep lookup tables (per-frame sigma) sigmas_full = noise_scheduler.sigmas.to( device=accelerator.device, dtype=torch.float64 ) schedule_timesteps_full = noise_scheduler.timesteps.to(accelerator.device) current_start_frame = 0 total_pred = torch.zeros_like(target_latent) full_seq_len = frame_seqlen * num_frames # 5) Block-by-block rollout — single-step x0 prediction per block for block_idx, current_num_frames in enumerate(all_num_frames): start_idx = current_start_frame end_idx = current_start_frame + current_num_frames noisy_input = noisy_input_full[:, :, start_idx:end_idx] timestep_block = timestep_full[:, start_idx:end_idx].to(torch.int64) flow_pred = transformer3d( x=[noisy_input[i] for i in range(bsz)], context=prompt_embeds, t=timestep_block, seq_len=full_seq_len, kv_cache=kv_cache, crossattn_cache=crossattn_cache, current_start=current_start_frame * frame_seqlen, cache_start=None, ) if isinstance(flow_pred, list): flow_pred = torch.stack(flow_pred, dim=0) # Per-frame sigma -> x0 = xt - sigma * flow_pred step_indices_block = torch.argmin( (schedule_timesteps_full.unsqueeze(0) - timestep_block.reshape(-1).unsqueeze(1)).abs(), dim=1 ) sigma_block = sigmas_full[step_indices_block].to(weight_dtype) # timestep=0 (clean context) must use sigma=0 exactly. sigma_block[timestep_block.reshape(-1) == 0] = 0.0 sigma_block = sigma_block.reshape(bsz, 1, current_num_frames, 1, 1) pred_block = noisy_input - sigma_block * flow_pred total_pred[:, :, start_idx:end_idx] = pred_block # 6) Update KV cache with student's pred_block + context_noise # (matches pipeline_wan_self_forcing L802-L839) if block_idx < len(all_num_frames) - 1: ctx_t = torch.full( [bsz, current_num_frames], args.context_noise, device=accelerator.device, dtype=torch.int64, ) with torch.no_grad(): transformer3d( x=[pred_block[i] for i in range(bsz)], context=prompt_embeds, t=ctx_t, seq_len=full_seq_len, kv_cache=kv_cache, crossattn_cache=crossattn_cache, current_start=current_start_frame * frame_seqlen, cache_start=None, ) current_start_frame += current_num_frames # 7) ODE-endpoint MSE loss (mask out clean timestep=0 frames) mask = (timestep_full != 0).unsqueeze(1).unsqueeze(-1).unsqueeze(-1) mask = mask.expand_as(target_latent).float() if mask.sum() > 0: loss = F.mse_loss(total_pred * mask, target_latent * mask, reduction="sum") / mask.sum() else: loss = F.mse_loss(total_pred, target_latent) else: # --- Baseline (one-shot causal-mask) preparation --- # Random timestep index per frame/block if flex_partition is not None: index = get_timestep_for_ode_flexible( 0, num_denoising_steps, bsz, flex_partition, accelerator.device, generator=torch_rng, clean_index=clean_step_index, ) # [B, F] else: index = get_timestep_for_ode( 0, num_denoising_steps, bsz, num_frames, args.num_frame_per_block, args.independent_first_frame, accelerator.device, generator=torch_rng, ) # [B, F] # Optional: force per-sample full-zero start to cover the # real inference rollout (all blocks starting from pure noise). if args.prob_full_zero_start > 0.0: zero_mask = ( torch.rand(bsz, device=accelerator.device, generator=torch_rng) < args.prob_full_zero_start ) if zero_mask.any(): index[zero_mask] = 0 # Gather noisy input from ODE trajectory gather_index = index.reshape(bsz, 1, 1, num_frames, 1, 1).expand( -1, -1, C_dim, -1, H_dim, W_dim ) noisy_input = torch.gather(ode_latent, dim=1, index=gather_index).squeeze(1) # Compute actual timestep values: [B, F] timestep = train_denoising_step_list[index] # [B, F] # Build causal block mask. With a Flex-Forcing partition # installed the model dispatches to its flexible builder # (3.1); otherwise this is the inherited uniform mask. if flex_partition is not None: accelerator.unwrap_model(transformer3d).set_flex_chunk_sizes(flex_partition) accelerator.unwrap_model(transformer3d).create_block_mask_for_training( num_frames=num_frames, frame_seqlen=frame_seqlen, num_frame_per_block=args.num_frame_per_block, independent_first_frame=args.independent_first_frame, device=accelerator.device ) # Convert to list format for transformer noisy_input_list = [noisy_input[i] for i in range(bsz)] # ============================================================ # Baseline: one-shot causal-mask ODE regression # ============================================================ flow_pred = transformer3d( x=noisy_input_list, context=prompt_embeds, t=timestep, seq_len=seq_len, ) # Convert flow prediction to x0 prediction (per-frame). sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=torch.float64) schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device) step_indices = torch.argmin( (schedule_timesteps.unsqueeze(0) - timestep.reshape(-1).unsqueeze(1)).abs(), dim=1 ) sigma = sigmas[step_indices].to(weight_dtype) # Fix: timestep=0 (clean context frames) should have sigma=0 exactly. sigma[timestep.reshape(-1) == 0] = 0.0 sigma = sigma.reshape(bsz, 1, num_frames, 1, 1) pred_x0 = noisy_input - sigma * flow_pred # MSE loss (mask t=0 frames) mask = (timestep != 0).unsqueeze(1).unsqueeze(-1).unsqueeze(-1) mask = mask.expand_as(target_latent).float() if mask.sum() > 0: loss = F.mse_loss(pred_x0 * mask, target_latent * mask, reduction="sum") / mask.sum() else: loss = F.mse_loss(pred_x0, target_latent) avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean() train_loss += avg_loss.item() / args.gradient_accumulation_steps accelerator.backward(loss) if accelerator.sync_gradients: accelerator.clip_grad_norm_(trainable_params, args.max_grad_norm) optimizer.step() lr_scheduler.step() optimizer.zero_grad() # Logging and checkpointing if accelerator.sync_gradients: progress_bar.update(1) global_step += 1 accelerator.log({"train_loss": train_loss}, step=global_step) train_loss = 0.0 if global_step % args.checkpointing_steps == 0: if accelerator.is_main_process or fsdp_stage == 3 or zero_stage == 3: # _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[: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 and leaves the freed blocks in # the allocator caches, so the cache flushes run inside the same window. with progress_bar.paused(): accelerator.save_state(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( transformer3d, args, config, accelerator, weight_dtype, global_step, fsdp_stage=fsdp_stage, zero_stage=zero_stage, ) logs = {"loss": 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( transformer3d, args, config, accelerator, weight_dtype, global_step, fsdp_stage=fsdp_stage, zero_stage=zero_stage, ) # 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 or fsdp_stage == 3 or zero_stage == 3: gc.collect() torch.cuda.empty_cache() torch.cuda.ipc_collect() save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") accelerator.save_state(save_path) logger.info(f"Saved state to {save_path}") accelerator.end_training() if __name__ == "__main__": main()