import numpy as np import cv2 import torch import numpy as np from PIL import Image import torch import torch.nn as nn import os import shutil from absl import logging import sys from pathlib import Path from tqdm import tqdm from omegaconf import OmegaConf from torch.utils.data import DataLoader, DistributedSampler import datetime import os.path as osp import torch.distributed as dist import builtins import accelerate import wandb import re from diffusers.training_utils import EMAModel from rich import print def get_obj_from_str(string, reload=False): import importlib module, cls = string.rsplit(".", 1) if reload: module_imp = importlib.import_module(module) importlib.reload(module_imp) return getattr(importlib.import_module(module, package=None), cls) def tensor_detail(t): assert type(t) == torch.Tensor print(f"shape: {t.shape} mean: {t.mean():.2f}, std: {t.std():.2f}, min: {t.min():.2f}, max: {t.max():.2f}") def instantiate_from_config(config): if not "target" in config: raise KeyError("Expected key `target` to instantiate.") model = get_obj_from_str(config["target"])(**config.get("params", dict())) if config.get("resume", False): print(f"resume from: {config.get('resume')}") if os.path.isfile(config.get("resume")): model.load_state_dict(torch.load(config["resume"], map_location="cpu")) elif os.path.isdir(config.get("resume")) and hasattr(model, "from_pretrained"): model.from_pretrained(config.get("resume")) else: raise Exception("could not resume") return model def set_logger(log_level='info', fname=None): import logging as _logging handler = logging.get_absl_handler() formatter = _logging.Formatter('%(asctime)s - %(filename)s - %(message)s') handler.setFormatter(formatter) logging.set_verbosity(log_level) if fname is not None: handler = _logging.FileHandler(fname) handler.setFormatter(formatter) logging.get_absl_logger().addHandler(handler) def dct2str(dct): return str({k: f'{v:.6g}' for k, v in dct.items()}) def copy_files_by_suffix(source_dir, target_dir, suffixes=[".py"], exclude_dirs=[]): # Walk through the directory tree for root, _, files in os.walk(source_dir): if any(exclude_dir in root for exclude_dir in exclude_dirs): continue for file in files: # Check if the file has one of the specified suffixes if any(file.endswith(suffix) for suffix in suffixes): # Construct the source and target paths source_path = os.path.join(root, file) relative_path = os.path.relpath(source_path, source_dir) target_path = os.path.join(target_dir, relative_path) # Ensure the target directory exists os.makedirs(os.path.dirname(target_path), exist_ok=True) # Copy the file shutil.copyfile(source_path, target_path) def find_latest_step(regex, ckpt_root): if not isinstance(regex, re.Pattern): regex = re.compile(regex) ints = [] for file in os.listdir(ckpt_root): if re.match(regex, file): ints.append(int(re.findall(r'\d+', file)[0])) if len(ints) == 0: raise FileNotFoundError(f"no file match {regex} in {ckpt_root}") return max(ints) def resume_from_workdir(config, accelerator, model_context, ema_context): if config.get("resume", False): with PrintContext(f"resume from {config.workdir}", accelerator.is_main_process): for name in config.save_models: max_step = find_latest_step(f"{name}-(\d+).pt", config.ckpt_root) print(f"resume from {name}-{max_step}.pt") model_context[name].load_state_dict( torch.load(osp.join(config.ckpt_root, f"{name}-{max_step}.pt"), map_location="cpu") ) for k, ema in ema_context.items(): max_step = find_latest_step(f"{k}-ema-(\d+).pt", config.ckpt_root) print(f"resume from {k}-ema-{max_step}.pt") ema.load_state_dict( torch.load(osp.join(config.ckpt_root, f"{k}-ema-{max_step}.pt"), map_location="cpu") ) ema.to(accelerator.device) return max_step else: return 0 def get_model_context(models, device, dtype): model_context = dict() for key, model_config in models.items(): model = instantiate_from_config(model_config) if hasattr(model, "device"): try: model.device = device except Exception as e: print(e) print('passing set device') if "t5" in type(model).__name__.lower() and isinstance(model, nn.Module): # T5 model has a bug that it when using fp16 print(f"{'passing t5 model':-^72}") model_context[key] = model.to(device) continue if isinstance(model, nn.Module): model_context[key] = model.to(device=device, dtype=dtype) else: model_context[key] = model model_context["device"] = device model_context["dtype"] = dtype return model_context def get_ema_context(model_context, emas): """given config of ema models and model context, return an ema_context contains all ema model in the current train process Args: model_context (dict): dict of names, point to pytroch models emas (dict): dict of name, point to ema model, name was same with """ ema_context = dict() if emas is None: return ema_context for ema_item in emas: name = ema_item["name"] ema_context[name] = EMAModel(model_context[name].parameters(), **ema_item.params) return ema_context def get_data_context(data, accelerator=None): data_context = dict() for key, data_config in data.items(): dataset = instantiate_from_config(data_config.dataset) if data_config.get("distributed_sampler", False): sampler_cls = get_obj_from_str(data_config.distributed_sampler.target) distributed_sampler = sampler_cls( dataset, num_replicas=accelerator.num_processes if accelerator is not None else 1, rank=accelerator.process_index if accelerator is not None else 0, **data_config.distributed_sampler.params ) dataloader = DataLoader(dataset, sampler=distributed_sampler, **data_config.dataloader) else: dataloader = DataLoader(dataset, **data_config.dataloader) data_context[key] = dataloader data_context[key + "_generator"] = get_data_generator(dataloader, accelerator.is_main_process if accelerator is not None else True, key) data_context[key + "_dataset"] = dataset return data_context class Unimodel(torch.nn.Module): def __init__(self, *args, **kwargs): super().__init__() self._module_list = nn.ModuleList(*args) for k, v in kwargs.items(): setattr(self, k, v) def config_optimizer(model_context, optimizer_models, default_opt_params): """ model_context: dict of model instances optimizer_models: list of dict, each dict contains model name and modules default_opt_params: dict of default optimizer parameters """ default_opt_params = dict(default_opt_params) param_groups = [] for model_config in optimizer_models: model = model_context[model_config["name"]] if model_config.get("modules", None) is None: # all model when no sub modules specified model.requires_grad_(True) print(f"using all modules of {model_config['name']}") para_dict = default_opt_params.copy() opt_params = model_config.get("opt_params", dict()) para_dict.update(opt_params) para_dict["params"] = list(model.parameters()) param_groups.append(para_dict) else: model.requires_grad_(False) for module_config in model_config["modules"]: para_dict = default_opt_params.copy() params = [] for name, param in model.named_parameters(): if module_config["name"] in name: print(name) param.requires_grad = True params.append(param) para_dict["params"] = params opt_params = model_config.get("opt_params", dict()) para_dict.update(opt_params) param_groups.append(para_dict) return param_groups def cnt_params(model): return sum(param.numel() for param in model.parameters()) def get_hparams(input_args=None): argv = sys.argv if input_args is None else input_args lst = [] for i in range(len(argv)): if argv[i].startswith('config.'): hparam_full, val = argv[i].split('=') hparam = hparam_full.split('.')[-1] lst.append(f'{hparam}={val}') hparams = '-'.join(lst) if hparams == '': hparams = 'default' return hparams def add_prefix(dct, prefix): return {f'{prefix}/{key}': val for key, val in dct.items()} def grad_norm(model): total_norm = 0. for p in model.parameters(): if p.grad is not None: param_norm = p.grad.data.norm(2) total_norm += param_norm.item() ** 2 total_norm = total_norm ** (1. / 2) return total_norm def param_norm(model): total_norm = 0. for p in model.parameters(): param_norm = p.data.norm(2) total_norm += param_norm.item() ** 2 total_norm = total_norm ** (1. / 2) return total_norm class PrintContext(object): def __init__(self, name, verbose=True): self.name = name self.verbose = verbose def __enter__(self): if self.verbose: print(f'{self.name} processing...') def __exit__(self, exc_type, exc_val, exc_tb): if self.verbose: print(f'{self.name} done') def time_to_tensor(now: datetime.datetime): return torch.tensor([now.year, now.month, now.day, now.hour, now.minute, now.second], dtype=torch.long) def tensor_to_time(t: torch.Tensor): return datetime.datetime(*t.tolist()) def setup(config, unk): accelerator = accelerate.Accelerator(gradient_accumulation_steps=config.gradient_accumulation_steps) device = accelerator.device accelerate.utils.set_seed(config.seed, device_specific=True) # sync time for all processes g_handler = dist.new_group(backend='gloo') now = time_to_tensor(datetime.datetime.now()) dist.broadcast(now, src=0, group=g_handler) now = tensor_to_time(now).strftime("%Y-%m-%dT%H-%M-%S") print("unknow args: ", unk, get_hparams(unk)) if config.get("workdir", None) is None: config.workdir = osp.join(config.logdir, f"{config.config_name}-{get_hparams(unk)}-{now}") print(f"{'workdir: ' + config.workdir:-^72}") config.ckpt_root = osp.join(config.workdir, 'ckpts') config.eval_root = osp.join(config.workdir, "eval") config.eval_root2 = osp.join(config.workdir, "eval2") if accelerator.is_main_process: os.makedirs(config.workdir, exist_ok=True) os.makedirs(config.ckpt_root, exist_ok=True) os.makedirs(config.eval_root, exist_ok=True) os.makedirs(config.eval_root2, exist_ok=True) config.meta_dir = osp.join(config.workdir, f"meta-{now}") copy_files_by_suffix(os.getcwd(), config.meta_dir, exclude_dirs=[config.logdir], suffixes=[".py", ".yaml"]) with open(osp.join(config.meta_dir, "config.yaml"), "w") as f: f.write(OmegaConf.to_yaml(config)) wandb.init(dir=os.path.abspath(config.workdir), project=config.project, config=dict(config), name=config.wandb_run_name, job_type='train', mode=config.wandb_mode, group="DDP") if accelerator.is_main_process: set_logger(log_level='info', fname=os.path.join(config.workdir, 'output.log')) print(OmegaConf.to_yaml(config)) else: set_logger(log_level='error') builtins.print = lambda *args: None assert not ('total_batch_size' in config and 'batch_size' in config) if 'total_batch_size' not in config: config.total_batch_size = config.batch_size * accelerator.num_processes if 'batch_size' not in config: assert config.total_batch_size % accelerator.num_processes == 0 config.batch_size = config.total_batch_size // accelerator.num_processes if 'total_logical_batch_size' not in config: config.total_logical_batch_size = config.total_batch_size * config.gradient_accumulation_steps logging.info(f'Run on {accelerator.num_processes} devices') return accelerator, device def get_data_generator(loader, enable_tqdm, desc): while True: for data in tqdm(loader, disable=not enable_tqdm, desc=desc): yield data def do_resize_content(original_image: Image, scale_rate): # resize image content wile retain the original image size if scale_rate != 1: # Calculate the new size after rescaling new_size = tuple(int(dim * scale_rate) for dim in original_image.size) # Resize the image while maintaining the aspect ratio resized_image = original_image.resize(new_size) # Create a new image with the original size and black background padded_image = Image.new("RGBA", original_image.size, (0, 0, 0, 0)) paste_position = ((original_image.width - resized_image.width) // 2, (original_image.height - resized_image.height) // 2) padded_image.paste(resized_image, paste_position) return padded_image else: return original_image def add_stroke(img, color=(255, 255, 255), stroke_radius=3): # color in R, G, B format if isinstance(img, Image.Image): assert img.mode == "RGBA" img = cv2.cvtColor(np.array(img), cv2.COLOR_RGBA2BGRA) else: assert img.shape[2] == 4 gray = img[:,:, 3] ret, binary = cv2.threshold(gray,127,255,cv2.THRESH_BINARY) contours, hierarchy = cv2.findContours(binary,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) res = cv2.drawContours(img, contours,-1, tuple(color)[::-1] + (255,), stroke_radius) return Image.fromarray(cv2.cvtColor(res,cv2.COLOR_BGRA2RGBA)) def make_blob(image_size=(512, 512), sigma=0.2): """ make 2D blob image with: I(x, y)=1-\exp \left(-\frac{(x-H / 2)^2+(y-W / 2)^2}{2 \sigma^2 HS}\right) """ import numpy as np H, W = image_size x = np.arange(0, W, 1, float) y = np.arange(0, H, 1, float) x, y = np.meshgrid(x, y) x0 = W // 2 y0 = H // 2 img = 1 - np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2 * H * W)) return (img * 255).astype(np.uint8)