mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-12 12:40:04 +08:00
190 lines
8.3 KiB
Python
190 lines
8.3 KiB
Python
import random
|
|
import logging
|
|
from pathlib import Path
|
|
import shutil
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
import torch.distributed as dist
|
|
from tqdm.auto import tqdm
|
|
import math
|
|
import torch.nn.functional as F
|
|
import os
|
|
|
|
def get_trainable_params(model):
|
|
params = model.parameters()
|
|
params = [p for p in params if p.requires_grad]
|
|
return params
|
|
|
|
|
|
def set_seeds(seed_list, device=None):
|
|
if isinstance(seed_list, (tuple, list)):
|
|
seed = sum(seed_list)
|
|
else:
|
|
seed = seed_list
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
return torch.Generator(device).manual_seed(seed)
|
|
|
|
def get_start_epoch(resume_path, ckpt, steps_per_epoch):
|
|
if 'epoch' in ckpt:
|
|
start_epoch = ckpt['epoch']
|
|
else:
|
|
start_epoch = 0
|
|
if 'steps' in ckpt:
|
|
train_steps = ckpt['steps']
|
|
else:
|
|
try:
|
|
train_steps = int(Path(resume_path).stem)
|
|
except:
|
|
train_steps = start_epoch * steps_per_epoch
|
|
|
|
start_epoch_step = train_steps % steps_per_epoch + 1
|
|
return start_epoch, start_epoch_step, train_steps
|
|
|
|
def assert_shape(*args):
|
|
if len(args) < 2:
|
|
return
|
|
cond = True
|
|
fail_str = f"{args[0] if isinstance(args[0], (list, tuple)) else args[0].shape}"
|
|
for i in range(1, len(args)):
|
|
shape1 = args[i] if isinstance(args[i], (list, tuple)) else args[i].shape
|
|
shape2 = args[i - 1] if isinstance(args[i - 1], (list, tuple)) else args[i - 1].shape
|
|
cond = cond and (shape1 == shape2)
|
|
fail_str += f" vs {args[i] if isinstance(args[i], (list, tuple)) else args[i].shape}"
|
|
assert cond, fail_str
|
|
|
|
|
|
def create_logger(logging_dir=None, logging_file=None, ddp=True):
|
|
"""
|
|
Create a logger that writes to a log file and stdout.
|
|
"""
|
|
if not ddp or (ddp and dist.get_rank() == 0): # real logger
|
|
if logging_file is not None:
|
|
file_handler = [logging.FileHandler(logging_file)]
|
|
elif logging_dir is not None:
|
|
file_handler = [logging.FileHandler(f"{logging_dir}/log.txt")]
|
|
else:
|
|
file_handler = []
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='[\033[34m%(asctime)s\033[0m] %(message)s',
|
|
datefmt='%Y-%m-%d %H:%M:%S',
|
|
handlers=[logging.StreamHandler()] + file_handler
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
else:
|
|
logger = logging.getLogger(__name__)
|
|
logger.addHandler(logging.NullHandler())
|
|
return logger
|
|
|
|
def create_exp_folder(args, rank):
|
|
if rank == 0:
|
|
os.makedirs(args.results_dir, exist_ok=True)
|
|
existed_experiments = list(Path(args.results_dir).glob("*dit*"))
|
|
if len(existed_experiments) == 0:
|
|
experiment_index = 1
|
|
else:
|
|
existed_experiments.sort()
|
|
print('existed_experiments', existed_experiments)
|
|
experiment_index = max([int(x.stem.split('-')[0]) for x in existed_experiments]) + 1
|
|
dist.barrier()
|
|
model_string_name = args.task_flag if args.task_flag else args.model.replace("/", "-")
|
|
experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}" # Create an experiment folder
|
|
checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints
|
|
if rank == 0:
|
|
os.makedirs(checkpoint_dir, exist_ok=True)
|
|
logger = create_logger(experiment_dir)
|
|
logger.info(f"Experiment directory created at {experiment_dir}")
|
|
else:
|
|
logger = create_logger()
|
|
experiment_dir = ""
|
|
|
|
return experiment_dir, checkpoint_dir, logger
|
|
|
|
|
|
def model_resume(args, model, ema, logger, len_loader):
|
|
"""
|
|
Load pretrained weights.
|
|
"""
|
|
start_epoch = 0
|
|
start_epoch_step = 0
|
|
train_steps = 0
|
|
resume_path = args.resume
|
|
if not Path(resume_path).exists():
|
|
raise FileNotFoundError(f" Cannot find checkpoint from {resume_path}")
|
|
|
|
logger.info(f" Resume deepspeed={args.resume_deepspeed}, "
|
|
f"Resume split={args.resume_split}, "
|
|
f"Resume from checkpoint {resume_path}")
|
|
# Resume model and ema states (not include optimizer states) from a checkpoint saved by Deepspeed version of DIT.
|
|
if args.resume_deepspeed:
|
|
assert 'mp_rank_00_model_states.pt' in os.listdir(resume_path), f' Cannot find dp chkpt from {resume_path}'
|
|
resume_ckpt = torch.load(os.path.join(resume_path, 'mp_rank_00_model_states.pt'),
|
|
map_location=lambda storage, loc: storage)
|
|
# Resume main model
|
|
if args.ema_to_module:
|
|
logger.info(" Resume main model from the ema states.")
|
|
model.load_state_dict(resume_ckpt['ema'], strict=args.strict)
|
|
else:
|
|
logger.info(" Resume main model from the main states.")
|
|
model.load_state_dict(resume_ckpt['module'], strict=args.strict)
|
|
# Resume EMA model
|
|
if args.use_ema:
|
|
if args.module_to_ema:
|
|
logger.info(" Resume EMA model from the main states.")
|
|
ema.load_state_dict(resume_ckpt['module'], strict=args.strict)
|
|
else:
|
|
logger.info(" Resume EMA model from the EMA states.")
|
|
ema.load_state_dict(resume_ckpt['ema'], strict=args.strict)
|
|
if not args.reset_loader:
|
|
start_epoch, start_epoch_step, train_steps = get_start_epoch(args.resume, resume_ckpt, len_loader)
|
|
# Resume model and ema states (not include optimizer states) from two checkpoints separated from DeepSpeed ckpt.
|
|
elif args.resume_split:
|
|
# Resume main model
|
|
if args.ema_to_module:
|
|
assert 'pytorch_model_ema.pt' in os.listdir(
|
|
resume_path), f' Cannot find pytorch_model_ema.pt from {resume_path}'
|
|
logger.info(f" Resume main model from ema states.")
|
|
resume_ckpt_ema = torch.load(os.path.join(resume_path, 'pytorch_model_ema.pt'),
|
|
map_location=lambda storage, loc: storage)
|
|
model.load_state_dict(resume_ckpt_ema, strict=args.strict)
|
|
else:
|
|
assert 'pytorch_model_module.pt' in os.listdir(
|
|
resume_path), f' Cannot find pytorch_model_module.pt from {resume_path}'
|
|
logger.info(f" Resume main model from main states.")
|
|
resume_ckpt_module = torch.load(os.path.join(resume_path, 'pytorch_model_module.pt'),
|
|
map_location=lambda storage, loc: storage)
|
|
model.load_state_dict(resume_ckpt_module, strict=args.strict)
|
|
# Resume ema model
|
|
if args.use_ema:
|
|
if args.module_to_ema:
|
|
if "resume_ckpt_module" in locals():
|
|
logger.info(f" Resume ema model from main states.")
|
|
ema.load_state_dict(resume_ckpt_module, strict=args.strict)
|
|
else:
|
|
assert 'pytorch_model_module.pt' in os.listdir(
|
|
resume_path), f' Cannot find pytorch_model_module.pt from {resume_path}'
|
|
logger.info(f" Resume ema model from module states.")
|
|
resume_ckpt_module = torch.load(os.path.join(resume_path, 'pytorch_model_module.pt'),
|
|
map_location=lambda storage, loc: storage)
|
|
ema.load_state_dict(resume_ckpt_module, strict=args.strict)
|
|
else:
|
|
if "resume_ckpt_ema" in locals():
|
|
logger.info(f" Resume ema model from EMA states.")
|
|
ema.load_state_dict(resume_ckpt_ema, strict=args.strict)
|
|
else:
|
|
assert 'pytorch_model_ema.pt' in os.listdir(
|
|
resume_path), f' Cannot find pytorch_model_ema.pt from {resume_path}'
|
|
logger.info(f" Resume ema model from EMA states.")
|
|
resume_ckpt_ema = torch.load(os.path.join(resume_path, 'pytorch_model_ema.pt'),
|
|
map_location=lambda storage, loc: storage)
|
|
ema.load_state_dict(resume_ckpt_ema, strict=args.strict)
|
|
else:
|
|
raise ValueError(" “If `resume` is True, then either `resume_split` or `resume_deepspeed` must be true.”")
|
|
|
|
return model, ema, start_epoch, start_epoch_step, train_steps |