nodes/aigc-apps-VideoX-Fun
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
sparkwan
aigc-apps-VideoX-Fun/scripts/sparkdiffusion/train_distill_wan2.2.py
T
bubbliiiing e94e7e0c78 Update sparkwan
2026-10-20 13:23:31 +08:00

3138 lines
163 KiB
Python
Raw Permalink Blame History

"""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.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,
SparkWan2_2Transformer3DModel,
SparkWanTransformer3DModel,
Wan2_2Transformer3DModel, WanT5EncoderModel,
WanTransformer3DModel)
from videox_fun.pipeline import (Wan2_2I2VPipeline, Wan2_2Pipeline,
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, transformer3d, args, config, accelerator, weight_dtype, global_step, init_rollout_transformer3d=None):
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):
# FSDP(CPUOffload)-wrapped modules own their placement: never .to() the sharded
# text_encoder here or on exit; _te_sharded marks it and the generator below.
_te_sharded = (
getattr(accelerator.state, "fsdp_plugin", None) is not None
or getattr(accelerator.state, "deepspeed_plugin", None) is not None
)
logger.info("Running validation... ")
scheduler = FlowMatchEulerDiscreteScheduler(
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
# Full joint few-step trajectory (both experts, boundary-routed): [1.0] + the joint grid knots.
val_inference_steps = len(args.joint_backward_timesteps) + 1
# Dual-expert (MoE) validation: the pipeline routes rf s >= boundary to transformer_2 (the
# HIGH-noise expert) and s < boundary to transformer (the LOW-noise expert). The expert being
# trained is the live (possibly FSDP-wrapped) generator; the complementary expert is loaded
# fresh from its noise-model subpath so the full trajectory can be sampled.
if args.boundary_type == "full":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
transformer3d_2 = None
else:
if args.boundary_type == "low":
transformer3d_1 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
# Reuse the FROZEN distilled high-noise student that produced the low expert's
# training rollout start latent, so the high-noise segment (rf >= boundary) of the
# validation trajectory matches training/deployment instead of the base many-step expert.
if init_rollout_transformer3d is not None:
transformer3d_2 = init_rollout_transformer3d
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
transformer3d_2 = (SparkWan2_2Transformer3DModel if args.sparse_student else Wan2_2Transformer3DModel).from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
if args.sparse_student:
transformer3d_2.set_student_topk(args.student_topk)
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
transformer3d_1 = (SparkWan2_2Transformer3DModel if args.sparse_student else Wan2_2Transformer3DModel).from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
if args.sparse_student:
transformer3d_1.set_student_topk(args.student_topk)
transformer3d_2 = accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d
if args.train_mode != "normal":
pipeline = Wan2_2I2VPipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
else:
pipeline = Wan2_2Pipeline(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer3d_1,
transformer_2=transformer3d_2,
scheduler=scheduler,
)
# Stage only non-FSDP components on GPU; the FSDP-wrapped generator/text_encoder
# self-stream in place. VAE gets an offload hook and pins the pipeline device.
from videox_fun.dist import offload_components_cpu
_fsdp_cls_name = 'FullyShardedDataParallel'
for _comp in (pipeline.transformer, pipeline.transformer_2):
if _comp is not None and type(_comp).__name__ != _fsdp_cls_name:
_comp.to(accelerator.device, dtype=weight_dtype)
if not _te_sharded:
text_encoder.to(accelerator.device, dtype=weight_dtype)
# VAE is never FSDP-wrapped; give it an onload/offload hook and force pipeline device -> GPU.
offload_components_cpu(pipeline, accelerator.device, ["vae"])
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=[height, width])
sample = pipeline(
args.validation_prompts[i],
num_frames = video_length,
negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
height = height,
width = width,
generator = generator,
video = input_video,
mask_video = input_video_mask,
num_inference_steps = val_inference_steps,
guidance_scale = 1.0,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.875),
).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:
# Match the fixed training geometry when --fix_sample_size is set; otherwise the
# square video_sample_size default diverges from the trained aspect ratio.
if args.fix_sample_size is not None:
val_height, val_width = args.fix_sample_size
else:
val_height = val_width = args.video_sample_size
sample = pipeline(
args.validation_prompts[i],
num_frames = args.video_sample_n_frames,
negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
height = val_height,
width = val_width,
generator = generator,
num_inference_steps = val_inference_steps,
guidance_scale = 1.0,
boundary = config['transformer_additional_kwargs'].get('boundary', 0.875),
).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"
)
)
# Detach the validation-only VAE offload hook before returning to training so the
# training loop keeps sole control of the VAE's device placement.
for _h in getattr(pipeline, "_cpu_offload_component_hooks", []):
_h.remove()
# Drop pipeline AND the local handles to the freshly-loaded complementary expert (a full
# 14B on GPU, not FSDP-managed). `del pipeline` alone leaves transformer3d_1/_2 pinning it,
# so empty_cache below reclaims nothing and the 28GB lingers/fragments into the next
# student backward. The live generator survives via the accelerator's own reference.
del pipeline, transformer3d_1, transformer3d_2
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 and not _te_sharded:
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:
if 'pipeline' in locals():
for _h in getattr(pipeline, "_cpu_offload_component_hooks", []):
_h.remove()
del pipeline
# Release the freshly-loaded expert before empty_cache so it is actually reclaimed here.
if 'transformer3d_1' in locals():
del transformer3d_1
if 'transformer3d_2' in locals():
del transformer3d_2
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)
_te_sharded_exc = (
getattr(accelerator.state, "fsdp_plugin", None) is not None
or getattr(accelerator.state, "deepspeed_plugin", None) is not None
)
if not args.enable_text_encoder_in_dataloader and not _te_sharded_exc:
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="config/wan2.2/wan_civitai_t2v.yaml",
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(
"--generator_transformer_path",
type=str,
default=None,
help=("Path to a (sparse stage-1) student checkpoint loaded into the generator. When "
"--sparse_student is off it is also loaded into the dense critic (fake_score); "
"with --sparse_student the critic stays dense (initialized from the teacher)."),
)
parser.add_argument(
"--sparse_student",
action="store_true",
help="Distill a sparse (RoLa/SLA-style) student against the dense teacher: generator and init_rollout use SparkWan2_2Transformer3DModel (sparse + linear self-attention), teacher and critic stay dense.",
)
parser.add_argument(
"--student_topk",
type=float,
default=0.05,
help="Ratio of key blocks kept per query in the sparse student self-attention (sparsity = 1 - topk). Must match the --student_topk used in train_finetune_sparse.py.",
)
parser.add_argument(
"--boundary_type",
type=str,
default="low",
help=(
'The training boundary type for the dual-Transformer (MoE) architecture. '
'Support `"low"` (train the low-noise expert), `"high"` (train the high-noise '
'expert), and `"full"` (single-model training). Selects which noise-model subpath '
'is loaded and how the joint rf grid is split at the config boundary.'
),
)
parser.add_argument(
"--joint_backward_timesteps",
nargs='+',
type=float,
default=[0.933781, 0.875, 0.608979],
help="Canonical joint few-step rf grid (rCM, descending); its MIDDLE knot is the MoE "
"boundary. Value = rf time tan(theta)/(1+tan(theta)) at theta=[1.5, boundary, 1.0]. "
"Each expert's denoising grid is auto-split from it at the config boundary "
"(high: [1.0] + knots >= boundary, low: knots < boundary), and the rf values are "
"used directly as Euler sigmas.",
)
parser.add_argument(
"--init_rollout_timesteps",
nargs='+',
type=float,
default=[],
help="rf interior knots of the HIGH segment for the low expert's init rollout path "
"[1.0, *init_rollout_timesteps, boundary]. Empty = auto: joint-grid knots above "
"the boundary.",
)
parser.add_argument(
"--init_rollout_path",
type=str,
default="",
help="Path to the trained high-noise student checkpoint used to roll the frozen high "
"student from pure noise down to the boundary, producing the low-noise expert's "
"rollout starting latent online (empty = the low expert cold-starts from pure noise).",
)
parser.add_argument(
"--fake_score_transformer_path",
type=str,
default=None,
help=("If you want to warm start only the fake score from a weight, 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.",
)
parser.add_argument(
"--index_jitter_ratio",
type=float,
default=0.3,
help="Symmetric jitter budget (fraction of the neighboring gap) applied to the "
"denoising step indices when --randomize_step_indices is enabled.",
)
parser.add_argument(
"--flow_euler_rollout",
action="store_true",
help="Simulate the normal flow-matching inference rollout in the generator's multi-step "
"self-rollout (LightX2V-style): keep the model prediction in flow/velocity space and "
"advance to the next noise level with a deterministic Euler ODE step "
"(x_next = x_t + (sigma_next - sigma_t) * v), instead of converting the prediction "
"to x0 and re-noising with fresh noise. The final step still converts to x0 since "
"the DMD objective is defined on x0.",
)
parser.add_argument(
"--dfd",
action="store_true",
help="whether to use DFD post-training on a DMD-pretrained generator.",
)
parser.add_argument(
"--dfd_teacher_replace_prob",
type=float,
default=0.5,
help="Probability of replacing the teacher-score input with paired real data.",
)
parser.add_argument(
"--dfd_start_step",
type=int,
default=0,
help="Switch on DFD from this global_step onward; earlier steps run plain DMD.",
)
parser.add_argument(
"--decoupled_dmd",
action="store_true",
help="Use Decoupled DMD (arXiv:2511.22677): decompose the DMD gradient into a CFG Augmentation "
"term re-noised at tau_CA cleaner than the generator's current step and a Distribution "
"Matching term re-noised over the full noise range (Decoupled-Hybrid schedule).",
)
parser.add_argument(
"--tdm",
action="store_true",
help="Use TDM (Trajectory Distribution Matching, the official SparkDiffusion student loss) instead "
"of plain DMD. Three deltas from DMD: (1) the last rollout segment does NOT collapse to x0 but "
"takes one more rf-Euler step to s_{i-1} (the next grid knot), so the anchor is G_x_{t_{i-1}}; "
"(2) D_time is clamped strictly above s_{i-1}; (3) the fake/teacher x0 preds are mapped back to "
"s_{i-1} via the rf PF-ODE analytic step and the surrogate is compared in the x_{t_{i-1}} domain. "
"When s_{i-1}==0 (grid exhausted) it reduces exactly to DMD. Requires --flow_euler_rollout; "
"mutually exclusive with --dfd / --decoupled_dmd.",
)
parser.add_argument(
"--div_mode",
type=str,
default="pcm_random",
help="CrossDistill diversity mode: pcm_fixed / pcm_random (PCM-style phase-0 boundary "
"alignment). Only used when --loss_scale_div > 0.",
)
parser.add_argument(
"--loss_scale_div",
type=float,
default=0.0,
help="Weight of the CrossDistill PCM diversity term (per-element-mean space). 0 disables "
"diversity and recovers the plain-DMD baseline.",
)
parser.add_argument(
"--dcm_total_steps",
type=int,
default=48,
help="Subdivision of the continuous rf grid used by the PCM diversity term (dw = 1 / N).",
)
parser.add_argument(
"--dcm_skipping_interval_steps",
type=int,
default=1,
help="Number of teacher Euler sub-steps K within phase 0 for the PCM diversity term.",
)
parser.add_argument(
"--teacher_timestep_shift",
type=float,
default=12.0,
help="rf timestep shift for the PCM diversity grid (Wan2.2 T2V A14B teacher shift = 12.0).",
)
parser.add_argument(
"--div_rf_boundary",
type=float,
default=0.933781,
help="Phase-0 right boundary rf_b for the PCM diversity term (official backward_timesteps[0]).",
)
parser.add_argument(
"--div_teacher_guidance",
type=float,
default=5.0,
help="CFG scale for the PCM diversity teacher rollout (official teacher_guidance); "
"independent of --real_guidance_scale.",
)
args = parser.parse_args()
if isinstance(args.report_to, str) and args.report_to.lower() == "none":
args.report_to = None
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.dfd:
if args.train_mode != "normal":
raise ValueError("DFD currently supports T2V training with --train_mode normal.")
if args.enable_text_encoder_in_dataloader:
raise ValueError("DFD does not support --enable_text_encoder_in_dataloader.")
if args.seed is None:
raise ValueError("DFD requires an explicit --seed for reproducible post-training.")
if args.dfd_start_step < 0:
raise ValueError("--dfd_start_step must be non-negative.")
if not 0 <= args.dfd_teacher_replace_prob <= 1:
raise ValueError("--dfd_teacher_replace_prob must be in [0, 1].")
if args.gen_update_interval <= 0:
raise ValueError("--gen_update_interval must be greater than zero.")
if args.trainable_modules is None:
args.trainable_modules = [""]
if args.decoupled_dmd:
if args.real_guidance_scale <= 1.0:
raise ValueError("--real_guidance_scale must be greater than 1.0 to keep the CA term of Decoupled DMD active.")
if args.dfd:
raise ValueError("Decoupled DMD does not support --dfd.")
if any(int(i) <= 1 for i in args.denoising_step_indices_list):
print("WARNING: --denoising_step_indices_list contains a step index below 2, where tau_CA degenerates to tau_CA == t.")
if args.tdm:
if not args.flow_euler_rollout:
raise ValueError("--tdm requires --flow_euler_rollout (the trajectory must be a deterministic rf-Euler rollout).")
if args.dfd:
raise ValueError("--tdm does not support --dfd.")
if args.decoupled_dmd:
raise ValueError("--tdm does not support --decoupled_dmd.")
if args.loss_scale_div > 0 and args.div_mode not in ("pcm_fixed", "pcm_random"):
raise NotImplementedError(
f"--div_mode '{args.div_mode}' is not implemented for CrossDistill diversity; only "
"'pcm_fixed' and 'pcm_random' are supported (set --loss_scale_div 0 to disable it)."
)
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)
# --- Derive this expert's discrete rf grid from --boundary_type and the model's MoE `boundary`
# (the same value the pipeline routes on: the high expert handles rf s >= boundary). The canonical
# joint grid --joint_backward_timesteps (rCM, descending; its MIDDLE knot is the boundary) is split
# at the boundary into this expert's segment, and the rf values are used DIRECTLY as Euler sigmas
# (the shift=12 scheduler grid does not contain the rCM knots, so the timestep-index map is bypassed). ---
_boundary = float(config['transformer_additional_kwargs'].get('boundary', 0.875))
_joint_rf = [float(k) for k in args.joint_backward_timesteps]
if not args.init_rollout_timesteps:
args.init_rollout_timesteps = [k for k in _joint_rf if k > _boundary]
# Frozen-high-student path that rolls pure noise down to rf=boundary, depositing the low expert's start latent.
init_rollout_grid = [1.0] + [float(k) for k in args.init_rollout_timesteps] + [_boundary]
if args.boundary_type == "high":
denoising_rf_grid = [1.0] + [k for k in _joint_rf if k >= _boundary]
elif args.boundary_type == "low":
denoising_rf_grid = [_boundary] + [k for k in _joint_rf if k < _boundary]
else: # full: a single model walks the whole joint grid from pure noise
denoising_rf_grid = [1.0] + [k for k in _joint_rf if 0.0 < k < 1.0]
if len(denoising_rf_grid) < 2:
raise ValueError(
f"boundary_type='{args.boundary_type}' produced a degenerate rf grid {denoising_rf_grid}; "
f"check --joint_backward_timesteps against the config boundary {_boundary}."
)
print(f"[boundary_type={args.boundary_type}] MoE boundary={_boundary} | denoising rf grid={denoising_rf_grid} | init_rollout grid={init_rollout_grid}")
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 Transformer
# Dual-expert (MoE) Wan2.2: --boundary_type selects which noise model is the student/teacher.
# The high expert handles rf s >= boundary, the low expert rf s < boundary (the same routing the
# inference pipeline uses). real_score (teacher) and fake_score (critic) stay DENSE Wan2_2 and are
# loaded from this expert's own noise-model subpath, so each expert is distilled against its own teacher.
if args.boundary_type == "low" or args.boundary_type == "full":
sub_path = config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')
else:
sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
# Sparse student (SparkDiffusion stage-1 output): block-sparse + linear self-attention
# (SparkWan2_2Transformer3DModel); the dense teacher and the dense critic stay untouched,
# matching the reference pipeline where fake_score initializes from the teacher weights.
def build_student_net(weights_dir):
if args.sparse_student:
net = SparkWan2_2Transformer3DModel.from_pretrained(
weights_dir,
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
# Materialize the zero-init sparse branch that the dense base cannot provide.
net.set_student_topk(args.student_topk)
return net
return Wan2_2Transformer3DModel.from_pretrained(
weights_dir,
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
generator_transformer3d = build_student_net(os.path.join(args.pretrained_model_name_or_path, sub_path))
real_score_transformer3d = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
fake_score_transformer3d = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(args.pretrained_model_name_or_path, sub_path),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
).to(weight_dtype)
# Low-noise expert only: a FROZEN copy of the trained high-noise student, rolled from pure noise
# down to the boundary to build the low expert's on-distribution rollout start latent (init_rollout).
init_rollout_transformer3d = None
if args.init_rollout_path and args.init_rollout_path.strip():
high_sub_path = config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')
init_rollout_transformer3d = build_student_net(os.path.join(args.pretrained_model_name_or_path, high_sub_path))
_rollout_ckpt = args.init_rollout_path
if os.path.isdir(_rollout_ckpt):
_rollout_ckpt = os.path.join(_rollout_ckpt, "transformer", "diffusion_pytorch_model.safetensors") \
if os.path.exists(os.path.join(_rollout_ckpt, "transformer", "diffusion_pytorch_model.safetensors")) \
else os.path.join(_rollout_ckpt, "diffusion_pytorch_model.safetensors")
if _rollout_ckpt.endswith("safetensors"):
from safetensors.torch import load_file
_rollout_sd = load_file(_rollout_ckpt)
else:
_rollout_sd = torch.load(_rollout_ckpt, map_location="cpu")
_rollout_sd = _rollout_sd["state_dict"] if "state_dict" in _rollout_sd else _rollout_sd
_rollout_sd = { (k.split(".", 1)[1] if k.startswith(("generator_transformer3d.", "transformer.", "model.")) else k): v for k, v in _rollout_sd.items() }
m, u = init_rollout_transformer3d.load_state_dict(_rollout_sd, strict=False)
print(f"init_rollout missing keys: {len(m)}, unexpected keys: {len(u)}")
init_rollout_transformer3d.requires_grad_(False)
init_rollout_transformer3d.eval()
# 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)
real_score_transformer3d.eval()
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.generator_transformer_path is not None:
print(f"From generator/fake-score checkpoint: {args.generator_transformer_path}")
if args.generator_transformer_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(args.generator_transformer_path)
else:
state_dict = torch.load(args.generator_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)
print(f"generator missing keys: {len(m)}, unexpected keys: {len(u)}")
assert len(u) == 0
if args.sparse_student:
# Sparse attn_op.* keys have no dense-critic counterpart; keep fake_score dense
# (teacher-init), as in the reference pipeline where the critic skips student weights.
print("sparse_student: keeping fake_score dense (skipping sparse checkpoint load)")
else:
m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False)
print(f"fake_score missing keys: {len(m)}, unexpected keys: {len(u)}")
assert len(u) == 0
if args.fake_score_transformer_path is not None:
print(f"From fake-score checkpoint: {args.fake_score_transformer_path}")
if args.fake_score_transformer_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(args.fake_score_transformer_path)
else:
state_dict = torch.load(args.fake_score_transformer_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = fake_score_transformer3d.load_state_dict(state_dict, strict=False)
print(f"fake_score 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:
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
load_model = Wan2_2Transformer3DModel.from_pretrained(
input_dir, subfolder="transformer",
low_cpu_mem_usage=True,
)
model.register_to_config(**load_model.config)
model.load_state_dict(load_model.state_dict())
del load_model
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
if os.path.exists(pkl_path):
with open(pkl_path, 'rb') as file:
loaded_number, _ = pickle.load(file)
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
accelerator.register_save_state_pre_hook(save_model_hook)
accelerator.register_load_state_pre_hook(load_model_hook)
accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
if args.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
# DFD needs paired videos for the teacher-score input; everything else reuses the DMD data path.
# CrossDistill's PCM diversity term anchors on the paired real-video latent too, so it also needs the
# real-video path; otherwise the text-only collate drops pixel_values, the div anchor never fires and
# div_loss stays 0 (silently degrading CrossDistill back to plain DMD).
need_real_video = args.train_mode != "normal" or args.dfd or \
(args.loss_scale_div > 0 and args.div_mode in ("pcm_fixed", "pcm_random"))
if need_real_video:
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,
text_drop_ratio=0.0,
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 need_real_video:
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 not need_real_video:
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
# Under --low_vram use FSDP CPUOffload for CPU<->GPU streaming instead of manual
# .to(device)/.to('cpu'), which fights the FSDP flat_param. The frozen init_rollout
# (the low expert's high-noise teacher) is sharded here too; it never enters
# accelerator.prepare, so it stays out of the optimizer group.
shard_fn = partial(
shard_model,
device_id=accelerator.device,
param_dtype=weight_dtype,
offload_to_cpu=args.low_vram,
)
real_score_transformer3d = shard_fn(real_score_transformer3d)
if init_rollout_transformer3d is not None:
init_rollout_transformer3d = shard_fn(init_rollout_transformer3d)
text_encoder = shard_fn(text_encoder)
# Move text_encoder and vae to gpu and cast to weight_dtype.
# vae is never sharded and needs explicit placement; sharded modules are owned by FSDP,
# so skip .to() on them to keep the flat_param attached.
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if fsdp_stage == 0 and zero_stage == 0:
real_score_transformer3d.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
if init_rollout_transformer3d is not None:
init_rollout_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")
# 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)
if args.dfd:
# The student timesteps are the denoising-step sigmas of the shift schedule (normalized RF times).
dfd_student_timesteps = noise_scheduler.timesteps[args.train_sampling_steps - torch.tensor(args.denoising_step_indices_list)].to(accelerator.device) / 1000
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
train_div_loss = 0.0
train_denoising_loss = 0.0
train_dfd_real_replace = 0.0
dfd_real_replace_now = 0.0
train_ca_scale = 0.0
train_dm_scale = 0.0
train_latent_std = 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
batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch)
for step, batch in enumerate(train_dataloader):
generator_update = step % args.gen_update_interval == 0
dfd_active = args.dfd and global_step >= args.dfd_start_step
# 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):
real_latents = None
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)
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)
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),
)
# DFD encodes the paired real video as the student anchor.
if dfd_active:
pixel_values = rearrange(batch["pixel_values"].to(accelerator.device, dtype=weight_dtype, non_blocking=True), "b f c h w -> b c f h w")
if args.low_vram:
vae.to(accelerator.device)
with torch.no_grad():
real_latents = vae.encode(pixel_values)[0].sample()
if dfd_active:
target_shape = real_latents.size()
# CrossDistill PCM diversity: encode the paired real video as the phase-0 anchor x0
# (the official _compute_pcm_div_loss anchors on the real data latent, not G.detach()).
if real_latents is None and args.loss_scale_div > 0 and args.div_mode in ("pcm_fixed", "pcm_random") \
and 'pixel_values' in batch and batch.get('pixel_values', None) is not None:
pixel_values = rearrange(batch["pixel_values"].to(accelerator.device, dtype=weight_dtype, non_blocking=True), "b f c h w -> b c f h w")
if args.low_vram:
vae.to(accelerator.device)
with torch.no_grad():
real_latents = vae.encode(pixel_values)[0].sample()
target_shape = real_latents.size()
if args.low_vram:
vae.to('cpu')
real_score_transformer3d = real_score_transformer3d.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()
# 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)
# --- rf-native grid helpers (dual-expert Wan2.2). The rectified-flow sigma equals the
# rf time s, and the rCM grid knots are NOT on the shift=12 scheduler grid, so the rollout
# below drives the nets with t = s * T_SCALE and steps/inverts with s directly, bypassing
# the scheduler timestep-index map used by get_sigmas / convert_flow_pred_to_x0. ---
def rf_flow_to_x0(flow_pred, xt, rf):
# Invert x_s = (1 - s) * x0 + s * eps => x0 = x_s - s * v (v = eps - x0), in double
# to match convert_flow_pred_to_x0, then cast back to the model output dtype.
return (xt.double() - float(rf) * flow_pred.double()).to(flow_pred.dtype)
def rf_add_noise(latents, noise, rf):
# rf-native forward interpolation x_s = (1 - s) * x0 + s * noise (s == sigma == rf time).
s = float(rf)
return (1.0 - s) * latents + s * noise
# --- CrossDistill PCM diversity helpers (faithful port of the official
# _compute_pcm_div_loss). rf time s == scheduler sigma == timestep / 1000, so the
# student/teacher nets are driven with t = s * 1000 exactly as the DMD path does. ---
T_SCALE = 1000.0
def shift_rf_time(u, shift):
return shift * u / (1.0 + (shift - 1.0) * u)
def inv_shift_rf_time(rf, shift):
return rf / (shift - (shift - 1.0) * rf) if shift > 0 else rf
def net_velocity(net, x_5d, s_B_1, embeds, slen):
t = (s_B_1.reshape(-1) * T_SCALE).to(torch.float32)
# The div forwards run outside the step-level autocast (the DMD rollout and DFD student
# each open their own), so wrap the net call here; otherwise the fp32 t / fp32 real-latent
# input hits the bf16 FSDP weights and F.linear raises "mat1 and mat2 must have same dtype".
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
return net(
x=x_5d,
context=embeds,
t=t,
seq_len=slen,
y=inpaint_latents if args.train_mode != "normal" else None,
).float()
def _euler(net, x, s_cur, s_next, embeds, slen):
v = net_velocity(net, x, s_cur, embeds, slen)
sc = s_cur.reshape(-1, 1, 1, 1, 1).to(x.dtype)
sn = s_next.reshape(-1, 1, 1, 1, 1).to(x.dtype)
return x + (sn - sc) * v.to(x.dtype)
def compute_pcm_div_loss(x0, embeds, neg_embeds, slen):
# PCM phase-0 boundary alignment on the real-x0 anchor: the student one-step jump
# G(x_t0, t0)->boundary must match the teacher K-step Euler + student target jump
# G(x_tK, tK)->boundary. Self-contained (own eps + own continuous rf grid), so it
# layers onto the discrete DMD rollout without touching it.
B = x0.shape[0]
K = args.dcm_skipping_interval_steps
dw = 1.0 / args.dcm_total_steps
shift = args.teacher_timestep_shift
dev = accelerator.device
eps = torch.randn(x0.size(), device=dev, generator=torch_rng, dtype=weight_dtype)
rf_b = float(args.div_rf_boundary)
s_b = inv_shift_rf_time(rf_b, shift)
valid_width = (1.0 - s_b) - K * dw
if args.div_mode == "pcm_fixed":
w = torch.zeros(B, 1, device=dev)
else:
assert valid_width > 0, f"PCM-div phase-0 width <= 0: {valid_width:.4f}"
w = torch.rand(B, 1, device=dev, generator=torch_rng) * valid_width
t_list = [shift_rf_time(1.0 - (w + k * dw), shift) for k in range(K + 1)]
t0, tK = t_list[0], t_list[-1]
t_bnd = torch.full((B, 1), rf_b, device=dev)
t0_5d = t0.reshape(-1, 1, 1, 1, 1).to(x0.dtype)
xt = (1.0 - t0_5d) * x0 + t0_5d * eps
# no_grad teacher rollout + student target BEFORE the grad student forward (FSDP-safe).
with torch.no_grad():
xk = xt
for k in range(K):
dt = (t_list[k] - t_list[k + 1]).reshape(-1, 1, 1, 1, 1).to(xk.dtype)
vc = net_velocity(real_score_transformer3d, xk, t_list[k], embeds, slen)
if args.div_teacher_guidance > 1.0:
vu = net_velocity(real_score_transformer3d, xk, t_list[k], neg_embeds, slen)
vc = vu + args.div_teacher_guidance * (vc - vu)
xk = xk - dt * vc.to(xk.dtype)
x_tgt = _euler(generator_transformer3d, xk, tK, t_bnd, embeds, slen).detach()
x_pred = _euler(generator_transformer3d, xt, t0, t_bnd, embeds, slen)
loss = (x_pred.double() - x_tgt.double()) ** 2
loss[torch.isnan(loss).flatten(start_dim=1).any(dim=1)] = 0
return loss.mean()
# --- Main Training Logic ---
bsz, channel, num_frames, height, width = target_shape
# 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)
if dfd_active:
# Sample one shared shifted RF timestep for both phases of this step.
score_indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
dfd_score_timestep = noise_scheduler.timesteps[score_indices].to(device=accelerator.device)
dfd_score_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=real_latents.dtype)
# Noise the paired real latent to one of the 4-step student timesteps and denoise it.
student_timestep_indices = torch.randint(0, dfd_student_timesteps.numel(), (bsz,), device=accelerator.device, generator=torch_rng)
student_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=real_latents.dtype)
student_timestep = dfd_student_timesteps[student_timestep_indices] * 1000
# Generator's current-step noise level t, used to bound tau_CA in Decoupled DMD.
generator_step_timestep = student_timestep
student_input = add_noise(real_latents, student_noise, student_timestep)
student_forward_context = contextlib.nullcontext() if generator_update else torch.no_grad()
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device), student_forward_context:
dfd_student_flow = generator_transformer3d(
x=student_input,
context=prompt_embeds,
t=student_timestep,
seq_len=seq_len,
y=None,
)
dfd_student_pred = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=dfd_student_flow,
xt=student_input,
timestep=student_timestep,
)
else:
# rf-native discrete grid for this expert, derived in main() from --joint_backward_timesteps
# split at the MoE boundary. The rf values are used DIRECTLY as Euler sigmas (the shift=12
# scheduler grid does not contain the rCM knots, so the timestep-index map is bypassed).
denoising_step_list = denoising_rf_grid
if generator_update:
if dfd_active:
generator_pred = dfd_student_pred
generator_timestep = dfd_score_timestep
dmd_noise = dfd_score_noise
else:
generator_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype)
# Low-noise expert: first roll the FROZEN high student from pure noise down to the
# boundary (init_rollout), so the low rollout starts on-distribution at rf=boundary
# instead of cold-starting from pure noise (which would break train/inference routing).
if init_rollout_transformer3d is not None:
with torch.no_grad():
for _s_cur, _s_next in zip(init_rollout_grid[:-1], init_rollout_grid[1:]):
_sc = torch.full((bsz, 1), float(_s_cur), device=accelerator.device, dtype=torch.float32)
_sn = torch.full((bsz, 1), float(_s_next), device=accelerator.device, dtype=torch.float32)
generator_noise = _euler(init_rollout_transformer3d, generator_noise, _sc, _sn, prompt_embeds, seq_len)
generator_noise = generator_noise.detach()
num_denoising_steps = len(denoising_step_list)
final_step_index = generate_and_sync_list(num_denoising_steps, device=generator_noise.device)[0]
for index, current_rf in enumerate(denoising_step_list):
is_final_step = (index == final_step_index)
rf_cur = float(current_rf)
timestep = torch.full(
generator_noise.shape[:1],
rf_cur * T_SCALE,
device=generator_noise.device,
dtype=torch.float32
)
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,
)
if is_final_step and args.tdm:
# TDM: the last segment does NOT collapse to x0; take one more rf-Euler
# step to s_{i-1} (the next knot, 0.0 when the grid is exhausted, which
# degenerates to DMD), so the anchor is G_x_{t_{i-1}} instead of G_x0.
s_im1 = float(denoising_step_list[index + 1]) if index + 1 < num_denoising_steps else 0.0
G_x_t_im1 = (
generator_noise.float() + (s_im1 - rf_cur) * generator_pred.float()
).to(generator_noise.dtype)
t_im1 = s_im1
generator_step_timestep = timestep
break
if not args.flow_euler_rollout or is_final_step:
generator_pred = rf_flow_to_x0(generator_pred, generator_noise, rf_cur)
if is_final_step:
# Generator's current-step noise level t, used to bound tau_CA in Decoupled DMD.
generator_step_timestep = timestep
break
rf_next = float(denoising_step_list[index + 1])
if args.flow_euler_rollout:
# rf-native Euler ODE step (sigma == rf time s): keep the prediction in
# flow/velocity space and advance x_next = x_t + (s_next - s_cur) * v (fp32).
generator_noise = (
generator_noise.float() + (rf_next - rf_cur) * generator_pred.float()
).to(generator_noise.dtype)
else:
generator_noise = rf_add_noise(
generator_pred,
torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng),
rf_next
)
indices = idx_sampling(bsz, generator=torch_rng, device=accelerator.device).long().cpu()
generator_timestep = noise_scheduler.timesteps[indices].to(device=accelerator.device)
if args.tdm:
# TDM: D_time clamped STRICTLY above s_{i-1} (=t_im1). Resample uniformly over the
# shift-grid timesteps > t_im1 (kept on-grid so convert_flow_pred_to_x0 stays exact);
# when t_im1 == 0 (grid exhausted) this is the full range and TDM degenerates to DMD.
_grid_ts = noise_scheduler.timesteps.to(device=accelerator.device)
_k = int((_grid_ts > float(t_im1) * T_SCALE).sum().item())
indices = torch.randint(0, max(_k, 1), (bsz,), generator=torch_rng, device=accelerator.device)
generator_timestep = _grid_ts[indices].to(device=accelerator.device)
dmd_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
if args.tdm:
# rf forward transition q(x_{s_D} | x_{s_{i-1}}) with G = G_x_{t_{i-1}} (NOT x0):
# D_xt = ((1-s_D)/(1-s1))*G + sqrt((s_D-s1)(s_D+s1-2 s_D s1))/(1-s1) * eps
_sD = (generator_timestep / T_SCALE).to(torch.float32).reshape(-1, 1, 1, 1, 1)
_s1 = torch.full_like(_sD, float(t_im1))
_oms1 = (1.0 - _s1).clamp(min=1e-6)
_sig = torch.sqrt(((_sD - _s1) * (_sD + _s1 - 2.0 * _sD * _s1)).clamp(min=0.0)) / _oms1
generator_denoised_input = (
((1.0 - _sD) / _oms1) * G_x_t_im1.to(torch.float32) + _sig * dmd_noise.to(torch.float32)
).detach().to(accelerator.device, dtype=weight_dtype)
else:
generator_denoised_input = add_noise(
generator_pred,
dmd_noise,
generator_timestep
).detach().to(accelerator.device, dtype=weight_dtype)
if args.decoupled_dmd:
# Decoupled DMD (arXiv:2511.22677, Sec. 4.3): the CA engine must re-noise at
# tau_CA > t, i.e. noise levels cleaner than the generator's current step, so it
# only enhances not-yet-resolved (higher-frequency) content. Scheduler timesteps
# are ordered from noisy to clean, hence sample indices strictly after t's position.
schedule_timesteps = noise_scheduler.timesteps.to(device=accelerator.device)
num_schedule_steps = schedule_timesteps.numel()
step_positions = torch.argmin(
(schedule_timesteps.unsqueeze(0).float() - generator_step_timestep.reshape(-1, 1).float()).abs(),
dim=1
)
# Uniform over the schedule positions cleaner than t, matching how tau_DM is
# sampled uniformly over the full set of schedule positions.
ca_low = (step_positions + 1).clamp(max=num_schedule_steps - 1)
ca_span = num_schedule_steps - ca_low
ca_offset = (
torch.rand(ca_low.shape, device=accelerator.device, generator=torch_rng) * ca_span
).long()
ca_positions = (ca_low + ca_offset).clamp(max=num_schedule_steps - 1)
ca_timestep = schedule_timesteps[ca_positions]
ca_noise = torch.randn(generator_pred.shape, dtype=generator_pred.dtype, device=generator_pred.device, generator=torch_rng)
generator_ca_input = add_noise(
generator_pred,
ca_noise,
ca_timestep
).detach().to(accelerator.device, dtype=weight_dtype)
# DFD may feed the frozen teacher the paired real latent instead.
real_score_input = generator_denoised_input
use_dfd_real = dfd_active and real_latents is not None and torch.rand((), device=accelerator.device, generator=torch_rng).item() < args.dfd_teacher_replace_prob
if use_dfd_real:
real_score_input = add_noise(
real_latents.to(generator_pred.dtype),
dmd_noise,
generator_timestep
).detach().to(accelerator.device, dtype=weight_dtype)
train_dfd_real_replace += 1.0
# 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,
)
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,
)
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 (conditional branch always evaluated on the DM input tau_DM)
real_score_main_cond = real_score_transformer3d(
x=real_score_input,
context=prompt_embeds,
t=generator_timestep,
seq_len=seq_len,
y=inpaint_latents 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=real_score_input,
timestep=generator_timestep
)
if args.decoupled_dmd:
# CA term: teacher cond/uncond evaluated at the cleaner re-noise level tau_CA.
real_score_ca_cond = real_score_transformer3d(
x=generator_ca_input,
context=prompt_embeds,
t=ca_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
)
real_score_ca_cond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=real_score_ca_cond,
xt=generator_ca_input,
timestep=ca_timestep
)
real_score_ca_uncond = real_score_transformer3d(
x=generator_ca_input,
context=neg_prompt_embeds,
t=ca_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
)
real_score_ca_uncond = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=real_score_ca_uncond,
xt=generator_ca_input,
timestep=ca_timestep
)
else:
real_score_main_uncond = real_score_transformer3d(
x=real_score_input,
context=neg_prompt_embeds,
t=generator_timestep,
seq_len=seq_len,
y=inpaint_latents 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=real_score_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
if args.decoupled_dmd:
# Decoupled DMD (arXiv:2511.22677, Eq. 8) splits the teacher target into
# DM term (shield/regularizer): s_cond(tau_DM) - s_fake(tau_DM), full noise range;
# CA term (spear/engine): (alpha - 1) * (s_cond(tau_CA) - s_uncond(tau_CA)), tau_CA > t.
# Composing them into one target keeps Eq. 8's (alpha - 1) relative weight under the
# single adaptive normalizer of DMD, and reduces exactly to s_cfg when tau_CA == tau_DM.
ca_delta = (
real_score_ca_cond - real_score_ca_uncond
) * (args.real_guidance_scale - 1.0)
real_score_target = real_score_main_cond + ca_delta
else:
real_score_target = real_score_main
ca_delta = (
real_score_main_cond - real_score_main_uncond
) * (args.real_guidance_scale - 1.0)
# Magnitudes of the two Eq. 6 terms, recomputed from the frozen score tensors so the
# monitor costs no extra forward pass and never touches the graph.
ca_scale = ca_delta.float().abs().mean()
dm_scale = (real_score_main_cond - fake_score_main).float().abs().mean()
if args.tdm:
# TDM: compare in the x_{t_{i-1}} domain. Map the fake/teacher x0 preds back to s_{i-1}
# via the rf PF-ODE analytic step, then reuse the DMD-style score-distillation surrogate
# with the adaptive weight_factor = mean|G - x_{t_{i-1}}^teacher| (clip 1e-5).
_sD = (generator_timestep / T_SCALE).to(torch.float32).reshape(-1, 1, 1, 1, 1)
_s1 = torch.full_like(_sD, float(t_im1))
def _to_t_im1(x_t, x0_pred):
eps_pred = (x_t - (1.0 - _sD) * x0_pred) / _sD.clamp(min=1e-6)
return (1.0 - _s1) * x0_pred + _s1 * eps_pred
x_t_im1_fake = _to_t_im1(generator_denoised_input.float(), fake_score_main.float())
x_t_im1_teacher = _to_t_im1(generator_denoised_input.float(), real_score_target.float())
Gd = G_x_t_im1.double()
weight_factor = (Gd - x_t_im1_teacher.double()).abs().mean(dim=[1, 2, 3, 4], keepdim=True).clip(min=1e-5)
tdm_grad = (x_t_im1_fake.double() - x_t_im1_teacher.double()) / weight_factor
dmd_loss = 0.5 * F.mse_loss(Gd, (Gd - tdm_grad).detach(), reduction="mean")
generator_pred = G_x_t_im1 # the monitor .std() below tracks the TDM anchor
else:
fake_to_real_grad = fake_score_main - real_score_target
if dfd_active:
normalizer = 1.0 / ((generator_pred.float() - real_score_target.float()).abs().mean(dim=[1, 2, 3, 4], keepdim=True) + 1e-6)
fake_to_real_grad = fake_to_real_grad * normalizer.to(fake_to_real_grad.dtype)
dmd_loss = 0.5 * F.mse_loss(
generator_pred,
(generator_pred - fake_to_real_grad).detach(),
reduction="mean"
)
else:
generator_to_real_norm = generator_pred - real_score_target
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
monitor = accelerator.gather(
torch.stack([ca_scale, dm_scale, generator_pred.float().std()])[None]
).mean(dim=0)
train_ca_scale += monitor[0].item()
train_dm_scale += monitor[1].item()
train_latent_std += monitor[2].item()
# --- CrossDistill PCM diversity (self-contained additive term; real-x0 anchor) ---
div_loss = torch.zeros((), device=accelerator.device, dtype=torch.float64)
if args.loss_scale_div > 0 and args.div_mode in ("pcm_fixed", "pcm_random") and real_latents is not None:
div_loss = compute_pcm_div_loss(real_latents, prompt_embeds, neg_prompt_embeds, seq_len)
train_div_loss += accelerator.gather(div_loss.float().repeat(args.train_batch_size)).mean().item()
if args.low_vram:
real_score_transformer3d = real_score_transformer3d.to("cpu")
fake_score_transformer3d = fake_score_transformer3d.to("cpu")
torch.cuda.empty_cache()
# Single merged backward: FSDP flat_param forbids two accelerator.backward calls
# on the same wrapped generator within one iteration.
generator_loss = dmd_loss + args.loss_scale_div * div_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.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():
if dfd_active:
fake_score_denoised_pred = dfd_student_pred.detach()
critic_timestep = dfd_score_timestep
critic_noise = dfd_score_noise
else:
fake_score_critic_noise = torch.randn(target_shape, device=accelerator.device, generator=torch_rng, dtype=weight_dtype)
# Low-noise expert: same init_rollout prefix as the generator phase (frozen high student,
# pure noise -> boundary) so the critic trains on the student's on-distribution start latent.
# The whole critic rollout already runs under torch.no_grad().
if init_rollout_transformer3d is not None:
for _s_cur, _s_next in zip(init_rollout_grid[:-1], init_rollout_grid[1:]):
_sc = torch.full((bsz, 1), float(_s_cur), device=accelerator.device, dtype=torch.float32)
_sn = torch.full((bsz, 1), float(_s_next), device=accelerator.device, dtype=torch.float32)
fake_score_critic_noise = _euler(init_rollout_transformer3d, fake_score_critic_noise, _sc, _sn, prompt_embeds, seq_len)
fake_score_critic_noise = fake_score_critic_noise.detach()
num_denoising_steps = len(denoising_step_list)
final_step_index = generate_and_sync_list(num_denoising_steps, device=fake_score_critic_noise.device)[0]
for index, current_rf in enumerate(denoising_step_list):
is_final_step = (index == final_step_index)
rf_cur = float(current_rf)
timestep = torch.full(
fake_score_critic_noise.shape[:1],
rf_cur * T_SCALE,
device=fake_score_critic_noise.device,
dtype=torch.float32
)
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,
)
fake_score_denoised_pred = rf_flow_to_x0(fake_score_denoised_pred, fake_score_critic_noise, rf_cur)
if is_final_step:
break
rf_next = float(denoising_step_list[index + 1])
fake_score_critic_noise = rf_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),
rf_next
)
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,
)
fake_score_denoised_output_x0 = convert_flow_pred_to_x0(
scheduler=noise_scheduler,
flow_pred=fake_score_denoised_output,
xt=fake_score_denoised_input,
timestep=critic_timestep
)
denoising_loss = F.mse_loss(
fake_score_denoised_output_x0.float(),
fake_score_denoised_pred.detach().float(),
reduction="mean"
)
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
gen_log_div = max(train_gen_log_count, 1)
tracker_logs = {"train_denoising_loss": train_denoising_loss, "train_dmd_loss": train_dmd_loss / gen_log_div}
tracker_logs["train_div_loss"] = train_div_loss / gen_log_div
tracker_logs["train_ca_scale"] = train_ca_scale / gen_log_div
tracker_logs["train_dm_scale"] = train_dm_scale / gen_log_div
tracker_logs["train_ca_dm_ratio"] = train_ca_scale / (train_dm_scale + 1e-12)
tracker_logs["train_latent_std"] = train_latent_std / gen_log_div
if args.dfd:
tracker_logs["train_dfd_real_replace"] = train_dfd_real_replace
dfd_real_replace_now = train_dfd_real_replace
accelerator.log(tracker_logs, step=global_step)
train_dmd_loss = 0.0
train_div_loss = 0.0
train_denoising_loss = 0.0
train_dfd_real_replace = 0.0
train_ca_scale = 0.0
train_dm_scale = 0.0
train_latent_std = 0.0
train_gen_log_count = 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,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
init_rollout_transformer3d,
)
if args.dfd:
logs = {"lr": (lr_scheduler if generator_update else fake_score_lr_scheduler).get_last_lr()[0], "denoising_loss": denoising_loss.detach().item()}
if generator_update:
logs["dmd_loss"] = dmd_loss.detach().item()
logs["ca"] = ca_scale.item()
logs["dm"] = dm_scale.item()
logs["dfd_real_replace"] = dfd_real_replace_now
else:
logs = {"denoising_loss": denoising_loss.detach().item(), "dmd_loss": dmd_loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
logs["ca"] = ca_scale.item()
logs["dm"] = dm_scale.item()
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,
generator_transformer3d,
args,
config,
accelerator,
weight_dtype,
global_step,
init_rollout_transformer3d,
)
# 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()
Reference in New Issue View Git Blame Copy Permalink