# 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 getpass
import hashlib
import json
import os
import os.path as osp
import random
import time
import types
import warnings
from dataclasses import asdict
from pathlib import Path
import numpy as np
import pyrallis
import torch
from accelerate import Accelerator, InitProcessGroupKwargs
from accelerate.utils import DistributedType
from PIL import Image
from termcolor import colored
warnings.filterwarnings("ignore") # ignore warning
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.utils.checkpoint import load_checkpoint, save_checkpoint
from diffusion.utils.config import SanaConfig
from diffusion.utils.data_sampler import AspectRatioBatchSampler
from diffusion.utils.dist_utils import clip_grad_norm_, flush, get_world_size
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, read_config, set_random_seed
from diffusion.utils.optimizer import auto_scale_lr, build_optimizer
os.environ["TOKENIZERS_PARALLELISM"] = "false"
def set_fsdp_env():
os.environ["ACCELERATE_USE_FSDP"] = "true"
os.environ["FSDP_AUTO_WRAP_POLICY"] = "TRANSFORMER_BASED_WRAP"
os.environ["FSDP_BACKWARD_PREFETCH"] = "BACKWARD_PRE"
os.environ["FSDP_TRANSFORMER_CLS_TO_WRAP"] = "SanaBlock"
@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([[image_size, image_size]], 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)
# Create sampling noise:
logger.info("Running validation... ")
image_logs = []
def run_sampling(init_z=None, label_suffix="", vae=None, sampler="dpm-solver"):
latents = []
current_image_logs = []
for prompt in validation_prompts:
z = (
torch.randn(1, config.vae.vae_latent_dim, latent_size, latent_size, device=device)
if init_z is None
else init_z
)
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)
# caption_embs = caption_embs[:, None]
# emb_masks = emb_masks[:, None]
model_kwargs = dict(data_info={"img_hw": hw, "aspect_ratio": ar}, mask=emb_masks)
if sampler == "dpm-solver":
dpm_solver = DPMS(
model.forward_with_dpmsolver,
condition=caption_embs,
uncondition=null_y,
cfg_scale=4.5,
model_kwargs=model_kwargs,
)
denoised = dpm_solver.sample(
z,
steps=14,
order=2,
skip_type="time_uniform",
method="multistep",
)
elif sampler == "flow_euler":
flow_solver = FlowEuler(
model, condition=caption_embs, uncondition=null_y, cfg_scale=4.5, model_kwargs=model_kwargs
)
denoised = flow_solver.sample(z, steps=28)
elif sampler == "flow_dpm-solver":
dpm_solver = DPMS(
model.forward_with_dpmsolver,
condition=caption_embs,
uncondition=null_y,
cfg_scale=4.5,
model_type="flow",
model_kwargs=model_kwargs,
schedule="FLOW",
)
denoised = dpm_solver.sample(
z,
steps=20,
order=2,
skip_type="time_uniform_flow",
method="multistep",
flow_shift=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).to(torch.float16)
for prompt, latent in zip(validation_prompts, latents):
latent = latent.to(torch.float16)
samples = vae_decode(config.vae.vae_type, vae, latent)
samples = (
torch.clamp(127.5 * samples + 128.0, 0, 255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()[0]
)
image = Image.fromarray(samples)
current_image_logs.append({"validation_prompt": prompt + label_suffix, "images": [image]})
return current_image_logs
# First run with original noise
image_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:
init_noise = torch.clone(init_noise).to(device)
image_logs += run_sampling(init_z=init_noise, label_suffix=" w/ init noise", vae=vae, sampler=vis_sampler)
formatted_images = []
for log in image_logs:
images = log["images"]
validation_prompt = log["validation_prompt"]
for image in images:
formatted_images.append((validation_prompt, np.asarray(image)))
for tracker in accelerator.trackers:
if tracker.name == "tensorboard":
for validation_prompt, image in formatted_images:
tracker.writer.add_images(validation_prompt, image[None, ...], step, dataformats="NHWC")
elif tracker.name == "wandb":
import wandb
wandb_images = []
for validation_prompt, image in formatted_images:
wandb_images.append(wandb.Image(image, caption=validation_prompt, file_type="jpg"))
tracker.log({"validation": wandb_images})
else:
logger.warn(f"image logging not implemented for {tracker.name}")
def concatenate_images(image_caption, images_per_row=5, image_format="webp"):
import io
images = [log["images"][0] for log in image_caption]
if images[0].size[0] > 1024:
images = [image.resize((1024, 1024)) for image in images]
widths, heights = zip(*(img.size for img in images))
max_width = max(widths)
total_height = sum(heights[i : i + images_per_row][0] for i in range(0, len(images), images_per_row))
new_im = Image.new("RGB", (max_width * images_per_row, total_height))
y_offset = 0
for i in range(0, len(images), images_per_row):
row_images = images[i : i + images_per_row]
x_offset = 0
for img in row_images:
new_im.paste(img, (x_offset, y_offset))
x_offset += max_width
y_offset += heights[i]
webp_image_bytes = io.BytesIO()
new_im.save(webp_image_bytes, format=image_format)
webp_image_bytes.seek(0)
new_im = Image.open(webp_image_bytes)
return new_im
if config.train.local_save_vis:
file_format = "webp"
local_vis_save_path = osp.join(config.work_dir, "log_vis")
os.umask(0o000)
os.makedirs(local_vis_save_path, exist_ok=True)
concatenated_image = concatenate_images(image_logs, images_per_row=5, image_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}")
)
concatenated_image.save(save_path)
del vae
flush()
return image_logs
def train(config, args, accelerator, model, optimizer, lr_scheduler, train_dataloader, train_diffusion, logger):
if getattr(config.train, "debug_nan", False):
DebugUnderflowOverflow(model)
logger.info("NaN debugger registered. Start to detect overflow during training.")
log_buffer = LogBuffer()
global_step = start_step + 1
skip_step = max(config.train.skip_step, global_step) % train_dataloader_len
skip_step = skip_step if skip_step < (train_dataloader_len - 20) else 0
loss_nan_timer = 0
# Cache Dataset for BatchSampler
if args.caching and config.model.multi_scale:
caching_start = time.time()
logger.info(
f"Start caching your dataset for batch_sampler at {cache_file}. \n"
f"This may take a lot of time...No training will launch"
)
train_dataloader.batch_sampler.sampler.set_start(max(train_dataloader.batch_sampler.exist_ids, 0))
accelerator.wait_for_everyone()
for index, _ in enumerate(train_dataloader):
accelerator.wait_for_everyone()
if index % 2000 == 0:
logger.info(
f"rank: {rank}, Cached file len: {len(train_dataloader.batch_sampler.cached_idx)} / {len(train_dataloader)}"
)
print(
f"rank: {rank}, Cached file len: {len(train_dataloader.batch_sampler.cached_idx)} / {len(train_dataloader)}"
)
if (time.time() - caching_start) / 3600 > 3.7:
json.dump(train_dataloader.batch_sampler.cached_idx, open(cache_file, "w"), indent=4)
accelerator.wait_for_everyone()
break
if len(train_dataloader.batch_sampler.cached_idx) == len(train_dataloader) - 1000:
logger.info(
f"Saving rank: {rank}, Cached file len: {len(train_dataloader.batch_sampler.cached_idx)} / {len(train_dataloader)}"
)
json.dump(train_dataloader.batch_sampler.cached_idx, open(cache_file, "w"), indent=4)
accelerator.wait_for_everyone()
continue
accelerator.wait_for_everyone()
print(f"Saving rank-{rank} Cached file len: {len(train_dataloader.batch_sampler.cached_idx)}")
json.dump(train_dataloader.batch_sampler.cached_idx, open(cache_file, "w"), indent=4)
return
# 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
)
sampler.set_epoch(epoch)
sampler.set_start(max((skip_step - 1) * config.train.train_batch_size, 0))
if skip_step > 1 and accelerator.is_main_process:
logger.info(f"Skipped Steps: {skip_step}")
skip_step = 1
data_time_start = time.time()
data_time_all = 0
lm_time_all = 0
vae_time_all = 0
model_time_all = 0
for step, batch in enumerate(train_dataloader):
# image, json_info, key = batch
accelerator.wait_for_everyone()
data_time_all += time.time() - data_time_start
vae_time_start = time.time()
if load_vae_feat:
z = batch[0].to(accelerator.device)
else:
with torch.no_grad():
with torch.amp.autocast(
"cuda",
enabled=(config.model.mixed_precision == "fp16" or config.model.mixed_precision == "bf16"),
):
z = vae_encode(
config.vae.vae_type, vae, batch[0], config.vae.sample_posterior, accelerator.device
)
accelerator.wait_for_everyone()
vae_time_all += time.time() - vae_time_start
clean_images = z
data_info = batch[3]
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 or "Qwen" 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_chi_prompt_tokens = len(tokenizer.encode(chi_prompt))
max_length_all = (
num_chi_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]
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"]:
# 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=None, # not used
)
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()
accelerator.backward(loss)
if accelerator.sync_gradients:
grad_norm = accelerator.clip_grad_norm_(model.parameters(), config.train.gradient_clip)
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())
log_buffer.update(logs)
if (step + 1) % config.train.log_interval == 0 or (step + 1) == 1:
accelerator.wait_for_everyone()
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()
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
info = (
f"Epoch: {epoch} | Global Step: {global_step} | Local Step: {current_step} // {train_dataloader_len}, "
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}, Cap: {batch[5][0]}, "
)
info += (
f"s:({model.module.h}, {model.module.w}), "
if hasattr(model, "module")
else f"s:({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.training_hours
):
accelerator.wait_for_everyone()
if accelerator.is_main_process:
os.umask(0o000)
ckpt_saved_path = save_checkpoint(
osp.join(config.work_dir, "checkpoints"),
epoch=epoch,
step=global_step,
model=accelerator.unwrap_model(model),
optimizer=optimizer,
lr_scheduler=lr_scheduler,
generator=generator,
add_symlink=True,
)
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.training_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):
accelerator.wait_for_everyone()
if accelerator.is_main_process:
if validation_noise is not None:
log_validation(
accelerator=accelerator,
config=config,
model=model,
logger=logger,
step=global_step,
device=accelerator.device,
vae=vae,
init_noise=validation_noise,
)
else:
log_validation(
accelerator=accelerator,
config=config,
model=model,
logger=logger,
step=global_step,
device=accelerator.device,
vae=vae,
)
# avoid dead-lock of multiscale data batch sampler
# for internal, refactor dataloader logic to remove the ad-hoc implementation
if (
config.model.multi_scale
and (train_dataloader_len - sampler.step_start // config.train.train_batch_size - step) < 30
):
global_step = epoch * train_dataloader_len
logger.info("Early stop current iteration")
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()
if accelerator.is_main_process:
# os.umask(0o000)
ckpt_saved_path = save_checkpoint(
osp.join(config.work_dir, "checkpoints"),
epoch=epoch,
step=global_step,
model=accelerator.unwrap_model(model),
optimizer=optimizer,
lr_scheduler=lr_scheduler,
generator=generator,
add_symlink=True,
)
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)
accelerator.wait_for_everyone()
@pyrallis.wrap()
def main(cfg: SanaConfig) -> None:
global train_dataloader_len, start_epoch, start_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
config = cfg
args = cfg
# config = read_config(args.config)
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=True,
)
if args.debug:
config.train.log_interval = 1
config.train.train_batch_size = min(64, config.train.train_batch_size)
args.report_to = "tensorboard"
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
if config.train.use_fsdp:
init_train = "FSDP"
from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullStateDictConfig
set_fsdp_env()
fsdp_plugin = FullyShardedDataParallelPlugin(
state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
)
else:
init_train = "DDP"
fsdp_plugin = None
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"),
fsdp_plugin=fsdp_plugin,
kwargs_handlers=[init_handler],
)
log_name = "train_log.log"
logger = get_root_logger(osp.join(config.work_dir, log_name))
logger.info(accelerator.state)
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)
logger.info(f"Config: \n{config}")
logger.info(f"World_size: {get_world_size()}, seed: {config.train.seed}")
logger.info(f"Initializing: {init_train} for training")
image_size = config.model.image_size
latent_size = int(image_size) // config.vae.vae_downsample_rate
pred_sigma = getattr(config.scheduler, "pred_sigma", True)
learn_sigma = getattr(config.scheduler, "learn_sigma", True) and pred_sigma
max_length = config.text_encoder.model_max_length
vae = None
validation_noise = (
torch.randn(1, config.vae.vae_latent_dim, latent_size, latent_size, device="cpu", generator=generator)
if getattr(config.train, "deterministic_validation", False)
else None
)
if not config.data.load_vae_feat:
vae = get_vae(config.vae.vae_type, config.vae.vae_pretrained, accelerator.device).to(torch.float16)
tokenizer = text_encoder = 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
)
text_embed_dim = text_encoder.config.hidden_size
else:
text_embed_dim = config.text_encoder.caption_channels
logger.info(f"vae type: {config.vae.vae_type}")
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",
)
if config.train.visualize and len(config.train.validation_prompts):
# preparing embeddings for visualization. We put it here for saving GPU memory
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_chi_prompt = hashlib.sha256(chi_prompt.encode()).hexdigest()
else:
uuid_chi_prompt = hashlib.sha256(b"").hexdigest()
config.train.valid_prompt_embed_root = osp.join(config.train.valid_prompt_embed_root, uuid_chi_prompt)
Path(config.train.valid_prompt_embed_root).mkdir(parents=True, exist_ok=True)
if config.text_encoder.chi_prompt:
# Save complex human instruct to a file
chi_prompt_file = osp.join(config.train.valid_prompt_embed_root, "chi_prompt.txt")
with open(chi_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("Preparing Visualization prompt embeddings...")
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 prompt in 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 or "Qwen" 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_chi_prompt_tokens = len(tokenizer.encode(chi_prompt))
max_length_all = (
num_chi_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]
else:
raise ValueError(f"{config.text_encoder.text_encoder_name} is not supported!!")
torch.save({"caption_embeds": caption_emb, "emb_mask": caption_emb_mask}, prompt_embed_path)
null_tokens = tokenizer(
"", max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
).to(accelerator.device)
if "T5" in config.text_encoder.text_encoder_name:
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 or "Qwen" in config.text_encoder.text_encoder_name:
null_token_emb = text_encoder(null_tokens.input_ids, attention_mask=null_tokens.attention_mask)[0]
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()
os.environ["AUTOCAST_LINEAR_ATTN"] = "true" if config.model.autocast_linear_attn else "false"
# 1. build scheduler
train_diffusion = Scheduler(
str(config.scheduler.train_sampling_steps),
noise_schedule=config.scheduler.noise_schedule,
predict_v=config.scheduler.predict_v,
learn_sigma=learn_sigma,
pred_sigma=pred_sigma,
snr=config.train.snr_loss,
flow_shift=config.scheduler.flow_shift,
)
predict_info = f"v-prediction: {config.scheduler.predict_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.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)
# 2. build models
model_kwargs = {
"pe_interpolation": config.model.pe_interpolation,
"config": config,
"model_max_length": max_length,
"qk_norm": config.model.qk_norm,
"micro_condition": config.model.micro_condition,
"caption_channels": text_embed_dim,
"y_norm": config.text_encoder.y_norm,
"attn_type": config.model.attn_type,
"ffn_type": config.model.ffn_type,
"mlp_ratio": config.model.mlp_ratio,
"mlp_acts": list(config.model.mlp_acts),
"in_channels": config.vae.vae_latent_dim,
"y_norm_scale_factor": config.text_encoder.y_norm_scale_factor,
"use_pe": config.model.use_pe,
"linear_head_dim": config.model.linear_head_dim,
"pred_sigma": pred_sigma,
"learn_sigma": learn_sigma,
}
model = build_model(
config.model.model,
config.train.grad_checkpointing,
getattr(config.model, "fp32_attention", False),
input_size=latent_size,
**model_kwargs,
).train()
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"],
)
)
# 2-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:
_, missing, unexpected, _ = load_checkpoint(
config.model.load_from,
model,
load_ema=config.model.resume_from.get("load_ema", False),
null_embed_path=null_embed_path,
)
logger.warning(f"Missing keys: {missing}")
logger.warning(f"Unexpected keys: {unexpected}")
# prepare for FSDP clip grad norm calculation
if accelerator.distributed_type == DistributedType.FSDP:
for m in accelerator._models:
m.clip_grad_norm_ = types.MethodType(clip_grad_norm_, m)
# 3. build dataloader
config.data.data_dir = config.data.data_dir if isinstance(config.data.data_dir, list) else [config.data.data_dir]
config.data.data_dir = [
data if data.startswith(("https://", "http://", "gs://", "/", "~")) else osp.abspath(osp.expanduser(data))
for data in config.data.data_dir
]
num_replicas = int(os.environ["WORLD_SIZE"])
rank = int(os.environ["RANK"])
dataset = build_dataset(
asdict(config.data),
resolution=image_size,
aspect_ratio_type=config.model.aspect_ratio_type,
real_prompt_ratio=config.train.real_prompt_ratio,
max_length=max_length,
config=config,
caption_proportion=config.data.caption_proportion,
sort_dataset=config.data.sort_dataset,
vae_downsample_rate=config.vae.vae_downsample_rate,
)
accelerator.wait_for_everyone()
if config.model.multi_scale:
drop_last = True
uuid = hashlib.sha256("-".join(config.data.data_dir).encode()).hexdigest()[:8]
cache_dir = osp.expanduser(f"~/.cache/_wids_batchsampler_cache")
os.makedirs(cache_dir, exist_ok=True)
base_pattern = (
f"{cache_dir}/{getpass.getuser()}-{uuid}-sort_dataset{config.data.sort_dataset}"
f"-hq_only{config.data.hq_only}-valid_num{config.data.valid_num}"
f"-aspect_ratio{len(dataset.aspect_ratio)}-droplast{drop_last}"
f"dataset_len{len(dataset)}"
)
cache_file = f"{base_pattern}-num_replicas{num_replicas}-rank{rank}"
for i in config.data.data_dir:
cache_file += f"-{i}"
cache_file += ".json"
sampler = DistributedRangedSampler(dataset, num_replicas=num_replicas, rank=rank)
batch_sampler = AspectRatioBatchSampler(
sampler=sampler,
dataset=dataset,
batch_size=config.train.train_batch_size,
aspect_ratios=dataset.aspect_ratio,
drop_last=drop_last,
ratio_nums=dataset.ratio_nums,
config=config,
valid_num=config.data.valid_num,
hq_only=config.data.hq_only,
cache_file=cache_file,
caching=args.caching,
)
train_dataloader = build_dataloader(dataset, batch_sampler=batch_sampler, num_workers=config.train.num_workers)
train_dataloader_len = len(train_dataloader)
logger.info(f"rank-{rank} Cached file len: {len(train_dataloader.batch_sampler.cached_idx)}")
else:
sampler = DistributedRangedSampler(dataset, num_replicas=num_replicas, rank=rank)
train_dataloader = build_dataloader(
dataset,
num_workers=config.train.num_workers,
batch_size=config.train.train_batch_size,
shuffle=False,
sampler=sampler,
)
train_dataloader_len = len(train_dataloader)
load_vae_feat = getattr(train_dataloader.dataset, "load_vae_feat", False)
load_text_feat = getattr(train_dataloader.dataset, "load_text_feat", False)
# 4. 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 Setting: ', '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}, qk norm: {config.model.qk_norm}, "
f"fp32 attn: {config.model.fp32_attention}, attn type: {config.model.attn_type}, 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
total_steps = train_dataloader_len * config.train.num_epochs
# Resume training
if config.model.resume_from is not None and config.model.resume_from["checkpoint"] is not None:
rng_state = None
ckpt_path = osp.join(config.work_dir, "checkpoints")
check_flag = osp.exists(ckpt_path) and len(os.listdir(ckpt_path)) != 0
if config.model.resume_from["checkpoint"] == "latest":
if check_flag:
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["checkpoint"] = config.model.load_from
if config.model.resume_from["checkpoint"] is not None:
_, missing, unexpected, rng_state = load_checkpoint(
**config.model.resume_from,
model=model,
optimizer=optimizer if check_flag else None,
lr_scheduler=lr_scheduler if check_flag else None,
null_embed_path=null_embed_path,
)
logger.warning(f"Missing keys: {missing}")
logger.warning(f"Unexpected keys: {unexpected}")
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
# resume randomise
if rng_state:
logger.info("resuming randomise")
torch.set_rng_state(rng_state["torch"])
np.random.set_state(rng_state["numpy"])
random.setstate(rng_state["python"])
generator.set_state(rng_state["generator"]) # resume generator status
try:
torch.cuda.set_rng_state_all(rng_state["torch_cuda"])
except:
logger.warning("Failed to resume torch_cuda rng state")
# 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)
optimizer, lr_scheduler = accelerator.prepare(optimizer, lr_scheduler)
# Start Training
train(
config=config,
args=args,
accelerator=accelerator,
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
train_dataloader=train_dataloader,
train_diffusion=train_diffusion,
logger=logger,
)
if __name__ == "__main__":
main()