utils.py

import os
import logging
import time
import copy
import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torchvision.datasets import ImageFolder
from torchvision import transforms
import numpy as np
from collections import OrderedDict
from PIL import Image
from copy import deepcopy
from glob import glob
from time import time

import importlib
import datetime
from omegaconf import OmegaConf
try:
    import wandb
except ImportError:  # pragma: no cover
    wandb = None


class WandbWriter:
    def __init__(self, run):
        self._run = run

    def add_scalar(self, key: str, value: float, step: int | None = None) -> None:
        if self._run is None:
            return
        data = {key: value}
        if step is not None:
            self._run.log(data, step=step)
        else:
            self._run.log(data)

    def finish(self) -> None:
        if self._run is not None:
            self._run.finish()
            self._run = None


def find_model(model_name, is_train=False):

    #start_time = time.time()
    print('model_name: ',model_name)
    checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage, mmap=True, weights_only=False)
    #end_time = time.time()
    # load_time = end_time - start_time
    # print(f"Model loading time: {load_time:.2f} seconds")
    
    if "ema" in checkpoint:  # supports checkpoints from train.py
        checkpoint = checkpoint["ema"] 
        print("load ema ckpt")
        return checkpoint
    elif ("model" in checkpoint) and is_train: 
        checkpoint = checkpoint["model"]
        print("load non-ema ckpt")

        return checkpoint




@torch.no_grad()
def update_ema(ema_model, model, decay=0.9999):
    """
    Step the EMA model towards the current model.
    """
    ema_params = OrderedDict(ema_model.named_parameters())
    model_params = OrderedDict(model.named_parameters())

    for name, param in model_params.items():
        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)

    ema_buffers = OrderedDict(ema_model.named_buffers())
    model_buffers = OrderedDict(model.named_buffers())

    for name, buffer in model_buffers.items():
        ema_buffers[name].mul_(decay).add_(buffer.data, alpha=1 - decay)


def requires_grad(model, flag=True):
    """
    Set requires_grad flag for all parameters in a model.
    """
    for p in model.parameters():
        p.requires_grad = flag


def cleanup():
    """
    End DDP training.
    """
    dist.destroy_process_group()


def create_logger(logging_dir):
    """
    Create a logger that writes to a log file and stdout.
    """
    if dist.get_rank() == 0:  # real logger
        logging.basicConfig(
            level=logging.INFO,
            format='[\033[34m%(asctime)s\033[0m] %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S',
            handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
        )
        logger = logging.getLogger(__name__)
    else:  # dummy logger (does nothing)
        logger = logging.getLogger(__name__)
        logger.addHandler(logging.NullHandler())
    return logger



def center_crop_arr(pil_image, image_size):
    """
    Center cropping implementation from ADM.
    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
    """
    while min(*pil_image.size) >= 2 * image_size:
        pil_image = pil_image.resize(
            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
        )

    scale = image_size / min(*pil_image.size)
    pil_image = pil_image.resize(
        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
    )

    arr = np.array(pil_image)
    crop_y = (arr.shape[0] - image_size) // 2
    crop_x = (arr.shape[1] - image_size) // 2
    return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])


def setup_ddp(args):
    # Setup DDP:
    dist.init_process_group("nccl")
    assert args.global_batch_size % dist.get_world_size() == 0, f"Batch size must be divisible by world size."
    rank = dist.get_rank()
    device = rank % torch.cuda.device_count()
    seed = args.global_seed * dist.get_world_size() + rank
    torch.manual_seed(seed)
    torch.cuda.set_device(device)
    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")

    return rank, device, seed


def setup_exp_dir(rank, args):
    config = copy.deepcopy(args)
    args = args.basic
    # Setup an experiment folder:
    if rank == 0:
        os.makedirs(args.results_dir, exist_ok=True)  

        experiment_index = len(glob(f"{args.results_dir}/*"))
        temp = glob(f"{args.results_dir}/*")
        temp = [int(os.path.basename(_).split("-")[0]) for _ in temp]
        experiment_index = max(temp) + 1 if len(temp) > 0 else 0


        model_string_name = args.exp_name

        experiment_dir = f"{args.results_dir}/{experiment_index:04d}-{model_string_name}-GPU{dist.get_world_size()}"  
        
        checkpoint_dir = f"{experiment_dir}/checkpoints"  

        os.makedirs(checkpoint_dir, exist_ok=True)
        logger = create_logger(experiment_dir)
        logger.info(f"Experiment directory created at {experiment_dir}")

        wandb_run = None
        wandb_project = getattr(args, "wandb_project", None)
        wandb_name = getattr(args, "wandb_name", model_string_name)
        wandb_entity = getattr(args, "wandb_entity", None)
        enable_wandb = getattr(args, "wandb", True) and wandb is not None
        if getattr(args, "wandb", True) and wandb is None:
            logger.warning("wandb package not found; skipping wandb logging.")
        if enable_wandb:
            wandb_kwargs = {
                "project": wandb_project or os.environ.get("WANDB_PROJECT", "fast-dit"),
                "name": wandb_name,
                "config": OmegaConf.to_container(config, resolve=True),
                "dir": experiment_dir,
            }
            if wandb_entity:
                wandb_kwargs["entity"] = wandb_entity
            try:
                wandb_run = wandb.init(**wandb_kwargs)
                logger.info("Initialized Weights & Biases run: %s", wandb_run.name)
            except Exception as exc:
                logger.warning("Failed to initialize wandb logging: %s", exc)
                wandb_run = None

        writer = WandbWriter(wandb_run)

        now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
        OmegaConf.save(config, os.path.join(experiment_dir, f"config_{now}.yaml"))

        return logger, writer, checkpoint_dir
    else:
        logger = create_logger(None)
        return logger, WandbWriter(None), None



def setup_data(rank, args, mode="dinov3"):

    if mode == "vae":
        transform = transforms.Compose([
            transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
        ])
    else:
        transform = transforms.Compose([
            # transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),  
            transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            # transforms.Normalize(
            #     mean=(0.485, 0.456, 0.406),  # ImageNet mean
            #     std=(0.229, 0.224, 0.225),   # ImageNet std
            #     inplace=True
            # ),
        ])

    dataset = ImageFolder(args.data_path, transform=transform)

    sampler = DistributedSampler(
        dataset,
        num_replicas=dist.get_world_size(),
        rank=rank,
        shuffle=True,
        seed=args.global_seed
    )
    loader = DataLoader(
        dataset,
        batch_size=int(args.global_batch_size // dist.get_world_size()),
        shuffle=False,
        sampler=sampler,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=True
    )
    return dataset, sampler, loader



def get_obj_from_str(string, reload=False):
    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 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_lr_scheduler_config(config, optimizer):
    new_dict = {}
    new_dict["target"] = config["target"]
    new_dict["params"] = {"optimizer": optimizer}
    for k,v in config.params.items():
        new_dict["params"][k] = v
    return new_dict