Upload libs/base_utils.py with huggingface_hub

libs/base_utils.py

@@ -1,84 +1,393 @@
-import numpy as np
-import cv2
-import torch
-import numpy as np
-from PIL import Image
-
-        
-def instantiate_from_config(config):
-    if not "target" in config:
-        raise KeyError("Expected key `target` to instantiate.")
-    return get_obj_from_str(config["target"])(**config.get("params", dict()))
-
-
-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 drawRoundRec(draw, color, x, y, w, h, r):
-    drawObject = draw
-
-    '''Rounds'''
-    drawObject.ellipse((x, y, x + r, y + r), fill=color)
-    drawObject.ellipse((x + w - r, y, x + w, y + r), fill=color)
-    drawObject.ellipse((x, y + h - r, x + r, y + h), fill=color)
-    drawObject.ellipse((x + w - r, y + h - r, x + w, y + h), fill=color)
-
-    '''rec.s'''
-    drawObject.rectangle((x + r / 2, y, x + w - (r / 2), y + h), fill=color)
-    drawObject.rectangle((x, y + r / 2, x + w, y + h - (r / 2)), fill=color)
-
-
-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))
+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)