wehub-resource-sync
1365 lines
61 KiB
Python
Executable File
1365 lines
61 KiB
Python
Executable File
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# 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
|
|
# limitations under the License.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import datetime
|
|
import gc
|
|
import hashlib
|
|
import os
|
|
import os.path as osp
|
|
import time
|
|
import warnings
|
|
from copy import deepcopy
|
|
from dataclasses import asdict
|
|
from pathlib import Path
|
|
|
|
warnings.filterwarnings("ignore") # ignore warning
|
|
|
|
import imageio
|
|
import numpy as np
|
|
import pyrallis
|
|
import torch
|
|
from accelerate import Accelerator, InitProcessGroupKwargs, skip_first_batches
|
|
from PIL import Image
|
|
from termcolor import colored
|
|
|
|
from diffusion import DPMS, FlowEuler, Scheduler
|
|
from diffusion.data.builder import build_dataloader, build_dataset
|
|
from diffusion.data.wids import DistributedRangedSampler
|
|
from diffusion.model.builder import build_model, get_tokenizer_and_text_encoder, get_vae, vae_decode, vae_encode
|
|
from diffusion.model.respace import compute_density_for_timestep_sampling
|
|
from diffusion.model.utils import get_weight_dtype
|
|
from diffusion.utils.checkpoint import load_checkpoint, save_checkpoint
|
|
from diffusion.utils.config import SanaVideoConfig, model_video_init_config
|
|
from diffusion.utils.data_sampler import AspectRatioBatchSampler, AspectRatioBatchSamplerVideo
|
|
from diffusion.utils.dist_utils import dist, flush, get_world_size
|
|
from diffusion.utils.git import save_git_snapshot
|
|
from diffusion.utils.logger import LogBuffer, get_root_logger
|
|
from diffusion.utils.lr_scheduler import build_lr_scheduler
|
|
from diffusion.utils.misc import DebugUnderflowOverflow, init_random_seed, set_random_seed
|
|
from diffusion.utils.optimizer import auto_scale_lr, build_optimizer
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
|
|
def set_fsdp_env():
|
|
# Basic FSDP settings
|
|
os.environ["ACCELERATE_USE_FSDP"] = "true"
|
|
|
|
# Auto wrapping policy
|
|
os.environ["FSDP_AUTO_WRAP_POLICY"] = "TRANSFORMER_BASED_WRAP"
|
|
os.environ["FSDP_TRANSFORMER_CLS_TO_WRAP"] = "SanaVideoMSBlock" # Your transformer block name
|
|
|
|
# Performance optimization settings
|
|
os.environ["FSDP_BACKWARD_PREFETCH"] = "BACKWARD_PRE"
|
|
os.environ["FSDP_FORWARD_PREFETCH"] = "false"
|
|
|
|
# State dict settings
|
|
os.environ["FSDP_STATE_DICT_TYPE"] = "FULL_STATE_DICT"
|
|
os.environ["FSDP_SYNC_MODULE_STATES"] = "true"
|
|
os.environ["FSDP_USE_ORIG_PARAMS"] = "true"
|
|
|
|
# Sharding strategy
|
|
os.environ["FSDP_SHARDING_STRATEGY"] = "HYBRID_SHARD" # FULL_SHARD
|
|
|
|
# Memory optimization settings (optional)
|
|
os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = "false"
|
|
os.environ["FSDP_OFFLOAD_PARAMS"] = "false"
|
|
|
|
# Precision settings
|
|
os.environ["FSDP_REDUCE_SCATTER_PRECISION"] = "fp32"
|
|
os.environ["FSDP_ALL_GATHER_PRECISION"] = "fp32"
|
|
os.environ["FSDP_OPTIMIZER_STATE_PRECISION"] = "fp32"
|
|
|
|
|
|
def ema_update(model_dest, model_src, rate):
|
|
param_dict_src = dict(model_src.named_parameters())
|
|
for p_name, p_dest in model_dest.named_parameters():
|
|
p_src = param_dict_src[p_name]
|
|
assert p_src is not p_dest
|
|
p_dest.data.mul_(rate).add_((1 - rate) * p_src.data)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def log_validation(accelerator, config, model, logger, step, device, vae=None, init_noise=None):
|
|
torch.cuda.empty_cache()
|
|
vis_sampler = config.scheduler.vis_sampler
|
|
model = accelerator.unwrap_model(model).eval()
|
|
hw = torch.tensor([[video_height, video_width]], dtype=torch.float, device=device).repeat(1, 1)
|
|
ar = torch.tensor([[1.0]], device=device).repeat(1, 1)
|
|
null_y = torch.load(null_embed_path, map_location="cpu")
|
|
null_y = null_y["uncond_prompt_embeds"].to(device)
|
|
cfg_scale = 4.5
|
|
|
|
# Create sampling noise:
|
|
logger.info("Running validation... ")
|
|
video_logs = []
|
|
|
|
def run_sampling(init_z=None, label_suffix="", vae=None, sampler="dpm-solver"):
|
|
latents = []
|
|
current_video_logs = []
|
|
for prompt in validation_prompts:
|
|
z = (
|
|
torch.randn(1, config.vae.vae_latent_dim, latent_temp, latent_height, latent_width, device=device)
|
|
if init_z is None
|
|
else init_z
|
|
)
|
|
logger.info(f"Loading embedding for prompt from: {config.train.valid_prompt_embed_root}")
|
|
embed = torch.load(
|
|
osp.join(config.train.valid_prompt_embed_root, f"{prompt[:50]}_{valid_prompt_embed_suffix}"),
|
|
map_location="cpu",
|
|
)
|
|
caption_embs, emb_masks = embed["caption_embeds"].to(device), embed["emb_mask"].to(device)
|
|
model_kwargs = dict(data_info={"img_hw": hw, "aspect_ratio": ar}, mask=emb_masks)
|
|
|
|
with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
|
|
if sampler == "flow_dpm-solver":
|
|
dpm_solver = DPMS(
|
|
model.forward_with_dpmsolver,
|
|
condition=caption_embs,
|
|
uncondition=null_y,
|
|
cfg_scale=cfg_scale,
|
|
model_type="flow",
|
|
model_kwargs=model_kwargs,
|
|
schedule="FLOW",
|
|
)
|
|
denoised = dpm_solver.sample(
|
|
z,
|
|
steps=50,
|
|
order=2,
|
|
skip_type="time_uniform_flow",
|
|
method="multistep",
|
|
flow_shift=(
|
|
config.scheduler.inference_flow_shift
|
|
if config.scheduler.inference_flow_shift is not None
|
|
else config.scheduler.flow_shift
|
|
),
|
|
)
|
|
else:
|
|
raise ValueError(f"{sampler} not implemented")
|
|
|
|
latents.append(denoised)
|
|
torch.cuda.empty_cache()
|
|
if vae is None:
|
|
vae = get_vae(
|
|
config.vae.vae_type, config.vae.vae_pretrained, accelerator.device, dtype=vae_dtype, config=config.vae
|
|
)
|
|
for prompt, latent in zip(validation_prompts, latents):
|
|
latent = latent.to(vae_dtype)
|
|
samples = vae_decode(config.vae.vae_type, vae, latent)
|
|
video = (
|
|
torch.clamp(127.5 * samples[0] + 127.5, 0, 255).permute(1, 0, 2, 3).to("cpu", dtype=torch.uint8).numpy()
|
|
) # C,T,H,W -> T,C,H,W
|
|
current_video_logs.append({"validation_prompt": prompt + label_suffix, "videos": video})
|
|
|
|
return current_video_logs
|
|
|
|
# First run with original noise
|
|
video_logs += run_sampling(init_z=None, label_suffix="", vae=vae, sampler=vis_sampler)
|
|
|
|
# Second run with init_noise if provided
|
|
if init_noise is not None:
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
init_noise = torch.clone(init_noise).to(device)
|
|
video_logs += run_sampling(init_z=init_noise, label_suffix=" w/ init noise", vae=vae, sampler=vis_sampler)
|
|
|
|
for tracker in accelerator.trackers:
|
|
if tracker.name == "wandb":
|
|
import wandb
|
|
|
|
wandb_items = []
|
|
for log_item in video_logs:
|
|
wandb_items.append(
|
|
wandb.Video(log_item["videos"], caption=log_item["validation_prompt"], fps=16, format="mp4")
|
|
)
|
|
tracker.log({"validation": wandb_items})
|
|
else:
|
|
logger.warn(f"Video logging not implemented for {tracker.name}")
|
|
|
|
def concatenate_videos(video_data, videos_per_row=3, video_format="mp4"):
|
|
videos = [torch.from_numpy(log["videos"]).to(torch.uint8) for log in video_data] # T,C,H,W
|
|
|
|
num_videos = len(videos)
|
|
num_rows = (num_videos + videos_per_row - 1) // videos_per_row
|
|
num_frames, num_channels, height, width = videos[0].shape
|
|
total_width = width * min(videos_per_row, num_videos)
|
|
total_height = height * num_rows
|
|
|
|
grid_video = torch.zeros((num_frames, num_channels, total_height, total_width), dtype=videos[0].dtype)
|
|
|
|
for i, video in enumerate(videos):
|
|
|
|
row = i // videos_per_row
|
|
col = i % videos_per_row
|
|
|
|
y_offset = row * height
|
|
x_offset = col * width
|
|
|
|
h, w = video.shape[2:]
|
|
|
|
grid_video[:, :, y_offset : y_offset + h, x_offset : x_offset + w] = video
|
|
|
|
return grid_video
|
|
|
|
if config.train.local_save_vis:
|
|
file_format = "mp4"
|
|
local_vis_save_path = osp.join(config.work_dir, "log_vis")
|
|
os.umask(0o000)
|
|
os.makedirs(local_vis_save_path, exist_ok=True)
|
|
concatenated_video = concatenate_videos(video_logs, videos_per_row=5, video_format=file_format)
|
|
save_path = (
|
|
osp.join(local_vis_save_path, f"vis_{step}.{file_format}")
|
|
if init_noise is None
|
|
else osp.join(local_vis_save_path, f"vis_{step}_w_init.{file_format}")
|
|
)
|
|
save_video = concatenated_video.permute(0, 2, 3, 1)
|
|
writer = imageio.v2.get_writer(save_path, fps=16, format="FFMPEG", codec="libx264", quality=8)
|
|
for frame in save_video.numpy():
|
|
writer.append_data(frame)
|
|
writer.close()
|
|
|
|
model.train()
|
|
del vae
|
|
flush()
|
|
return video_logs
|
|
|
|
|
|
def train(
|
|
config,
|
|
args,
|
|
accelerator,
|
|
model,
|
|
model_ema,
|
|
optimizer,
|
|
lr_scheduler,
|
|
train_dataloader,
|
|
train_dataloader_image,
|
|
train_diffusion,
|
|
logger,
|
|
):
|
|
if getattr(config.train, "debug_nan", False):
|
|
DebugUnderflowOverflow(model, max_frames_to_save=100)
|
|
logger.info("NaN debugger registered. Start to detect overflow during training.")
|
|
log_buffer = LogBuffer()
|
|
|
|
global_step = start_step
|
|
video_step = start_video_step # Track video steps separately
|
|
image_step = start_image_step # Track image steps separately
|
|
|
|
skip_step = max(config.train.skip_step, video_step) % train_dataloader_len
|
|
skip_step = skip_step if skip_step < (train_dataloader_len - 20) else 0
|
|
skip_step_image = max(config.train.skip_step, image_step) % train_dataloader_image_len
|
|
skip_step_image = skip_step_image if skip_step_image < (train_dataloader_image_len - 20) else 0
|
|
loss_nan_timer = 0
|
|
model_instance.to(accelerator.device)
|
|
|
|
# Now you train the model
|
|
for epoch in range(start_epoch + 1, config.train.num_epochs + 1):
|
|
time_start, last_tic = time.time(), time.time()
|
|
sampler = (
|
|
train_dataloader.batch_sampler.sampler
|
|
if (num_replicas > 1 or config.model.multi_scale)
|
|
else train_dataloader.sampler
|
|
)
|
|
if joint_training_interval > 0:
|
|
image_sampler = (
|
|
train_dataloader_image.batch_sampler.sampler
|
|
if (num_replicas > 1 or config.model.multi_scale)
|
|
else train_dataloader_image.sampler
|
|
)
|
|
else:
|
|
image_sampler = None
|
|
if train_dataloader.dataset.shuffle_dataset:
|
|
logger.info(f"Shuffled dataset, no skip step")
|
|
else:
|
|
set_start_value = max((skip_step - 1) * config.train.train_batch_size, 0)
|
|
os.environ[f"CURRENT_VIDEO_STEP_START_RANK_{rank}"] = str(set_start_value)
|
|
sampler.set_epoch(epoch)
|
|
sampler.set_start(set_start_value)
|
|
if image_sampler is not None:
|
|
set_image_start_value = max((skip_step_image - 1) * config.train.train_batch_size_image, 0)
|
|
os.environ[f"CURRENT_IMAGE_STEP_START_RANK_{rank}"] = str(set_image_start_value)
|
|
image_sampler.set_epoch(epoch)
|
|
image_sampler.set_start(set_image_start_value)
|
|
|
|
if skip_step > 1 and accelerator.is_main_process:
|
|
logger.info(f"Skipped video training Steps: {skip_step}")
|
|
if image_sampler is not None:
|
|
logger.info(f"Skipped image training Steps: {skip_step_image}")
|
|
skip_step = 1
|
|
data_time_start = time.time()
|
|
data_time_all = 0
|
|
lm_time_all = 0
|
|
vae_time_all = 0
|
|
model_time_all = 0
|
|
|
|
# Create dataloader iterators for joint training
|
|
video_dataloader_iter = iter(train_dataloader)
|
|
image_dataloader_iter = iter(train_dataloader_image)
|
|
|
|
# Use range instead of enumerating train_dataloader
|
|
for step in range(train_dataloader_len):
|
|
|
|
# Determine if this is an image training step
|
|
is_image_step = (
|
|
(joint_training_interval > 0) and (global_step % joint_training_interval == 0) and (global_step > 0)
|
|
)
|
|
|
|
if is_image_step:
|
|
# Get image batch for joint training
|
|
try:
|
|
batch = next(image_dataloader_iter)
|
|
except StopIteration:
|
|
# Reset image dataloader iterator if exhausted
|
|
image_dataloader_iter = iter(train_dataloader_image)
|
|
batch = next(image_dataloader_iter)
|
|
is_video_data = False
|
|
image_step += 1 # Increment image step counter
|
|
else:
|
|
# Get video batch
|
|
try:
|
|
batch = next(video_dataloader_iter)
|
|
except StopIteration:
|
|
# Reset video dataloader iterator if exhausted
|
|
logger.info(f"Reset video dataloader iterator")
|
|
sampler.set_start(0)
|
|
video_dataloader_iter = iter(train_dataloader)
|
|
batch = next(video_dataloader_iter)
|
|
is_video_data = True
|
|
video_step += 1 # Increment video step counter
|
|
|
|
# if epoch > config.train.num_epochs:
|
|
# logger.info(f"Stopping training at epoch {epoch}, step {global_step} due to num_epochs limit.")
|
|
# return
|
|
|
|
# image, json_info, key = batch
|
|
accelerator.wait_for_everyone()
|
|
data_time_all += time.time() - data_time_start
|
|
vae_time_start = time.time()
|
|
data_info = batch[3]
|
|
|
|
with torch.no_grad():
|
|
if is_video_data:
|
|
if load_vae_feat:
|
|
z = batch[0].to(accelerator.device)
|
|
# Video data processing (original code)
|
|
else:
|
|
try:
|
|
z = vae_encode(
|
|
config.vae.vae_type,
|
|
vae,
|
|
batch[0].permute(0, 2, 1, 3, 4).to(vae_dtype),
|
|
device=accelerator.device,
|
|
cache_key=data_info["cache_key"],
|
|
if_cache=config.vae.if_cache,
|
|
data_info=data_info,
|
|
) # B,F,C,H,W -> B,C,F,H,W
|
|
except Exception as e:
|
|
print(f"Error in vae_encode: {e}")
|
|
print(f"Data info: {data_info}")
|
|
|
|
else:
|
|
# Image data processing (similar to stage1)
|
|
if batch[0].dim() == 4:
|
|
batch[0] = batch[0][:, :, None] # B,C,H,W -> B,C,1,H,W
|
|
z = vae_encode(
|
|
config.vae.vae_type,
|
|
vae,
|
|
batch[0].to(vae_dtype),
|
|
device=accelerator.device,
|
|
)
|
|
|
|
accelerator.wait_for_everyone()
|
|
vae_time_all += time.time() - vae_time_start
|
|
|
|
clean_images = z
|
|
|
|
lm_time_start = time.time()
|
|
if load_text_feat:
|
|
y = batch[1] # bs, 1, N, C
|
|
y_mask = batch[2] # bs, 1, 1, N
|
|
else:
|
|
if "T5" in config.text_encoder.text_encoder_name:
|
|
with torch.no_grad():
|
|
txt_tokens = tokenizer(
|
|
batch[1], max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
|
|
).to(accelerator.device)
|
|
y = text_encoder(txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask)[0][:, None]
|
|
y_mask = txt_tokens.attention_mask[:, None, None]
|
|
elif "gemma" in config.text_encoder.text_encoder_name:
|
|
with torch.no_grad():
|
|
if not config.text_encoder.chi_prompt:
|
|
max_length_all = config.text_encoder.model_max_length
|
|
prompt = batch[1]
|
|
else:
|
|
chi_prompt = "\n".join(config.text_encoder.chi_prompt)
|
|
prompt = [chi_prompt + i for i in batch[1]]
|
|
num_sys_prompt_tokens = len(tokenizer.encode(chi_prompt))
|
|
max_length_all = (
|
|
num_sys_prompt_tokens + config.text_encoder.model_max_length - 2
|
|
) # magic number 2: [bos], [_]
|
|
txt_tokens = tokenizer(
|
|
prompt,
|
|
padding="max_length",
|
|
max_length=max_length_all,
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
).to(accelerator.device)
|
|
select_index = [0] + list(
|
|
range(-config.text_encoder.model_max_length + 1, 0)
|
|
) # first bos and end N-1
|
|
y = text_encoder(txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask)[0][:, None][
|
|
:, :, select_index
|
|
]
|
|
y_mask = txt_tokens.attention_mask[:, None, None][:, :, :, select_index]
|
|
elif "Qwen" in config.text_encoder.text_encoder_name:
|
|
with torch.no_grad():
|
|
y, y_mask = text_encoder.get_prompt_embeds(prompt)
|
|
y_mask = y_mask[:, None, None]
|
|
else:
|
|
print("error")
|
|
exit()
|
|
|
|
# Sample a random timestep for each image
|
|
bs = clean_images.shape[0]
|
|
timesteps = torch.randint(
|
|
0, config.scheduler.train_sampling_steps, (bs,), device=clean_images.device
|
|
).long()
|
|
if config.scheduler.weighting_scheme in ["logit_normal", "mode"]:
|
|
# adapting from diffusers.training_utils
|
|
u = compute_density_for_timestep_sampling(
|
|
weighting_scheme=config.scheduler.weighting_scheme,
|
|
batch_size=bs,
|
|
logit_mean=config.scheduler.logit_mean,
|
|
logit_std=config.scheduler.logit_std,
|
|
mode_scale=config.scheduler.mode_scale,
|
|
)
|
|
timesteps = (u * config.scheduler.train_sampling_steps).long().to(clean_images.device)
|
|
grad_norm = None
|
|
accelerator.wait_for_everyone()
|
|
lm_time_all += time.time() - lm_time_start
|
|
model_time_start = time.time()
|
|
with accelerator.accumulate(model):
|
|
# Predict the noise residual
|
|
optimizer.zero_grad()
|
|
loss_term = train_diffusion.training_losses(
|
|
model, clean_images, timesteps, model_kwargs=dict(y=y, mask=y_mask, data_info=data_info)
|
|
)
|
|
loss = loss_term["loss"].mean()
|
|
|
|
# Temporal coherence loss (frame-difference consistency)
|
|
if (
|
|
is_video_data
|
|
and hasattr(config.train, "temporal_coherence_loss")
|
|
and config.train.temporal_coherence_loss
|
|
and getattr(config.train, "temporal_coherence_weight", 0.0) > 0.0
|
|
):
|
|
try:
|
|
model_output = loss_term["output"]
|
|
# noise = loss_term["noise"]
|
|
x_t = loss_term["x_t"]
|
|
|
|
pred_x0 = x_t - timesteps.view(-1, 1, 1, 1, 1) / 1000.0 * model_output
|
|
|
|
if pred_x0.dim() == 5:
|
|
pred_diff = pred_x0[:, :, 1:] - pred_x0[:, :, :-1]
|
|
gt_diff = clean_images[:, :, 1:] - clean_images[:, :, :-1]
|
|
tc_loss = (pred_diff - gt_diff).pow(2).mean() * config.train.temporal_coherence_weight
|
|
loss = loss + tc_loss
|
|
loss_term["tc"] = tc_loss.detach()
|
|
except Exception:
|
|
pass
|
|
|
|
accelerator.backward(loss)
|
|
|
|
if accelerator.sync_gradients:
|
|
grad_norm = accelerator.clip_grad_norm_(model.parameters(), config.train.gradient_clip)
|
|
if not config.train.use_fsdp and config.train.ema_update and model_ema is not None:
|
|
ema_update(model_ema, model, config.train.ema_rate)
|
|
|
|
optimizer.step()
|
|
lr_scheduler.step()
|
|
accelerator.wait_for_everyone()
|
|
model_time_all += time.time() - model_time_start
|
|
|
|
if torch.any(torch.isnan(loss)):
|
|
loss_nan_timer += 1
|
|
lr = lr_scheduler.get_last_lr()[0]
|
|
logs = {args.loss_report_name: accelerator.gather(loss).mean().item()}
|
|
if grad_norm is not None:
|
|
logs.update(grad_norm=accelerator.gather(grad_norm).mean().item())
|
|
if "tc" in loss_term:
|
|
logs.update(tc=accelerator.gather(loss_term["tc"]).mean().item())
|
|
log_buffer.update(logs)
|
|
if (global_step + 1) % config.train.log_interval == 0 or (step + 1) == 1:
|
|
accelerator.wait_for_everyone()
|
|
if args.debug:
|
|
print(f"Rank {rank}: current_batch_id: {batch[4]}")
|
|
|
|
t = (time.time() - last_tic) / config.train.log_interval
|
|
t_d = data_time_all / config.train.log_interval
|
|
t_m = model_time_all / config.train.log_interval
|
|
t_lm = lm_time_all / config.train.log_interval
|
|
t_vae = vae_time_all / config.train.log_interval
|
|
avg_time = (time.time() - time_start) / (step + 1)
|
|
eta = str(datetime.timedelta(seconds=int(avg_time * (total_steps - global_step - 1))))
|
|
eta_epoch = str(
|
|
datetime.timedelta(
|
|
seconds=int(
|
|
avg_time
|
|
* (train_dataloader_len - sampler.step_start // config.train.train_batch_size - step - 1)
|
|
)
|
|
)
|
|
)
|
|
log_buffer.average()
|
|
|
|
if joint_training_interval > 0:
|
|
current_step = (
|
|
global_step
|
|
- sampler.step_start // config.train.train_batch_size
|
|
- image_sampler.step_start // config.train.train_batch_size_image
|
|
) % train_dataloader_len
|
|
else:
|
|
current_step = (
|
|
global_step - sampler.step_start // config.train.train_batch_size
|
|
) % train_dataloader_len
|
|
|
|
current_step = train_dataloader_len if current_step == 0 else current_step
|
|
|
|
data_type = "Image" if not is_video_data else "Video"
|
|
id_info = (
|
|
f"{batch[4][-1]}:{'/'.join(data_info['zip_file'][-1].split('/')[-2:])}"
|
|
if "zip_file" in data_info
|
|
else f"{batch[4][-1]}"
|
|
)
|
|
info = (
|
|
f"Epoch: {epoch} | Global Step: {global_step + 1} / {train_dataloader_len}, "
|
|
f"Video Step: {video_step} | Image Step: {image_step} | id: {id_info}, "
|
|
f"total_eta: {eta}, epoch_eta:{eta_epoch}, time: all:{t:.3f}, model:{t_m:.3f}, data:{t_d:.3f}, "
|
|
f"lm:{t_lm:.3f}, vae:{t_vae:.3f}, lr:{lr:.3e}, DataType: {data_type}, Cap: {batch[5][0]}, "
|
|
)
|
|
info += (
|
|
f"s:({model.module.f}, {model.module.h}, {model.module.w}), "
|
|
if hasattr(model, "module")
|
|
else f"s:({model.f}, {model.h}, {model.w}), "
|
|
)
|
|
|
|
info += ", ".join([f"{k}:{v:.4f}" for k, v in log_buffer.output.items()])
|
|
last_tic = time.time()
|
|
log_buffer.clear()
|
|
data_time_all = 0
|
|
model_time_all = 0
|
|
lm_time_all = 0
|
|
vae_time_all = 0
|
|
if accelerator.is_main_process:
|
|
logger.info(info)
|
|
|
|
logs.update(lr=lr)
|
|
if accelerator.is_main_process:
|
|
accelerator.log(logs, step=global_step)
|
|
|
|
global_step += 1
|
|
|
|
if loss_nan_timer > 20:
|
|
raise ValueError("Loss is NaN too much times. Break here.")
|
|
if (
|
|
global_step % config.train.save_model_steps == 0
|
|
or (time.time() - training_start_time) / 3600 > config.train.early_stop_hours
|
|
):
|
|
torch.cuda.synchronize()
|
|
accelerator.wait_for_everyone()
|
|
|
|
# Choose different saving methods based on whether FSDP is used
|
|
if config.train.use_fsdp:
|
|
# FSDP mode
|
|
os.umask(0o000)
|
|
saved_info = {
|
|
"video_step": video_step,
|
|
"image_step": image_step,
|
|
}
|
|
ckpt_saved_path = save_checkpoint(
|
|
work_dir=osp.join(config.work_dir, "checkpoints"),
|
|
epoch=epoch,
|
|
model=model,
|
|
accelerator=accelerator,
|
|
optimizer=optimizer,
|
|
lr_scheduler=lr_scheduler,
|
|
step=global_step,
|
|
saved_info=saved_info,
|
|
add_symlink=True,
|
|
)
|
|
else:
|
|
# DDP mode
|
|
if accelerator.is_main_process:
|
|
os.umask(0o000)
|
|
saved_info = {
|
|
"video_step": video_step,
|
|
"image_step": image_step,
|
|
}
|
|
ckpt_saved_path = save_checkpoint(
|
|
work_dir=osp.join(config.work_dir, "checkpoints"),
|
|
epoch=epoch,
|
|
model=accelerator.unwrap_model(model),
|
|
model_ema=accelerator.unwrap_model(model_ema) if model_ema is not None else None,
|
|
optimizer=optimizer,
|
|
lr_scheduler=lr_scheduler,
|
|
step=global_step,
|
|
saved_info=saved_info,
|
|
generator=generator,
|
|
add_symlink=True,
|
|
)
|
|
|
|
if accelerator.is_main_process:
|
|
if config.train.online_metric and global_step % config.train.eval_metric_step == 0 and step > 1:
|
|
online_metric_monitor_dir = osp.join(config.work_dir, config.train.online_metric_dir)
|
|
os.makedirs(online_metric_monitor_dir, exist_ok=True)
|
|
with open(f"{online_metric_monitor_dir}/{ckpt_saved_path.split('/')[-1]}.txt", "w") as f:
|
|
f.write(osp.join(config.work_dir, "config.py") + "\n")
|
|
f.write(ckpt_saved_path)
|
|
|
|
if (time.time() - training_start_time) / 3600 > config.train.early_stop_hours:
|
|
logger.info(f"Stopping training at epoch {epoch}, step {global_step} due to time limit.")
|
|
return
|
|
|
|
if config.train.visualize and (global_step % config.train.eval_sampling_steps == 0 or (step + 1) == 1):
|
|
if config.train.use_fsdp:
|
|
merged_state_dict = accelerator.get_state_dict(model)
|
|
|
|
accelerator.wait_for_everyone()
|
|
if accelerator.is_main_process:
|
|
if config.train.use_fsdp:
|
|
model_instance.load_state_dict(merged_state_dict)
|
|
|
|
if validation_noise is not None:
|
|
log_validation(
|
|
accelerator=accelerator,
|
|
config=config,
|
|
model=model_instance,
|
|
logger=logger,
|
|
step=global_step,
|
|
device=accelerator.device,
|
|
vae=vae,
|
|
init_noise=validation_noise,
|
|
)
|
|
else:
|
|
log_validation(
|
|
accelerator=accelerator,
|
|
config=config,
|
|
model=model_instance,
|
|
logger=logger,
|
|
step=global_step,
|
|
device=accelerator.device,
|
|
vae=vae,
|
|
)
|
|
|
|
# avoid dead-lock of multiscale data batch sampler
|
|
if (
|
|
config.model.multi_scale
|
|
and (train_dataloader_len - sampler.step_start // config.train.train_batch_size - step) < 30
|
|
):
|
|
global_step = (
|
|
(global_step + train_dataloader_len - 1) // train_dataloader_len
|
|
) * train_dataloader_len + 1
|
|
logger.info("Early stop current iteration")
|
|
skip_first_batches(train_dataloader, True)
|
|
break
|
|
|
|
data_time_start = time.time()
|
|
|
|
if epoch % config.train.save_model_epochs == 0 or epoch == config.train.num_epochs and not config.debug:
|
|
accelerator.wait_for_everyone()
|
|
torch.cuda.synchronize()
|
|
|
|
# Choose different saving methods based on whether FSDP is used
|
|
if config.train.use_fsdp:
|
|
# FSDP mode
|
|
os.umask(0o000)
|
|
saved_info = {
|
|
"video_step": video_step,
|
|
"image_step": image_step,
|
|
}
|
|
ckpt_saved_path = save_checkpoint(
|
|
work_dir=osp.join(config.work_dir, "checkpoints"),
|
|
epoch=epoch,
|
|
model=model,
|
|
accelerator=accelerator,
|
|
optimizer=optimizer,
|
|
lr_scheduler=lr_scheduler,
|
|
step=global_step,
|
|
saved_info=saved_info,
|
|
add_symlink=True,
|
|
)
|
|
else:
|
|
# DDP mode
|
|
if accelerator.is_main_process:
|
|
os.umask(0o000)
|
|
saved_info = {
|
|
"video_step": video_step,
|
|
"image_step": image_step,
|
|
}
|
|
ckpt_saved_path = save_checkpoint(
|
|
osp.join(config.work_dir, "checkpoints"),
|
|
epoch=epoch,
|
|
step=global_step,
|
|
saved_info=saved_info,
|
|
model=accelerator.unwrap_model(model),
|
|
model_ema=accelerator.unwrap_model(model_ema) if model_ema is not None else None,
|
|
optimizer=optimizer,
|
|
lr_scheduler=lr_scheduler,
|
|
generator=generator,
|
|
add_symlink=True,
|
|
)
|
|
|
|
if accelerator.is_main_process:
|
|
online_metric_monitor_dir = osp.join(config.work_dir, config.train.online_metric_dir)
|
|
os.makedirs(online_metric_monitor_dir, exist_ok=True)
|
|
with open(f"{online_metric_monitor_dir}/{ckpt_saved_path.split('/')[-1]}.txt", "w") as f:
|
|
f.write(osp.join(config.work_dir, "config.py") + "\n")
|
|
f.write(ckpt_saved_path)
|
|
|
|
if epoch > config.train.num_epochs:
|
|
logger.info(f"Stopping training at epoch {epoch}, step {global_step} due to num_epochs limit.")
|
|
return
|
|
|
|
|
|
@pyrallis.wrap()
|
|
def main(cfg: SanaVideoConfig) -> None:
|
|
global train_dataloader_len, start_epoch, start_step, start_video_step, start_image_step, vae, generator, num_replicas, rank, training_start_time
|
|
global load_vae_feat, load_text_feat, validation_noise, text_encoder, tokenizer
|
|
global max_length, validation_prompts, latent_size, valid_prompt_embed_suffix, null_embed_path
|
|
global image_size, cache_file, total_steps, vae_dtype, model_instance
|
|
global video_width, video_height, num_frames, latent_temp, latent_height, latent_width
|
|
global image_encoder, image_processor, joint_training_interval, train_dataloader_image_len
|
|
|
|
config = cfg
|
|
args = cfg
|
|
|
|
# 1.Initialize training mode
|
|
if config.train.use_fsdp:
|
|
set_fsdp_env()
|
|
init_train = "FSDP"
|
|
else:
|
|
init_train = "DDP"
|
|
|
|
training_start_time = time.time()
|
|
load_from = True
|
|
|
|
if args.resume_from or config.model.resume_from:
|
|
load_from = False
|
|
config.model.resume_from = dict(
|
|
checkpoint=args.resume_from or config.model.resume_from,
|
|
load_ema=False,
|
|
resume_optimizer=True,
|
|
resume_lr_scheduler=config.train.resume_lr_scheduler,
|
|
)
|
|
|
|
if args.debug:
|
|
config.train.train_batch_size = min(64, config.train.train_batch_size)
|
|
if config.train.use_fsdp:
|
|
os.environ["FSDP_SHARDING_STRATEGY"] = "FULL_SHARD"
|
|
config.data.data_dir = {"video_toy_data": "data/video_toy_data"}
|
|
config.train.validation_prompts = [
|
|
"the opening scene begins with a dynamic view of a bustling cityscape captured in vibrant detail. towering skyscrapers dominate the skyline, while the streets below are alive with motion. people from diverse cultures fill the sidewalks, engaging in daily activities, their vibrant attire adding splashes of color to the scene. vehicles, including cars and buses, weave through the busy roads in a synchronized rhythm. bright billboards in various languages flash advertisements, reflecting the multicultural essence of the city. thecamera smoothly pans upward from the busy streets to focus on a sleek, modern office building. its reflective glass facade shimmers in the sunlight, hinting at its importance as a central location in the story. the atmosphere is energetic and cosmopolitan, setting the stage for an international narrative."
|
|
]
|
|
|
|
os.umask(0o000)
|
|
os.makedirs(config.work_dir, exist_ok=True)
|
|
|
|
init_handler = InitProcessGroupKwargs()
|
|
init_handler.timeout = datetime.timedelta(seconds=5400) # change timeout to avoid a strange NCCL bug
|
|
|
|
# Initialize accelerator and tensorboard logging
|
|
accelerator = Accelerator(
|
|
mixed_precision=config.model.mixed_precision,
|
|
gradient_accumulation_steps=config.train.gradient_accumulation_steps,
|
|
log_with=args.report_to,
|
|
project_dir=osp.join(config.work_dir, "logs"),
|
|
kwargs_handlers=[init_handler],
|
|
)
|
|
|
|
log_name = "train_log.log"
|
|
logger = get_root_logger(osp.join(config.work_dir, log_name))
|
|
logger.info(accelerator.state)
|
|
|
|
# save git snapshot
|
|
if not args.debug and accelerator.is_main_process:
|
|
job_name = osp.basename(config.work_dir)
|
|
save_git_snapshot(config.work_dir, job_name, logger)
|
|
|
|
config.train.seed = init_random_seed(getattr(config.train, "seed", None))
|
|
set_random_seed(config.train.seed + int(os.environ["LOCAL_RANK"]))
|
|
generator = torch.Generator(device="cpu").manual_seed(config.train.seed)
|
|
|
|
if accelerator.is_main_process:
|
|
pyrallis.dump(config, open(osp.join(config.work_dir, "config.yaml"), "w"), sort_keys=False, indent=4)
|
|
if args.report_to == "wandb":
|
|
import wandb
|
|
|
|
wandb.init(project=args.tracker_project_name, name=args.name, resume="allow", id=args.name)
|
|
|
|
config.global_world_size = get_world_size()
|
|
logger.info(f"Config: \n{config}")
|
|
logger.info(f"World_size: {config.global_world_size}, seed: {config.train.seed}")
|
|
logger.info(f"Initializing: {init_train} for training")
|
|
# scheduler
|
|
pred_sigma = getattr(config.scheduler, "pred_sigma", True)
|
|
learn_sigma = getattr(config.scheduler, "learn_sigma", True) and pred_sigma
|
|
|
|
# VAE
|
|
vae = None
|
|
vae_dtype = get_weight_dtype(config.vae.weight_dtype)
|
|
vae = get_vae(
|
|
config.vae.vae_type, config.vae.vae_pretrained, accelerator.device, dtype=vae_dtype, config=config.vae
|
|
)
|
|
|
|
logger.info(f"vae type: {config.vae.vae_type}, path: {config.vae.vae_pretrained}, weight_dtype: {vae_dtype}")
|
|
|
|
# Text encoder
|
|
max_length = config.text_encoder.model_max_length
|
|
tokenizer = text_encoder = text_handler = None
|
|
if not config.data.load_text_feat:
|
|
tokenizer, text_encoder = get_tokenizer_and_text_encoder(
|
|
name=config.text_encoder.text_encoder_name, device=accelerator.device
|
|
)
|
|
if "Qwen" in config.text_encoder.text_encoder_name:
|
|
text_handler = text_encoder
|
|
text_encoder = text_handler.text_encoder
|
|
text_embed_dim = text_encoder.config.hidden_size
|
|
else:
|
|
text_embed_dim = config.text_encoder.caption_channels
|
|
|
|
if config.text_encoder.chi_prompt:
|
|
chi_prompt = "\n".join(config.text_encoder.chi_prompt)
|
|
logger.info(f"Complex Human Instruct: {chi_prompt}")
|
|
|
|
os.makedirs(config.train.null_embed_root, exist_ok=True)
|
|
null_embed_path = osp.join(
|
|
config.train.null_embed_root,
|
|
f"null_embed_diffusers_{config.text_encoder.text_encoder_name}_{max_length}token_{text_embed_dim}.pth",
|
|
)
|
|
|
|
image_encoder, image_processor = None, None
|
|
|
|
# 2. build scheduler
|
|
train_diffusion = Scheduler(
|
|
str(config.scheduler.train_sampling_steps),
|
|
noise_schedule=config.scheduler.noise_schedule,
|
|
predict_flow_v=config.scheduler.predict_flow_v,
|
|
learn_sigma=learn_sigma,
|
|
pred_sigma=pred_sigma,
|
|
snr=config.train.snr_loss,
|
|
flow_shift=config.scheduler.flow_shift,
|
|
)
|
|
predict_info = (
|
|
f"flow-prediction: {config.scheduler.predict_flow_v}, noise schedule: {config.scheduler.noise_schedule}"
|
|
)
|
|
if "flow" in config.scheduler.noise_schedule:
|
|
predict_info += f", flow shift: {config.scheduler.flow_shift}"
|
|
if config.scheduler.inference_flow_shift is not None:
|
|
predict_info += f", inference flow shift: {config.scheduler.inference_flow_shift}"
|
|
if config.scheduler.weighting_scheme in ["logit_normal", "mode"]:
|
|
predict_info += (
|
|
f", flow weighting: {config.scheduler.weighting_scheme}, "
|
|
f"logit-mean: {config.scheduler.logit_mean}, logit-std: {config.scheduler.logit_std}"
|
|
)
|
|
logger.info(predict_info)
|
|
|
|
# 3. build dataloader
|
|
config.data.data_dir = (
|
|
config.data.data_dir if isinstance(config.data.data_dir, dict) else {"default": config.data.data_dir}
|
|
)
|
|
config.data.data_dir = {
|
|
k: data if data.startswith(("https://", "http://", "gs://", "/", "~")) else osp.abspath(osp.expanduser(data))
|
|
for k, data in config.data.data_dir.items()
|
|
}
|
|
config.image_data.data_dir = (
|
|
config.image_data.data_dir if isinstance(config.image_data.data_dir, list) else [config.image_data.data_dir]
|
|
)
|
|
config.image_data.data_dir = [
|
|
data if data.startswith(("https://", "http://", "gs://", "/", "~")) else osp.abspath(osp.expanduser(data))
|
|
for data in config.image_data.data_dir
|
|
]
|
|
|
|
num_replicas = int(os.environ["WORLD_SIZE"])
|
|
rank = int(os.environ["RANK"])
|
|
joint_training_interval = config.train.joint_training_interval
|
|
|
|
# video dataset
|
|
|
|
set_random_seed(int(time.time()) % (2**31) + int(os.environ["LOCAL_RANK"]))
|
|
|
|
if config.model.aspect_ratio_type is not None:
|
|
config.data.aspect_ratio_type = config.model.aspect_ratio_type
|
|
dataset = build_dataset(
|
|
asdict(config.data),
|
|
resolution=config.data.image_size,
|
|
max_length=max_length,
|
|
config=config,
|
|
caption_proportion=config.data.caption_proportion,
|
|
sort_dataset=config.data.sort_dataset,
|
|
vae_downsample_rate=config.vae.vae_stride[-1],
|
|
num_frames=config.data.num_frames,
|
|
)
|
|
sampler = DistributedRangedSampler(dataset, num_replicas=num_replicas, rank=rank)
|
|
|
|
if joint_training_interval > 0:
|
|
# image dataset
|
|
if config.model.aspect_ratio_type is not None:
|
|
config.image_data.aspect_ratio_type = config.model.aspect_ratio_type
|
|
dataset_image = build_dataset(
|
|
asdict(config.image_data),
|
|
resolution=config.image_data.image_size,
|
|
max_length=max_length,
|
|
config=config,
|
|
caption_proportion=config.image_data.caption_proportion,
|
|
sort_dataset=config.image_data.sort_dataset,
|
|
vae_downsample_rate=config.vae.vae_stride[-1],
|
|
num_frames=config.image_data.num_frames,
|
|
)
|
|
|
|
image_sampler = DistributedRangedSampler(dataset_image, num_replicas=num_replicas, rank=rank)
|
|
|
|
if config.model.multi_scale:
|
|
batch_sampler = AspectRatioBatchSamplerVideo(
|
|
sampler=sampler,
|
|
dataset=dataset,
|
|
batch_size=config.train.train_batch_size,
|
|
aspect_ratios=dataset.aspect_ratio,
|
|
drop_last=True,
|
|
ratio_nums=dataset.ratio_nums,
|
|
config=config,
|
|
valid_num=config.data.valid_num,
|
|
)
|
|
train_dataloader = build_dataloader(
|
|
dataset, batch_sampler=batch_sampler, num_workers=config.train.num_workers, dataloader_type="video"
|
|
)
|
|
train_dataloader_len = len(train_dataloader)
|
|
|
|
if joint_training_interval > 0:
|
|
batch_sampler_image = AspectRatioBatchSampler(
|
|
sampler=image_sampler,
|
|
dataset=dataset_image,
|
|
batch_size=config.train.train_batch_size_image,
|
|
aspect_ratios=dataset_image.aspect_ratio,
|
|
drop_last=True,
|
|
ratio_nums=dataset_image.ratio_nums,
|
|
config=config,
|
|
clipscore_filter_thres=args.data.del_img_clip_thr,
|
|
)
|
|
train_dataloader_image = build_dataloader(
|
|
dataset_image,
|
|
batch_sampler=batch_sampler_image,
|
|
num_workers=config.train.num_workers,
|
|
dataloader_type="image",
|
|
)
|
|
train_dataloader_image_len = len(train_dataloader_image)
|
|
else:
|
|
train_dataloader_image = iter([])
|
|
train_dataloader_image_len = 1
|
|
|
|
else:
|
|
train_dataloader = build_dataloader(
|
|
dataset,
|
|
num_workers=config.train.num_workers,
|
|
batch_size=config.train.train_batch_size,
|
|
shuffle=False,
|
|
sampler=sampler,
|
|
dataloader_type="video",
|
|
)
|
|
train_dataloader_len = len(train_dataloader)
|
|
|
|
if joint_training_interval > 0:
|
|
# Build image dataloader for joint training
|
|
train_dataloader_image = build_dataloader(
|
|
dataset_image,
|
|
num_workers=config.train.num_workers,
|
|
batch_size=config.train.train_batch_size_image,
|
|
shuffle=False,
|
|
sampler=image_sampler,
|
|
dataloader_type="image",
|
|
)
|
|
train_dataloader_image_len = len(train_dataloader_image)
|
|
else:
|
|
train_dataloader_image = iter([])
|
|
train_dataloader_image_len = 1
|
|
|
|
logger.info(
|
|
f"Video set DataLoader length: {train_dataloader_len}, Image DataLoader length: {train_dataloader_image_len}"
|
|
)
|
|
logger.info(
|
|
colored(
|
|
f"Joint training mode enabled: Image data will be trained with every {joint_training_interval} video iterations",
|
|
"red",
|
|
)
|
|
)
|
|
load_vae_feat = getattr(train_dataloader.dataset, "load_vae_feat", False)
|
|
load_text_feat = getattr(train_dataloader.dataset, "load_text_feat", False)
|
|
|
|
# prepare input for visualization during training
|
|
# aspect_ratio_key = random.choice(list(dataset.aspect_ratio.keys()))
|
|
aspect_ratio_key = "0.57"
|
|
video_height, video_width = map(int, dataset.aspect_ratio[aspect_ratio_key])
|
|
num_frames = config.data.num_frames
|
|
latent_width = int(video_width) // config.vae.vae_stride[2]
|
|
latent_height = int(video_height) // config.vae.vae_stride[1]
|
|
latent_temp = int(num_frames - 1) // config.vae.vae_stride[0] + 1
|
|
|
|
validation_noise = (
|
|
torch.randn(
|
|
1, config.vae.vae_latent_dim, latent_temp, latent_height, latent_width, device="cpu", generator=generator
|
|
)
|
|
if getattr(config.train, "deterministic_validation", False)
|
|
else None
|
|
)
|
|
|
|
if not config.data.load_vae_feat and config.vae.cache_dir is not None:
|
|
vae_cache_dir = os.path.join(
|
|
config.vae.cache_dir,
|
|
f"{config.vae.vae_type}_{num_frames}x{video_height}x{video_width}",
|
|
)
|
|
os.makedirs(vae_cache_dir, exist_ok=True)
|
|
vae.cfg.cache_dir = vae_cache_dir
|
|
logger.info(f"Cache VAE latent of {num_frames}x{video_height}x{video_width} to {vae_cache_dir}")
|
|
|
|
# 4.preparing embeddings for visualization. We put it here for saving GPU memory
|
|
if config.train.visualize and len(config.train.validation_prompts):
|
|
valid_prompt_embed_suffix = f"{max_length}token_{config.text_encoder.text_encoder_name}_{text_embed_dim}.pth"
|
|
validation_prompts = config.train.validation_prompts
|
|
skip = True
|
|
if config.text_encoder.chi_prompt:
|
|
uuid_sys_prompt = hashlib.sha256(chi_prompt.encode()).hexdigest()
|
|
else:
|
|
uuid_sys_prompt = hashlib.sha256(b"").hexdigest()
|
|
config.train.valid_prompt_embed_root = osp.join(
|
|
config.train.valid_prompt_embed_root,
|
|
f"{uuid_sys_prompt}_{config.task}_{latent_height}x{latent_width}_{config.vae.vae_type}_{config.model.image_latent_mode}",
|
|
)
|
|
Path(config.train.valid_prompt_embed_root).mkdir(parents=True, exist_ok=True)
|
|
|
|
if config.text_encoder.chi_prompt:
|
|
# Save system prompt to a file
|
|
system_prompt_file = osp.join(config.train.valid_prompt_embed_root, "system_prompt.txt")
|
|
with open(system_prompt_file, "w", encoding="utf-8") as f:
|
|
f.write(chi_prompt)
|
|
|
|
for prompt in validation_prompts:
|
|
prompt_embed_path = osp.join(
|
|
config.train.valid_prompt_embed_root, f"{prompt[:50]}_{valid_prompt_embed_suffix}"
|
|
)
|
|
if not (osp.exists(prompt_embed_path) and osp.exists(null_embed_path)):
|
|
skip = False
|
|
logger.info(f"Preparing Visualization prompt embeddings at: {config.train.valid_prompt_embed_root}")
|
|
break
|
|
if accelerator.is_main_process and not skip:
|
|
if config.data.load_text_feat and (tokenizer is None or text_encoder is None):
|
|
logger.info(f"Loading text encoder and tokenizer from {config.text_encoder.text_encoder_name} ...")
|
|
tokenizer, text_encoder = get_tokenizer_and_text_encoder(name=config.text_encoder.text_encoder_name)
|
|
|
|
for i, prompt in enumerate(validation_prompts):
|
|
prompt_embed_path = osp.join(
|
|
config.train.valid_prompt_embed_root, f"{prompt[:50]}_{valid_prompt_embed_suffix}"
|
|
)
|
|
if "T5" in config.text_encoder.text_encoder_name:
|
|
txt_tokens = tokenizer(
|
|
prompt, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
|
|
).to(accelerator.device)
|
|
caption_emb = text_encoder(txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask)[0]
|
|
caption_emb_mask = txt_tokens.attention_mask
|
|
elif "gemma" in config.text_encoder.text_encoder_name:
|
|
if not config.text_encoder.chi_prompt:
|
|
max_length_all = config.text_encoder.model_max_length
|
|
else:
|
|
chi_prompt = "\n".join(config.text_encoder.chi_prompt)
|
|
prompt = chi_prompt + prompt
|
|
num_sys_prompt_tokens = len(tokenizer.encode(chi_prompt))
|
|
max_length_all = (
|
|
num_sys_prompt_tokens + config.text_encoder.model_max_length - 2
|
|
) # magic number 2: [bos], [_]
|
|
|
|
txt_tokens = tokenizer(
|
|
prompt,
|
|
max_length=max_length_all,
|
|
padding="max_length",
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
).to(accelerator.device)
|
|
select_index = [0] + list(range(-config.text_encoder.model_max_length + 1, 0))
|
|
caption_emb = text_encoder(txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask)[0][
|
|
:, select_index
|
|
]
|
|
caption_emb_mask = txt_tokens.attention_mask[:, select_index]
|
|
elif "Qwen" in config.text_encoder.text_encoder_name:
|
|
with torch.no_grad():
|
|
y, y_mask = text_encoder.get_prompt_embeds(prompt)
|
|
y_mask = y_mask[:, None, None]
|
|
else:
|
|
raise ValueError(f"{config.text_encoder.text_encoder_name} is not supported!!")
|
|
|
|
save_dict = {"caption_embeds": caption_emb, "emb_mask": caption_emb_mask}
|
|
torch.save(save_dict, prompt_embed_path)
|
|
|
|
if "T5" in config.text_encoder.text_encoder_name:
|
|
null_tokens = tokenizer(
|
|
"", max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
|
|
).to(accelerator.device)
|
|
null_token_emb = text_encoder(null_tokens.input_ids, attention_mask=null_tokens.attention_mask)[0]
|
|
elif "gemma" in config.text_encoder.text_encoder_name:
|
|
null_tokens = tokenizer(
|
|
"", max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
|
|
).to(accelerator.device)
|
|
null_token_emb = text_encoder(null_tokens.input_ids, attention_mask=null_tokens.attention_mask)[0]
|
|
elif "Qwen" in config.text_encoder.text_encoder_name:
|
|
with torch.no_grad():
|
|
null_token_emb, null_token_mask = text_encoder.get_prompt_embeds("")
|
|
null_token_mask = null_token_mask[:, None, None]
|
|
else:
|
|
raise ValueError(f"{config.text_encoder.text_encoder_name} is not supported!!")
|
|
torch.save(
|
|
{"uncond_prompt_embeds": null_token_emb, "uncond_prompt_embeds_mask": null_tokens.attention_mask},
|
|
null_embed_path,
|
|
)
|
|
if config.data.load_text_feat:
|
|
del tokenizer
|
|
del text_encoder
|
|
del null_token_emb
|
|
del null_tokens
|
|
flush()
|
|
|
|
# 5. build models
|
|
os.environ["AUTOCAST_LINEAR_ATTN"] = "true" if config.model.autocast_linear_attn else "false"
|
|
image_size = config.model.image_size
|
|
latent_size = int(image_size) // config.vae.vae_stride[-1]
|
|
model_kwargs = model_video_init_config(config, latent_size=latent_size)
|
|
model = build_model(
|
|
config.model.model,
|
|
config.train.grad_checkpointing,
|
|
getattr(config.model, "fp32_attention", False),
|
|
null_embed_path=null_embed_path,
|
|
**model_kwargs,
|
|
).train()
|
|
|
|
if (not config.train.use_fsdp) and config.train.ema_update:
|
|
model_ema = deepcopy(model).eval()
|
|
logger.info("Creating EMA model for DDP mode")
|
|
elif config.train.use_fsdp and config.train.ema_update:
|
|
logger.warning("EMA update is not supported in FSDP mode. Setting model_ema to None.")
|
|
model_ema = None
|
|
else:
|
|
model_ema = None
|
|
|
|
logger.info(
|
|
colored(
|
|
f"{model.__class__.__name__}:{config.model.model}, "
|
|
f"Model Parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M",
|
|
"green",
|
|
attrs=["bold"],
|
|
)
|
|
)
|
|
|
|
if config.train.use_fsdp:
|
|
model_instance = deepcopy(model)
|
|
elif model_ema is not None:
|
|
model_instance = deepcopy(model_ema)
|
|
else:
|
|
model_instance = model
|
|
|
|
# 5-1. load model
|
|
if args.load_from is not None:
|
|
config.model.load_from = args.load_from
|
|
if config.model.load_from is not None and load_from:
|
|
|
|
load_result = load_checkpoint(
|
|
checkpoint=config.model.load_from,
|
|
model=model,
|
|
model_ema=model_ema,
|
|
FSDP=config.train.use_fsdp,
|
|
load_ema=config.model.resume_from.get("load_ema", False),
|
|
null_embed_path=null_embed_path,
|
|
)
|
|
|
|
_, missing, unexpected, _, _ = load_result
|
|
logger.warning(colored(f"Missing keys: {missing}", "red"))
|
|
logger.warning(colored(f"Unexpected keys: {unexpected}", "red"))
|
|
|
|
if config.train.ema_update and not config.train.use_fsdp and model_ema is not None:
|
|
ema_update(model_ema, model, 0.0)
|
|
|
|
# 5-2. model growth
|
|
if config.model_growth is not None:
|
|
from diffusion.model.model_growth_utils import ModelGrowthInitializer
|
|
|
|
assert config.model.load_from is None
|
|
model_growth_initializer = ModelGrowthInitializer(model, config.model_growth)
|
|
model = model_growth_initializer.initialize(
|
|
strategy=config.model_growth.init_strategy, **config.model_growth.init_params
|
|
)
|
|
|
|
# 6. build optimizer and lr scheduler
|
|
lr_scale_ratio = 1
|
|
if getattr(config.train, "auto_lr", None):
|
|
lr_scale_ratio = auto_scale_lr(
|
|
config.train.train_batch_size * get_world_size() * config.train.gradient_accumulation_steps,
|
|
config.train.optimizer,
|
|
**config.train.auto_lr,
|
|
)
|
|
optimizer = build_optimizer(model, config.train.optimizer)
|
|
|
|
if config.train.lr_schedule_args and config.train.lr_schedule_args.get("num_warmup_steps", None):
|
|
config.train.lr_schedule_args["num_warmup_steps"] = (
|
|
config.train.lr_schedule_args["num_warmup_steps"] * num_replicas
|
|
)
|
|
lr_scheduler = build_lr_scheduler(config.train, optimizer, train_dataloader, lr_scale_ratio)
|
|
logger.warning(
|
|
f"{colored(f'Basic Training Settings: ', 'green', attrs=['bold'])}"
|
|
f"lr: {config.train.optimizer['lr']:.5f}, bs: {config.train.train_batch_size}, gc: {config.train.grad_checkpointing}, "
|
|
f"gc_accum_step: {config.train.gradient_accumulation_steps}."
|
|
)
|
|
logger.info(
|
|
f"{colored(f'Model Settings: ', 'green', attrs=['bold'])}"
|
|
f"qk norm: {config.model.qk_norm}, fp32 attn: {config.model.fp32_attention}, attn type: {config.model.attn_type}, linear_head_dim: {config.model.linear_head_dim}, ffn type: {config.model.ffn_type}, "
|
|
f"text encoder: {config.text_encoder.text_encoder_name}, captions: {config.data.caption_proportion}, precision: {config.model.mixed_precision}."
|
|
)
|
|
|
|
timestamp = time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())
|
|
|
|
if accelerator.is_main_process:
|
|
tracker_config = dict(vars(config))
|
|
try:
|
|
accelerator.init_trackers(args.tracker_project_name, tracker_config)
|
|
except:
|
|
accelerator.init_trackers(f"tb_{timestamp}")
|
|
|
|
start_epoch = 0
|
|
start_step = 0
|
|
start_video_step = 0 # Initialize video step counter
|
|
start_image_step = 0 # Initialize image step counter
|
|
total_steps = train_dataloader_len * config.train.num_epochs
|
|
|
|
# 7. Resume training
|
|
if config.model.resume_from is not None and config.model.resume_from["checkpoint"] is not None:
|
|
rng_state = None
|
|
loaded_image_step = None
|
|
loaded_video_step = None
|
|
ckpt_path = osp.join(config.work_dir, "checkpoints")
|
|
check_flag = osp.exists(ckpt_path) and len(os.listdir(ckpt_path)) != 0
|
|
remove_state_dict_keys = config.model.remove_state_dict_keys
|
|
|
|
if config.model.resume_from["checkpoint"] == "latest":
|
|
if check_flag:
|
|
remove_state_dict_keys = None
|
|
config.model.resume_from["resume_optimizer"] = True
|
|
config.model.resume_from["resume_lr_scheduler"] = True
|
|
checkpoints = os.listdir(ckpt_path)
|
|
if "latest.pth" in checkpoints and osp.exists(osp.join(ckpt_path, "latest.pth")):
|
|
config.model.resume_from["checkpoint"] = osp.realpath(osp.join(ckpt_path, "latest.pth"))
|
|
else:
|
|
checkpoints = [i for i in checkpoints if i.startswith("epoch_")]
|
|
checkpoints = sorted(checkpoints, key=lambda x: int(x.replace(".pth", "").split("_")[3]))
|
|
config.model.resume_from["checkpoint"] = osp.join(ckpt_path, checkpoints[-1])
|
|
else:
|
|
config.model.resume_from["resume_optimizer"] = config.train.load_from_optimizer
|
|
config.model.resume_from["resume_lr_scheduler"] = config.train.load_from_lr_scheduler
|
|
config.model.resume_from["checkpoint"] = config.model.load_from
|
|
|
|
if config.model.resume_from["checkpoint"] is not None:
|
|
|
|
load_result = load_checkpoint(
|
|
**config.model.resume_from,
|
|
model=model,
|
|
model_ema=model_ema if not config.train.use_fsdp else None,
|
|
FSDP=config.train.use_fsdp,
|
|
optimizer=optimizer,
|
|
lr_scheduler=lr_scheduler,
|
|
null_embed_path=null_embed_path,
|
|
remove_state_dict_keys=remove_state_dict_keys,
|
|
)
|
|
|
|
# Handle both old and new return formats
|
|
epoch, missing, unexpected, rng_state, saved_info = load_result
|
|
loaded_video_step = saved_info.get("video_step", None)
|
|
loaded_image_step = saved_info.get("image_step", None)
|
|
|
|
logger.warning(colored(f"Missing keys: {missing}", "red"))
|
|
logger.warning(colored(f"Unexpected keys: {unexpected}", "red"))
|
|
|
|
path = osp.basename(config.model.resume_from["checkpoint"])
|
|
try:
|
|
start_epoch = int(path.replace(".pth", "").split("_")[1]) - 1
|
|
start_step = int(path.replace(".pth", "").split("_")[3])
|
|
except:
|
|
pass
|
|
|
|
# Set video_step and image_step based on availability
|
|
if loaded_video_step is not None:
|
|
start_video_step = loaded_video_step
|
|
logger.info(f"Loaded video_step: {start_video_step} from checkpoint")
|
|
else:
|
|
# If no video_step in checkpoint, use global_step as video_step
|
|
start_video_step = start_step
|
|
logger.info(f"No video_step in checkpoint, using global_step as video_step: {start_video_step}")
|
|
|
|
if loaded_image_step is not None:
|
|
start_image_step = loaded_image_step
|
|
logger.info(f"Loaded image_step: {start_image_step} from checkpoint")
|
|
else:
|
|
# If no image_step in checkpoint, start from 0
|
|
start_image_step = 0
|
|
logger.info(f"No image_step in checkpoint, starting image_step from 0")
|
|
|
|
# 8. Prepare everything
|
|
# There is no specific order to remember, you just need to unpack the
|
|
# objects in the same order you gave them to the prepare method.
|
|
model = accelerator.prepare(model)
|
|
if model_ema is not None and not config.train.use_fsdp:
|
|
model_ema = accelerator.prepare(model_ema)
|
|
optimizer, lr_scheduler = accelerator.prepare(optimizer, lr_scheduler)
|
|
|
|
# load everything except model when resume
|
|
if (
|
|
config.train.use_fsdp
|
|
and config.model.resume_from is not None
|
|
and config.model.resume_from["checkpoint"] is not None
|
|
and config.model.resume_from["resume_optimizer"]
|
|
and config.model.resume_from["resume_lr_scheduler"]
|
|
):
|
|
logger.info(f"FSDP resume: Loading optimizer, scheduler, scaler, random_states...")
|
|
accelerator.load_state(
|
|
os.path.join(config.model.resume_from["checkpoint"], "model"),
|
|
state_dict_key=["optimizer", "scheduler", "scaler", "random_states"],
|
|
)
|
|
|
|
set_random_seed((start_step + 1) // config.train.save_model_steps + int(os.environ["LOCAL_RANK"]))
|
|
logger.info(f'Set seed: {(start_step + 1) // config.train.save_model_steps + int(os.environ["LOCAL_RANK"])}')
|
|
|
|
# Start Training
|
|
train(
|
|
config=config,
|
|
args=args,
|
|
accelerator=accelerator,
|
|
model=model,
|
|
model_ema=model_ema,
|
|
optimizer=optimizer,
|
|
lr_scheduler=lr_scheduler,
|
|
train_dataloader=train_dataloader,
|
|
train_dataloader_image=train_dataloader_image,
|
|
train_diffusion=train_diffusion,
|
|
logger=logger,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|