3009 lines
154 KiB
Python
3009 lines
154 KiB
Python
"""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, CLIPModel, SparkWanTransformer3DModel,
|
|
WanT5EncoderModel, WanTransformer3DModel)
|
|
from videox_fun.pipeline import WanI2VPipeline, WanPipeline
|
|
from videox_fun.utils.discrete_sampler import DiscreteSampling
|
|
from videox_fun.utils.tqdm_bar import PauseAwareTqdm
|
|
from videox_fun.utils.utils import (calculate_dimensions, get_image_latent,
|
|
get_image_to_video_latent,
|
|
save_videos_grid)
|
|
|
|
if is_wandb_available():
|
|
import wandb
|
|
|
|
|
|
def filter_kwargs(cls, kwargs):
|
|
import inspect
|
|
sig = inspect.signature(cls.__init__)
|
|
valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
|
|
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
|
|
return filtered_kwargs
|
|
|
|
def get_random_downsample_ratio(sample_size, image_ratio=[],
|
|
all_choices=False, rng=None):
|
|
def _create_special_list(length):
|
|
if length == 1:
|
|
return [1.0]
|
|
if length >= 2:
|
|
first_element = 0.75
|
|
remaining_sum = 1.0 - first_element
|
|
other_elements_value = remaining_sum / (length - 1)
|
|
special_list = [first_element] + [other_elements_value] * (length - 1)
|
|
return special_list
|
|
|
|
if sample_size >= 1536:
|
|
number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio
|
|
elif sample_size >= 1024:
|
|
number_list = [1, 1.25, 1.5, 2] + image_ratio
|
|
elif sample_size >= 768:
|
|
number_list = [1, 1.25, 1.5] + image_ratio
|
|
elif sample_size >= 512:
|
|
number_list = [1] + image_ratio
|
|
else:
|
|
number_list = [1]
|
|
|
|
if all_choices:
|
|
return number_list
|
|
|
|
number_list_prob = np.array(_create_special_list(len(number_list)))
|
|
if rng is None:
|
|
return np.random.choice(number_list, p = number_list_prob)
|
|
else:
|
|
return rng.choice(number_list, p = number_list_prob)
|
|
|
|
def resize_mask(mask, latent, process_first_frame_only=True):
|
|
latent_size = latent.size()
|
|
batch_size, channels, num_frames, height, width = mask.shape
|
|
|
|
if process_first_frame_only:
|
|
target_size = list(latent_size[2:])
|
|
target_size[0] = 1
|
|
first_frame_resized = F.interpolate(
|
|
mask[:, :, 0:1, :, :],
|
|
size=target_size,
|
|
mode='trilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
target_size = list(latent_size[2:])
|
|
target_size[0] = target_size[0] - 1
|
|
if target_size[0] != 0:
|
|
remaining_frames_resized = F.interpolate(
|
|
mask[:, :, 1:, :, :],
|
|
size=target_size,
|
|
mode='trilinear',
|
|
align_corners=False
|
|
)
|
|
resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2)
|
|
else:
|
|
resized_mask = first_frame_resized
|
|
else:
|
|
target_size = list(latent_size[2:])
|
|
resized_mask = F.interpolate(
|
|
mask,
|
|
size=target_size,
|
|
mode='trilinear',
|
|
align_corners=False
|
|
)
|
|
return resized_mask
|
|
|
|
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
|
check_min_version("0.18.0.dev0")
|
|
|
|
logger = get_logger(__name__, log_level="INFO")
|
|
|
|
def log_validation(vae, text_encoder, tokenizer, clip_image_encoder, transformer3d, args, config, accelerator, weight_dtype, global_step):
|
|
try:
|
|
is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine'
|
|
if is_deepspeed:
|
|
origin_config = transformer3d.config
|
|
transformer3d.config = accelerator.unwrap_model(transformer3d).config
|
|
with torch.no_grad(), torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
|
|
logger.info("Running validation... ")
|
|
scheduler = FlowMatchEulerDiscreteScheduler(
|
|
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
|
|
)
|
|
|
|
if args.train_mode != "normal":
|
|
pipeline = WanI2VPipeline(
|
|
vae=vae,
|
|
text_encoder=text_encoder,
|
|
tokenizer=tokenizer,
|
|
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
|
|
scheduler=scheduler,
|
|
clip_image_encoder=clip_image_encoder,
|
|
)
|
|
else:
|
|
pipeline = WanPipeline(
|
|
vae=vae,
|
|
text_encoder=text_encoder,
|
|
tokenizer=tokenizer,
|
|
transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d,
|
|
scheduler=scheduler,
|
|
)
|
|
pipeline = pipeline.to(accelerator.device)
|
|
|
|
if args.seed is None:
|
|
generator = None
|
|
else:
|
|
rank_seed = args.seed + accelerator.process_index
|
|
generator = torch.Generator(device=accelerator.device).manual_seed(rank_seed)
|
|
logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}")
|
|
|
|
for i in range(len(args.validation_prompts)):
|
|
if args.train_mode != "normal":
|
|
start_image = Image.open(args.validation_paths[i])
|
|
width, height = start_image.width, start_image.height
|
|
width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height)
|
|
|
|
video_length = int((args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if args.video_sample_n_frames != 1 else 1
|
|
input_video, input_video_mask, _ = get_image_to_video_latent(args.validation_paths[i], None, video_length=video_length, sample_size=[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 = args.student_num_steps,
|
|
guidance_scale = 1.0,
|
|
).videos
|
|
|
|
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
|
|
save_videos_grid(
|
|
sample,
|
|
os.path.join(
|
|
args.output_dir,
|
|
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.mp4"
|
|
)
|
|
)
|
|
else:
|
|
# 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 = args.student_num_steps,
|
|
guidance_scale = 1.0,
|
|
).videos
|
|
os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True)
|
|
save_videos_grid(
|
|
sample,
|
|
os.path.join(
|
|
args.output_dir,
|
|
f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.mp4"
|
|
)
|
|
)
|
|
|
|
del pipeline
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
|
|
if not args.enable_text_encoder_in_dataloader:
|
|
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
|
|
if is_deepspeed:
|
|
transformer3d.config = origin_config
|
|
except Exception as e:
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
print(f"Eval error on rank {accelerator.process_index} with info {e}")
|
|
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
|
|
if not args.enable_text_encoder_in_dataloader:
|
|
text_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
|
|
|
|
def linear_decay(initial_value, final_value, total_steps, current_step):
|
|
if current_step >= total_steps:
|
|
return final_value
|
|
current_step = max(0, current_step)
|
|
step_size = (final_value - initial_value) / total_steps
|
|
current_value = initial_value + step_size * current_step
|
|
return current_value
|
|
|
|
def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=None):
|
|
u = torch.normal(mean=0.0, std=1.0, size=shape, device=device, generator=generator)
|
|
t = 1 / (1 + torch.exp(-u)) * (high - low) + low
|
|
return torch.clip(t.to(torch.int32), low, high - 1)
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
|
parser.add_argument(
|
|
"--input_perturbation", type=float, default=0, help="The scale of input perturbation. Recommended 0.1."
|
|
)
|
|
parser.add_argument(
|
|
"--pretrained_model_name_or_path",
|
|
type=str,
|
|
default=None,
|
|
required=True,
|
|
help="Path to pretrained model or model identifier from huggingface.co/models.",
|
|
)
|
|
parser.add_argument(
|
|
"--revision",
|
|
type=str,
|
|
default=None,
|
|
required=False,
|
|
help="Revision of pretrained model identifier from huggingface.co/models.",
|
|
)
|
|
parser.add_argument(
|
|
"--variant",
|
|
type=str,
|
|
default=None,
|
|
help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
|
|
)
|
|
parser.add_argument(
|
|
"--train_data_dir",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"A folder containing the training data. "
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--train_data_meta",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"A csv containing the training data. "
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--max_train_samples",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For debugging purposes or quicker training, truncate the number of training examples to this "
|
|
"value if set."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--validation_prompts",
|
|
type=str,
|
|
default=None,
|
|
nargs="+",
|
|
help=("A set of prompts evaluated every `--validation_epochs` and logged to `--report_to`."),
|
|
)
|
|
parser.add_argument(
|
|
"--validation_paths",
|
|
type=str,
|
|
default=None,
|
|
nargs="+",
|
|
help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."),
|
|
)
|
|
parser.add_argument(
|
|
"--negative_prompt",
|
|
type=str,
|
|
default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
|
|
help=("The negative prompt of cfg distill"),
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir",
|
|
type=str,
|
|
default="sd-model-finetuned",
|
|
help="The output directory where the model predictions and checkpoints will be written.",
|
|
)
|
|
parser.add_argument(
|
|
"--cache_dir",
|
|
type=str,
|
|
default=None,
|
|
help="The directory where the downloaded models and datasets will be stored.",
|
|
)
|
|
parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
|
|
parser.add_argument(
|
|
"--random_flip",
|
|
action="store_true",
|
|
help="whether to randomly flip images horizontally",
|
|
)
|
|
parser.add_argument(
|
|
"--use_came",
|
|
action="store_true",
|
|
help="whether to use came",
|
|
)
|
|
parser.add_argument(
|
|
"--multi_stream",
|
|
action="store_true",
|
|
help="whether to use cuda multi-stream",
|
|
)
|
|
parser.add_argument(
|
|
"--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
|
|
)
|
|
parser.add_argument(
|
|
"--vae_mini_batch", type=int, default=32, help="mini batch size for vae."
|
|
)
|
|
parser.add_argument("--num_train_epochs", type=int, default=100)
|
|
parser.add_argument(
|
|
"--max_train_steps",
|
|
type=int,
|
|
default=None,
|
|
help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
|
|
)
|
|
parser.add_argument(
|
|
"--gradient_accumulation_steps",
|
|
type=int,
|
|
default=1,
|
|
help="Number of updates steps to accumulate before performing a backward/update pass.",
|
|
)
|
|
parser.add_argument(
|
|
"--gradient_checkpointing",
|
|
action="store_true",
|
|
help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
|
|
)
|
|
parser.add_argument(
|
|
"--learning_rate",
|
|
type=float,
|
|
default=1e-4,
|
|
help="Initial learning rate (after the potential warmup period) to use.",
|
|
)
|
|
parser.add_argument(
|
|
"--learning_rate_critic",
|
|
type=float,
|
|
default=1e-4,
|
|
help="Initial learning rate (after the potential warmup period) to use.",
|
|
)
|
|
parser.add_argument(
|
|
"--scale_lr",
|
|
action="store_true",
|
|
default=False,
|
|
help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
|
|
)
|
|
parser.add_argument(
|
|
"--lr_scheduler",
|
|
type=str,
|
|
default="constant",
|
|
help=(
|
|
'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
|
|
' "constant", "constant_with_warmup"]'
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
|
|
)
|
|
parser.add_argument(
|
|
"--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
|
|
)
|
|
parser.add_argument(
|
|
"--allow_tf32",
|
|
action="store_true",
|
|
help=(
|
|
"Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
|
|
" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--non_ema_revision",
|
|
type=str,
|
|
default=None,
|
|
required=False,
|
|
help=(
|
|
"Revision of pretrained non-ema model identifier. Must be a branch, tag or git identifier of the local or"
|
|
" remote repository specified with --pretrained_model_name_or_path."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--dataloader_num_workers",
|
|
type=int,
|
|
default=0,
|
|
help=(
|
|
"Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
|
|
),
|
|
)
|
|
parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
|
|
parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
|
|
parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
|
|
parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
|
|
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
|
|
parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
|
|
parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
|
|
parser.add_argument(
|
|
"--prediction_type",
|
|
type=str,
|
|
default=None,
|
|
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.",
|
|
)
|
|
parser.add_argument(
|
|
"--hub_model_id",
|
|
type=str,
|
|
default=None,
|
|
help="The name of the repository to keep in sync with the local `output_dir`.",
|
|
)
|
|
parser.add_argument(
|
|
"--logging_dir",
|
|
type=str,
|
|
default="logs",
|
|
help=(
|
|
"[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
|
|
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--report_model_info", action="store_true", help="Whether or not to report more info about model (such as norm, grad)."
|
|
)
|
|
parser.add_argument(
|
|
"--mixed_precision",
|
|
type=str,
|
|
default=None,
|
|
choices=["no", "fp16", "bf16"],
|
|
help=(
|
|
"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
|
|
" 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the"
|
|
" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--report_to",
|
|
type=str,
|
|
default="tensorboard",
|
|
help=(
|
|
'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
|
|
' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
|
|
),
|
|
)
|
|
parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
|
|
parser.add_argument(
|
|
"--checkpointing_steps",
|
|
type=int,
|
|
default=500,
|
|
help=(
|
|
"Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming"
|
|
" training using `--resume_from_checkpoint`."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--checkpoints_total_limit",
|
|
type=int,
|
|
default=None,
|
|
help=("Max number of checkpoints to store."),
|
|
)
|
|
parser.add_argument(
|
|
"--resume_from_checkpoint",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"Whether training should be resumed from a previous checkpoint. Use a path saved by"
|
|
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
|
|
),
|
|
)
|
|
parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.")
|
|
parser.add_argument(
|
|
"--validation_epochs",
|
|
type=int,
|
|
default=5,
|
|
help="Run validation every X epochs.",
|
|
)
|
|
parser.add_argument(
|
|
"--validation_steps",
|
|
type=int,
|
|
default=2000,
|
|
help="Run validation every X steps.",
|
|
)
|
|
parser.add_argument(
|
|
"--tracker_project_name",
|
|
type=str,
|
|
default="text2image-fine-tune",
|
|
help=(
|
|
"The `project_name` argument passed to Accelerator.init_trackers for"
|
|
" more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator"
|
|
),
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--snr_loss", action="store_true", help="Whether or not to use snr_loss."
|
|
)
|
|
parser.add_argument(
|
|
"--uniform_sampling", action="store_true", help="Whether or not to use uniform_sampling."
|
|
)
|
|
parser.add_argument(
|
|
"--enable_text_encoder_in_dataloader", action="store_true", help="Whether or not to use text encoder in dataloader."
|
|
)
|
|
parser.add_argument(
|
|
"--enable_bucket", action="store_true", help="Whether enable bucket sample in datasets."
|
|
)
|
|
parser.add_argument(
|
|
"--random_ratio_crop", action="store_true", help="Whether enable random ratio crop sample in datasets."
|
|
)
|
|
parser.add_argument(
|
|
"--random_frame_crop", action="store_true", help="Whether enable random frame crop sample in datasets."
|
|
)
|
|
parser.add_argument(
|
|
"--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets."
|
|
)
|
|
parser.add_argument(
|
|
"--training_with_video_token_length", action="store_true", help="The training stage of the model in training.",
|
|
)
|
|
parser.add_argument(
|
|
"--auto_tile_batch_size", action="store_true", help="Whether to auto tile batch size.",
|
|
)
|
|
parser.add_argument(
|
|
"--motion_sub_loss", action="store_true", help="Whether enable motion sub loss."
|
|
)
|
|
parser.add_argument(
|
|
"--motion_sub_loss_ratio", type=float, default=0.25, help="The ratio of motion sub loss."
|
|
)
|
|
parser.add_argument(
|
|
"--train_sampling_steps",
|
|
type=int,
|
|
default=1000,
|
|
help="Run train_sampling_steps.",
|
|
)
|
|
parser.add_argument(
|
|
"--keep_all_node_same_token_length",
|
|
action="store_true",
|
|
help="Reference of the length token.",
|
|
)
|
|
parser.add_argument(
|
|
"--token_sample_size",
|
|
type=int,
|
|
default=512,
|
|
help="Sample size of the token.",
|
|
)
|
|
parser.add_argument(
|
|
"--video_sample_size",
|
|
type=int,
|
|
default=512,
|
|
help="Sample size of the video.",
|
|
)
|
|
parser.add_argument(
|
|
"--image_sample_size",
|
|
type=int,
|
|
default=512,
|
|
help="Sample size of the image.",
|
|
)
|
|
parser.add_argument(
|
|
"--fix_sample_size",
|
|
nargs=2, type=int, default=None,
|
|
help="Fix Sample size [height, width] when using bucket and collate_fn."
|
|
)
|
|
parser.add_argument(
|
|
"--video_sample_stride",
|
|
type=int,
|
|
default=4,
|
|
help="Sample stride of the video.",
|
|
)
|
|
parser.add_argument(
|
|
"--video_sample_n_frames",
|
|
type=int,
|
|
default=17,
|
|
help="Num frame of video.",
|
|
)
|
|
parser.add_argument(
|
|
"--video_repeat",
|
|
type=int,
|
|
default=0,
|
|
help="Num of repeat video.",
|
|
)
|
|
parser.add_argument(
|
|
"--config_path",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"The config of the model in training."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--transformer_path",
|
|
type=str,
|
|
default=None,
|
|
help=("If you want to load the weight from other transformers, input its path."),
|
|
)
|
|
parser.add_argument(
|
|
"--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: the generator uses SparkWanTransformer3DModel (sparse + linear self-attention), teacher and critic stay dense.",
|
|
)
|
|
parser.add_argument(
|
|
"--student_topk",
|
|
type=float,
|
|
default=0.10,
|
|
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(
|
|
"--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(
|
|
"--joint_backward_timesteps",
|
|
type=float,
|
|
nargs="+",
|
|
default=None,
|
|
help="Optional rf knots (descending, in [0,1]) that define the few-step denoising grid DIRECTLY, "
|
|
"replacing --denoising_step_indices_list. Wan2.1 is a single expert (no MoE boundary split), so "
|
|
"the whole list is used: the grid becomes [1.0] + [k for k in knots if 0<k<1] and the rf values "
|
|
"serve directly as Euler sigmas (sigma == rf == timestep/1000). Official Wan2.1 rCM knots: "
|
|
"0.933781 0.852895 0.608979.",
|
|
)
|
|
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, 0.0 when exhausted), 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 it reduces exactly to DMD. Requires --flow_euler_rollout; "
|
|
"mutually exclusive with --dfd / --decoupled_dmd.",
|
|
)
|
|
|
|
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(
|
|
"--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=5.0,
|
|
help="rf timestep shift for the PCM diversity grid (matches the Wan2.1 scheduler shift).",
|
|
)
|
|
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)
|
|
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]
|
|
|
|
# --- Few-step rf grid. Wan2.1 is a SINGLE expert (no MoE boundary), so --joint_backward_timesteps is used
|
|
# whole: grid = [1.0] + [k for k in knots if 0<k<1], the rf values serving directly as Euler sigmas. When it
|
|
# is not given the grid falls back to --denoising_step_indices_list (mapped to rf inside the train loop). ---
|
|
if args.joint_backward_timesteps:
|
|
denoising_rf_grid = [1.0] + [float(k) for k in args.joint_backward_timesteps if 0.0 < float(k) < 1.0]
|
|
if len(denoising_rf_grid) < 2:
|
|
raise ValueError(
|
|
f"--joint_backward_timesteps produced a degenerate rf grid {denoising_rf_grid}; "
|
|
f"provide at least one knot strictly inside (0,1)."
|
|
)
|
|
args.student_num_steps = len(denoising_rf_grid)
|
|
else:
|
|
denoising_rf_grid = None
|
|
args.student_num_steps = len(args.denoising_step_indices_list)
|
|
print(f"[wan2.1] denoising rf grid={denoising_rf_grid if denoising_rf_grid is not None else 'from --denoising_step_indices_list'} "
|
|
f"| student_num_steps={args.student_num_steps} | tdm={args.tdm}")
|
|
|
|
# Load scheduler, tokenizer and models.
|
|
noise_scheduler = FlowMatchEulerDiscreteScheduler(
|
|
**filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
|
|
)
|
|
|
|
# Get Tokenizer
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
|
|
)
|
|
|
|
def deepspeed_zero_init_disabled_context_manager():
|
|
"""
|
|
returns either a context list that includes one that will disable zero.Init or an empty context list
|
|
"""
|
|
deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None
|
|
if deepspeed_plugin is None:
|
|
return []
|
|
|
|
return [deepspeed_plugin.zero3_init_context_manager(enable=False)]
|
|
|
|
# Currently Accelerate doesn't know how to handle multiple models under Deepspeed ZeRO stage 3.
|
|
# For this to work properly all models must be run through `accelerate.prepare`. But accelerate
|
|
# will try to assign the same optimizer with the same weights to all models during
|
|
# `deepspeed.initialize`, which of course doesn't work.
|
|
#
|
|
# For now the following workaround will partially support Deepspeed ZeRO-3, by excluding the 2
|
|
# frozen models from being partitioned during `zero.Init` which gets called during
|
|
# `from_pretrained` So CLIPTextModel and AutoencoderKL will not enjoy the parameter sharding
|
|
# across multiple gpus and only UNet2DConditionModel will get ZeRO sharded.
|
|
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
|
|
# Get Text encoder
|
|
text_encoder = WanT5EncoderModel.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
|
|
additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
|
|
low_cpu_mem_usage=True,
|
|
torch_dtype=weight_dtype,
|
|
)
|
|
text_encoder = text_encoder.eval()
|
|
# Get Vae
|
|
vae = AutoencoderKLWan.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
|
|
additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
|
|
)
|
|
vae.eval()
|
|
# Get Clip Image Encoder
|
|
if args.train_mode != "normal":
|
|
clip_image_encoder = CLIPModel.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
|
|
)
|
|
clip_image_encoder = clip_image_encoder.eval()
|
|
else:
|
|
clip_image_encoder = None
|
|
|
|
# Get Transformer
|
|
# Sparse student (SparkDiffusion stage-1 output): block-sparse + linear self-attention
|
|
# (SparkWanTransformer3DModel); the dense teacher and the dense critic stay untouched,
|
|
# matching the reference pipeline where fake_score initializes from the teacher weights.
|
|
if args.sparse_student:
|
|
generator_transformer3d = SparkWanTransformer3DModel.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
|
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
|
low_cpu_mem_usage=True,
|
|
).to(weight_dtype)
|
|
# Materialize the zero-init sparse branch that the dense base cannot provide.
|
|
generator_transformer3d.set_student_topk(args.student_topk)
|
|
else:
|
|
generator_transformer3d = WanTransformer3DModel.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
|
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
|
low_cpu_mem_usage=True,
|
|
).to(weight_dtype)
|
|
real_score_transformer3d = WanTransformer3DModel.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
|
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
|
low_cpu_mem_usage=True,
|
|
).to(weight_dtype)
|
|
fake_score_transformer3d = WanTransformer3DModel.from_pretrained(
|
|
os.path.join(args.pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
|
|
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
|
|
low_cpu_mem_usage=True,
|
|
).to(weight_dtype)
|
|
|
|
# Freeze vae and text_encoder and set generator_transformer3d to trainable
|
|
vae.requires_grad_(False)
|
|
text_encoder.requires_grad_(False)
|
|
generator_transformer3d.requires_grad_(False)
|
|
real_score_transformer3d.requires_grad_(False)
|
|
fake_score_transformer3d.requires_grad_(False)
|
|
if args.train_mode != "normal":
|
|
clip_image_encoder.requires_grad_(False)
|
|
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 = WanTransformer3DModel.from_pretrained(
|
|
input_dir, subfolder="transformer",
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
model.register_to_config(**load_model.config)
|
|
|
|
model.load_state_dict(load_model.state_dict())
|
|
del load_model
|
|
|
|
pkl_path = os.path.join(input_dir, "sampler_pos_start.pkl")
|
|
if os.path.exists(pkl_path):
|
|
with open(pkl_path, 'rb') as file:
|
|
loaded_number, _ = pickle.load(file)
|
|
batch_sampler.sampler._pos_start = max(loaded_number - args.dataloader_num_workers * accelerator.num_processes * 2, 0)
|
|
print(f"Load pkl from {pkl_path}. Get loaded_number = {loaded_number}.")
|
|
|
|
accelerator.register_save_state_pre_hook(save_model_hook)
|
|
accelerator.register_load_state_pre_hook(load_model_hook)
|
|
accelerator_fake_score_transformer3d.register_save_state_pre_hook(save_model_hook)
|
|
accelerator_fake_score_transformer3d.register_load_state_pre_hook(load_model_hook)
|
|
|
|
if args.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
|
|
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
|
|
real_score_transformer3d = shard_fn(real_score_transformer3d)
|
|
if fsdp_stage != 0 or zero_stage != 0:
|
|
from functools import partial
|
|
|
|
from videox_fun.dist import set_multi_gpus_devices, shard_model
|
|
shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype)
|
|
text_encoder = shard_fn(text_encoder)
|
|
|
|
# Move text_encode and vae to gpu and cast to weight_dtype
|
|
vae.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
|
|
real_score_transformer3d.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
|
|
if not args.enable_text_encoder_in_dataloader:
|
|
text_encoder.to(accelerator.device if not args.low_vram else "cpu")
|
|
if args.train_mode != "normal":
|
|
clip_image_encoder.to(accelerator.device if not args.low_vram else "cpu", dtype=weight_dtype)
|
|
|
|
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
|
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
|
|
if overrode_max_train_steps:
|
|
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
|
|
# Afterwards we recalculate our number of training epochs
|
|
args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
|
|
|
|
# We need to initialize the trackers we use, and also store our configuration.
|
|
# The trackers initializes automatically on the main process.
|
|
if accelerator.is_main_process:
|
|
tracker_config = dict(vars(args))
|
|
keys_to_pop = [k for k, v in tracker_config.items() if isinstance(v, list)]
|
|
for k in keys_to_pop:
|
|
tracker_config.pop(k)
|
|
print(f"Removed tracker_config['{k}']")
|
|
accelerator.init_trackers(args.tracker_project_name, tracker_config)
|
|
|
|
# Function for unwrapping if model was compiled with `torch.compile`.
|
|
def unwrap_model(model):
|
|
model = accelerator.unwrap_model(model)
|
|
model = model._orig_mod if is_compiled_module(model) else model
|
|
return model
|
|
|
|
# Train!
|
|
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
|
|
|
|
logger.info("***** Running training *****")
|
|
logger.info(f" Num examples = {len(train_dataset)}")
|
|
logger.info(f" Num Epochs = {args.num_train_epochs}")
|
|
logger.info(f" Instantaneous batch size per device = {args.train_batch_size}")
|
|
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
|
|
logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
|
|
logger.info(f" Total optimization steps = {args.max_train_steps}")
|
|
global_step = 0
|
|
first_epoch = 0
|
|
|
|
# Potentially load in the weights and states from a previous save
|
|
if args.resume_from_checkpoint:
|
|
if args.resume_from_checkpoint != "latest":
|
|
path = os.path.basename(args.resume_from_checkpoint)
|
|
else:
|
|
# Get the most recent checkpoint
|
|
dirs = os.listdir(args.output_dir)
|
|
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
|
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
|
path = dirs[-1] if len(dirs) > 0 else None
|
|
|
|
if path is None:
|
|
accelerator.print(
|
|
f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
|
|
)
|
|
args.resume_from_checkpoint = None
|
|
initial_global_step = 0
|
|
else:
|
|
global_step = int(path.split("-")[1])
|
|
|
|
initial_global_step = global_step
|
|
|
|
pkl_path = os.path.join(os.path.join(args.output_dir, path), "sampler_pos_start.pkl")
|
|
if os.path.exists(pkl_path):
|
|
with open(pkl_path, 'rb') as file:
|
|
_, first_epoch = pickle.load(file)
|
|
else:
|
|
first_epoch = global_step // num_update_steps_per_epoch
|
|
print(f"Load pkl from {pkl_path}. Get first_epoch = {first_epoch}.")
|
|
|
|
accelerator.print(f"Resuming from checkpoint {path}")
|
|
fake_score_path = os.path.join(path, "fake_score")
|
|
accelerator.load_state(os.path.join(args.output_dir, path))
|
|
accelerator_fake_score_transformer3d.load_state(os.path.join(args.output_dir, fake_score_path))
|
|
else:
|
|
initial_global_step = 0
|
|
|
|
progress_bar = PauseAwareTqdm(
|
|
range(0, args.max_train_steps),
|
|
initial=initial_global_step,
|
|
desc="Steps",
|
|
# Only show the progress bar once on each machine.
|
|
disable=not accelerator.is_local_main_process,
|
|
)
|
|
|
|
if args.multi_stream and args.train_mode != "normal":
|
|
# create extra cuda streams to speedup inpaint vae computation
|
|
vae_stream_1 = torch.cuda.Stream()
|
|
vae_stream_2 = torch.cuda.Stream()
|
|
else:
|
|
vae_stream_1 = None
|
|
vae_stream_2 = None
|
|
|
|
idx_sampling = DiscreteSampling(args.train_sampling_steps, uniform_sampling=args.uniform_sampling)
|
|
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)
|
|
clip_image_encoder.to(accelerator.device)
|
|
real_score_transformer3d = real_score_transformer3d.to("cpu")
|
|
if not args.enable_text_encoder_in_dataloader:
|
|
text_encoder.to("cpu")
|
|
|
|
with torch.no_grad():
|
|
# This way is quicker when batch grows up
|
|
def _batch_encode_vae(pixel_values):
|
|
pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w")
|
|
bs = args.vae_mini_batch
|
|
new_pixel_values = []
|
|
for i in range(0, pixel_values.shape[0], bs):
|
|
pixel_values_bs = pixel_values[i : i + bs]
|
|
pixel_values_bs = vae.encode(pixel_values_bs)[0]
|
|
pixel_values_bs = pixel_values_bs.sample()
|
|
new_pixel_values.append(pixel_values_bs)
|
|
return torch.cat(new_pixel_values, dim = 0)
|
|
|
|
# Encode inpaint latents.
|
|
mask_latents = _batch_encode_vae(mask_pixel_values)
|
|
if vae_stream_2 is not None:
|
|
torch.cuda.current_stream().wait_stream(vae_stream_2)
|
|
|
|
mask = rearrange(mask, "b f c h w -> b c f h w")
|
|
mask = torch.concat(
|
|
[
|
|
torch.repeat_interleave(mask[:, :, 0:1], repeats=4, dim=2),
|
|
mask[:, :, 1:]
|
|
], dim=2
|
|
)
|
|
mask = mask.view(mask.shape[0], mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4])
|
|
mask = mask.transpose(1, 2)
|
|
mask = resize_mask(1 - mask, mask_latents)
|
|
|
|
inpaint_latents = torch.concat([mask, mask_latents], dim=1)
|
|
|
|
clip_context = []
|
|
for clip_pixel_value in clip_pixel_values:
|
|
clip_image = Image.fromarray(np.uint8(clip_pixel_value.float().cpu().numpy()))
|
|
clip_image = TF.to_tensor(clip_image).sub_(0.5).div_(0.5).to(clip_image_encoder.device, weight_dtype)
|
|
_clip_context = clip_image_encoder([clip_image[:, None, :, :]])
|
|
clip_context.append(_clip_context)
|
|
clip_context = torch.cat(clip_context)
|
|
|
|
target_shape = mask_latents.size()
|
|
else:
|
|
text = batch['text']
|
|
if args.fix_sample_size is not None:
|
|
local_sample_size = [int(x / 16) * 16 for x in args.fix_sample_size]
|
|
num_frames = args.video_sample_n_frames
|
|
else:
|
|
if args.random_hw_adapt and args.training_with_video_token_length:
|
|
# Get token length
|
|
target_token_length = args.video_sample_n_frames * args.token_sample_size * args.token_sample_size
|
|
length_to_frame_num = get_length_to_frame_num(target_token_length)
|
|
|
|
if rng is None:
|
|
local_length = np.random.choice(list(length_to_frame_num.keys()))
|
|
else:
|
|
local_length = rng.choice(list(length_to_frame_num.keys()))
|
|
num_frames = length_to_frame_num[local_length]
|
|
|
|
aspect_ratio_sample_size = {key : [x / 512 * local_length for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
|
|
if rng is None:
|
|
aspect_ratio_key = np.random.choice(list(aspect_ratio_sample_size.keys()))
|
|
else:
|
|
aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys()))
|
|
local_sample_size = aspect_ratio_sample_size[aspect_ratio_key]
|
|
else:
|
|
num_frames = args.video_sample_n_frames
|
|
|
|
aspect_ratio_sample_size = {key : [x / 512 * args.video_sample_size for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
|
|
if rng is None:
|
|
aspect_ratio_key = np.random.choice(list(aspect_ratio_sample_size.keys()))
|
|
else:
|
|
aspect_ratio_key = rng.choice(list(aspect_ratio_sample_size.keys()))
|
|
local_sample_size = aspect_ratio_sample_size[aspect_ratio_key]
|
|
local_sample_size = [int(x / 16) * 16 for x in local_sample_size]
|
|
|
|
target_shape = (
|
|
len(text),
|
|
vae.latent_channels,
|
|
int((num_frames - 1) // vae.temporal_compression_ratio + 1),
|
|
int(local_sample_size[0] // vae.spatial_compression_ratio),
|
|
int(local_sample_size[1] // vae.spatial_compression_ratio),
|
|
)
|
|
|
|
# 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")
|
|
if args.train_mode != "normal":
|
|
clip_image_encoder.to('cpu')
|
|
torch.cuda.empty_cache()
|
|
if not args.enable_text_encoder_in_dataloader:
|
|
text_encoder.to(accelerator.device)
|
|
|
|
if args.enable_text_encoder_in_dataloader:
|
|
prompt_embeds = batch['encoder_hidden_states'].to(device=accelerator.device)
|
|
neg_prompt_embeds = batch['neg_encoder_hidden_states'].to(device=accelerator.device)
|
|
else:
|
|
with torch.no_grad():
|
|
prompt_ids = tokenizer(
|
|
batch['text'],
|
|
padding="max_length",
|
|
max_length=args.tokenizer_max_length,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_tensors="pt"
|
|
)
|
|
text_input_ids = prompt_ids.input_ids
|
|
prompt_attention_mask = prompt_ids.attention_mask
|
|
|
|
seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
|
|
prompt_embeds = text_encoder(text_input_ids.to(accelerator.device), attention_mask=prompt_attention_mask.to(accelerator.device))[0]
|
|
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
|
|
|
|
neg_txt = [
|
|
"色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" for text in batch['text']
|
|
]
|
|
neg_prompt_ids = tokenizer(
|
|
neg_txt,
|
|
padding="max_length",
|
|
max_length=args.tokenizer_max_length,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_tensors="pt"
|
|
)
|
|
neg_text_input_ids = neg_prompt_ids.input_ids
|
|
neg_prompt_attention_mask = neg_prompt_ids.attention_mask
|
|
|
|
neg_seq_lens = neg_prompt_attention_mask.gt(0).sum(dim=1).long()
|
|
neg_prompt_embeds = text_encoder(neg_text_input_ids.to(accelerator.device), attention_mask=neg_prompt_attention_mask.to(accelerator.device))[0]
|
|
neg_prompt_embeds = [u[:v] for u, v in zip(neg_prompt_embeds, neg_seq_lens)]
|
|
|
|
if args.low_vram:
|
|
generator_transformer3d = generator_transformer3d.to(accelerator.device)
|
|
real_score_transformer3d = real_score_transformer3d.to(accelerator.device)
|
|
fake_score_transformer3d = fake_score_transformer3d.to(accelerator.device)
|
|
if not args.enable_text_encoder_in_dataloader:
|
|
text_encoder.to('cpu')
|
|
torch.cuda.empty_cache()
|
|
|
|
# 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. The rectified-flow sigma equals the rf time s (== timestep / 1000
|
|
# on this scheduler), and --joint_backward_timesteps knots are used DIRECTLY as sigmas, so the
|
|
# rollout drives the nets with t = s * T_SCALE and steps/inverts with s directly, bypassing the
|
|
# scheduler timestep-index map (get_sigmas / convert_flow_pred_to_x0) that would quantize knots. ---
|
|
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,
|
|
clip_fea=clip_context 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,
|
|
)
|
|
elif denoising_rf_grid is not None:
|
|
# rf-native grid from --joint_backward_timesteps (single expert: the whole list, no MoE
|
|
# boundary split). Used DIRECTLY as Euler sigmas (sigma == rf == timestep / 1000 here).
|
|
denoising_step_list = denoising_rf_grid
|
|
else:
|
|
# Create discrete denoising steps (per-step, with optional randomization), mapped to rf
|
|
# values (sigma == timestep / 1000) so the rollout below is uniformly rf-native.
|
|
if getattr(args, 'randomize_step_indices', False):
|
|
random_indices = randomize_denoising_step_indices(
|
|
args.denoising_step_indices_list,
|
|
args.train_sampling_steps,
|
|
torch_rng,
|
|
accelerator,
|
|
jitter_ratio=getattr(args, 'index_jitter_ratio', 0.30),
|
|
)
|
|
else:
|
|
random_indices = torch.tensor(args.denoising_step_indices_list)
|
|
|
|
denoising_step_list = (
|
|
noise_scheduler.timesteps[args.train_sampling_steps - random_indices] / T_SCALE
|
|
).tolist()
|
|
|
|
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)
|
|
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,
|
|
clip_fea=clip_context 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,
|
|
clip_fea=clip_context if args.train_mode != "normal" else None,
|
|
)
|
|
fake_score_main_cond = convert_flow_pred_to_x0(
|
|
scheduler=noise_scheduler,
|
|
flow_pred=fake_score_main_cond,
|
|
xt=generator_denoised_input,
|
|
timestep=generator_timestep
|
|
)
|
|
|
|
if args.fake_guidance_scale != 0.0:
|
|
fake_score_main_uncond = fake_score_transformer3d(
|
|
x=generator_denoised_input,
|
|
context=neg_prompt_embeds,
|
|
t=generator_timestep,
|
|
seq_len=seq_len,
|
|
y=inpaint_latents if args.train_mode != "normal" else None,
|
|
clip_fea=clip_context if args.train_mode != "normal" else None,
|
|
)
|
|
fake_score_main_uncond = convert_flow_pred_to_x0(
|
|
scheduler=noise_scheduler,
|
|
flow_pred=fake_score_main_uncond,
|
|
xt=generator_denoised_input,
|
|
timestep=generator_timestep
|
|
)
|
|
fake_score_main = fake_score_main_uncond + (
|
|
fake_score_main_cond - fake_score_main_uncond
|
|
) * args.fake_guidance_scale
|
|
else:
|
|
fake_score_main = fake_score_main_cond
|
|
|
|
# Compute real score (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,
|
|
clip_fea=clip_context if args.train_mode != "normal" else None,
|
|
)
|
|
real_score_main_cond = convert_flow_pred_to_x0(
|
|
scheduler=noise_scheduler,
|
|
flow_pred=real_score_main_cond,
|
|
xt=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,
|
|
clip_fea=clip_context 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,
|
|
clip_fea=clip_context 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,
|
|
clip_fea=clip_context if args.train_mode != "normal" else None,
|
|
)
|
|
real_score_main_uncond = convert_flow_pred_to_x0(
|
|
scheduler=noise_scheduler,
|
|
flow_pred=real_score_main_uncond,
|
|
xt=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)
|
|
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,
|
|
clip_fea=clip_context 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,
|
|
clip_fea=clip_context 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,
|
|
clip_image_encoder,
|
|
generator_transformer3d,
|
|
args,
|
|
config,
|
|
accelerator,
|
|
weight_dtype,
|
|
global_step,
|
|
)
|
|
|
|
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,
|
|
clip_image_encoder,
|
|
generator_transformer3d,
|
|
args,
|
|
config,
|
|
accelerator,
|
|
weight_dtype,
|
|
global_step,
|
|
)
|
|
|
|
# Close the bar before the end-of-run checkpoint: tqdm keeps redrawing a live bar whenever
|
|
# something else writes to the console. PauseAwareTqdm.close() rebases the closing line onto
|
|
# the smoothed rate, so the worker warm-up and the first dataloader fetch do not dilute it.
|
|
progress_bar.close()
|
|
|
|
# Create the pipeline using the trained modules and save it.
|
|
accelerator.wait_for_everyone()
|
|
if accelerator.is_main_process:
|
|
generator_transformer3d = unwrap_model(generator_transformer3d)
|
|
|
|
if args.use_deepspeed or args.use_fsdp or accelerator.is_main_process:
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
|
|
fake_score_save_path = os.path.join(save_path, "fake_score")
|
|
accelerator.save_state(save_path)
|
|
accelerator_fake_score_transformer3d.save_state(fake_score_save_path)
|
|
logger.info(f"Saved state to {save_path}")
|
|
|
|
accelerator.end_training()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|