fred-dev/comfy_ui_ali

1from contextlib import contextmanager2import hashlib3import math4from pathlib import Path5import shutil6import urllib7import warnings8 9from PIL import Image10import torch11from torch import nn, optim12from torch.utils import data13 14 15def hf_datasets_augs_helper(examples, transform, image_key, mode='RGB'):16    """Apply passed in transforms for HuggingFace Datasets."""17    images = [transform(image.convert(mode)) for image in examples[image_key]]18    return {image_key: images}19 20 21def append_dims(x, target_dims):22    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""23    dims_to_append = target_dims - x.ndim24    if dims_to_append < 0:25        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')26    expanded = x[(...,) + (None,) * dims_to_append]27    # MPS will get inf values if it tries to index into the new axes, but detaching fixes this.28    # https://github.com/pytorch/pytorch/issues/8436429    return expanded.detach().clone() if expanded.device.type == 'mps' else expanded30 31 32def n_params(module):33    """Returns the number of trainable parameters in a module."""34    return sum(p.numel() for p in module.parameters())35 36 37def download_file(path, url, digest=None):38    """Downloads a file if it does not exist, optionally checking its SHA-256 hash."""39    path = Path(path)40    path.parent.mkdir(parents=True, exist_ok=True)41    if not path.exists():42        with urllib.request.urlopen(url) as response, open(path, 'wb') as f:43            shutil.copyfileobj(response, f)44    if digest is not None:45        file_digest = hashlib.sha256(open(path, 'rb').read()).hexdigest()46        if digest != file_digest:47            raise OSError(f'hash of {path} (url: {url}) failed to validate')48    return path49 50 51@contextmanager52def train_mode(model, mode=True):53    """A context manager that places a model into training mode and restores54    the previous mode on exit."""55    modes = [module.training for module in model.modules()]56    try:57        yield model.train(mode)58    finally:59        for i, module in enumerate(model.modules()):60            module.training = modes[i]61 62 63def eval_mode(model):64    """A context manager that places a model into evaluation mode and restores65    the previous mode on exit."""66    return train_mode(model, False)67 68 69@torch.no_grad()70def ema_update(model, averaged_model, decay):71    """Incorporates updated model parameters into an exponential moving averaged72    version of a model. It should be called after each optimizer step."""73    model_params = dict(model.named_parameters())74    averaged_params = dict(averaged_model.named_parameters())75    assert model_params.keys() == averaged_params.keys()76 77    for name, param in model_params.items():78        averaged_params[name].mul_(decay).add_(param, alpha=1 - decay)79 80    model_buffers = dict(model.named_buffers())81    averaged_buffers = dict(averaged_model.named_buffers())82    assert model_buffers.keys() == averaged_buffers.keys()83 84    for name, buf in model_buffers.items():85        averaged_buffers[name].copy_(buf)86 87 88class EMAWarmup:89    """Implements an EMA warmup using an inverse decay schedule.90    If inv_gamma=1 and power=1, implements a simple average. inv_gamma=1, power=2/3 are91    good values for models you plan to train for a million or more steps (reaches decay92    factor 0.999 at 31.6K steps, 0.9999 at 1M steps), inv_gamma=1, power=3/4 for models93    you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at94    215.4k steps).95    Args:96        inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.97        power (float): Exponential factor of EMA warmup. Default: 1.98        min_value (float): The minimum EMA decay rate. Default: 0.99        max_value (float): The maximum EMA decay rate. Default: 1.100        start_at (int): The epoch to start averaging at. Default: 0.101        last_epoch (int): The index of last epoch. Default: 0.102    """103 104    def __init__(self, inv_gamma=1., power=1., min_value=0., max_value=1., start_at=0,105                 last_epoch=0):106        self.inv_gamma = inv_gamma107        self.power = power108        self.min_value = min_value109        self.max_value = max_value110        self.start_at = start_at111        self.last_epoch = last_epoch112 113    def state_dict(self):114        """Returns the state of the class as a :class:`dict`."""115        return dict(self.__dict__.items())116 117    def load_state_dict(self, state_dict):118        """Loads the class's state.119        Args:120            state_dict (dict): scaler state. Should be an object returned121                from a call to :meth:`state_dict`.122        """123        self.__dict__.update(state_dict)124 125    def get_value(self):126        """Gets the current EMA decay rate."""127        epoch = max(0, self.last_epoch - self.start_at)128        value = 1 - (1 + epoch / self.inv_gamma) ** -self.power129        return 0. if epoch < 0 else min(self.max_value, max(self.min_value, value))130 131    def step(self):132        """Updates the step count."""133        self.last_epoch += 1134 135 136class InverseLR(optim.lr_scheduler._LRScheduler):137    """Implements an inverse decay learning rate schedule with an optional exponential138    warmup. When last_epoch=-1, sets initial lr as lr.139    inv_gamma is the number of steps/epochs required for the learning rate to decay to140    (1 / 2)**power of its original value.141    Args:142        optimizer (Optimizer): Wrapped optimizer.143        inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1.144        power (float): Exponential factor of learning rate decay. Default: 1.145        warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)146            Default: 0.147        min_lr (float): The minimum learning rate. Default: 0.148        last_epoch (int): The index of last epoch. Default: -1.149        verbose (bool): If ``True``, prints a message to stdout for150            each update. Default: ``False``.151    """152 153    def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., min_lr=0.,154                 last_epoch=-1, verbose=False):155        self.inv_gamma = inv_gamma156        self.power = power157        if not 0. <= warmup < 1:158            raise ValueError('Invalid value for warmup')159        self.warmup = warmup160        self.min_lr = min_lr161        super().__init__(optimizer, last_epoch, verbose)162 163    def get_lr(self):164        if not self._get_lr_called_within_step:165            warnings.warn("To get the last learning rate computed by the scheduler, "166                          "please use `get_last_lr()`.")167 168        return self._get_closed_form_lr()169 170    def _get_closed_form_lr(self):171        warmup = 1 - self.warmup ** (self.last_epoch + 1)172        lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power173        return [warmup * max(self.min_lr, base_lr * lr_mult)174                for base_lr in self.base_lrs]175 176 177class ExponentialLR(optim.lr_scheduler._LRScheduler):178    """Implements an exponential learning rate schedule with an optional exponential179    warmup. When last_epoch=-1, sets initial lr as lr. Decays the learning rate180    continuously by decay (default 0.5) every num_steps steps.181    Args:182        optimizer (Optimizer): Wrapped optimizer.183        num_steps (float): The number of steps to decay the learning rate by decay in.184        decay (float): The factor by which to decay the learning rate every num_steps185            steps. Default: 0.5.186        warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)187            Default: 0.188        min_lr (float): The minimum learning rate. Default: 0.189        last_epoch (int): The index of last epoch. Default: -1.190        verbose (bool): If ``True``, prints a message to stdout for191            each update. Default: ``False``.192    """193 194    def __init__(self, optimizer, num_steps, decay=0.5, warmup=0., min_lr=0.,195                 last_epoch=-1, verbose=False):196        self.num_steps = num_steps197        self.decay = decay198        if not 0. <= warmup < 1:199            raise ValueError('Invalid value for warmup')200        self.warmup = warmup201        self.min_lr = min_lr202        super().__init__(optimizer, last_epoch, verbose)203 204    def get_lr(self):205        if not self._get_lr_called_within_step:206            warnings.warn("To get the last learning rate computed by the scheduler, "207                          "please use `get_last_lr()`.")208 209        return self._get_closed_form_lr()210 211    def _get_closed_form_lr(self):212        warmup = 1 - self.warmup ** (self.last_epoch + 1)213        lr_mult = (self.decay ** (1 / self.num_steps)) ** self.last_epoch214        return [warmup * max(self.min_lr, base_lr * lr_mult)215                for base_lr in self.base_lrs]216 217 218def rand_log_normal(shape, loc=0., scale=1., device='cpu', dtype=torch.float32):219    """Draws samples from an lognormal distribution."""220    return (torch.randn(shape, device=device, dtype=dtype) * scale + loc).exp()221 222 223def rand_log_logistic(shape, loc=0., scale=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):224    """Draws samples from an optionally truncated log-logistic distribution."""225    min_value = torch.as_tensor(min_value, device=device, dtype=torch.float64)226    max_value = torch.as_tensor(max_value, device=device, dtype=torch.float64)227    min_cdf = min_value.log().sub(loc).div(scale).sigmoid()228    max_cdf = max_value.log().sub(loc).div(scale).sigmoid()229    u = torch.rand(shape, device=device, dtype=torch.float64) * (max_cdf - min_cdf) + min_cdf230    return u.logit().mul(scale).add(loc).exp().to(dtype)231 232 233def rand_log_uniform(shape, min_value, max_value, device='cpu', dtype=torch.float32):234    """Draws samples from an log-uniform distribution."""235    min_value = math.log(min_value)236    max_value = math.log(max_value)237    return (torch.rand(shape, device=device, dtype=dtype) * (max_value - min_value) + min_value).exp()238 239 240def rand_v_diffusion(shape, sigma_data=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):241    """Draws samples from a truncated v-diffusion training timestep distribution."""242    min_cdf = math.atan(min_value / sigma_data) * 2 / math.pi243    max_cdf = math.atan(max_value / sigma_data) * 2 / math.pi244    u = torch.rand(shape, device=device, dtype=dtype) * (max_cdf - min_cdf) + min_cdf245    return torch.tan(u * math.pi / 2) * sigma_data246 247 248def rand_split_log_normal(shape, loc, scale_1, scale_2, device='cpu', dtype=torch.float32):249    """Draws samples from a split lognormal distribution."""250    n = torch.randn(shape, device=device, dtype=dtype).abs()251    u = torch.rand(shape, device=device, dtype=dtype)252    n_left = n * -scale_1 + loc253    n_right = n * scale_2 + loc254    ratio = scale_1 / (scale_1 + scale_2)255    return torch.where(u < ratio, n_left, n_right).exp()256 257 258class FolderOfImages(data.Dataset):259    """Recursively finds all images in a directory. It does not support260    classes/targets."""261 262    IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp'}263 264    def __init__(self, root, transform=None):265        super().__init__()266        self.root = Path(root)267        self.transform = nn.Identity() if transform is None else transform268        self.paths = sorted(path for path in self.root.rglob('*') if path.suffix.lower() in self.IMG_EXTENSIONS)269 270    def __repr__(self):271        return f'FolderOfImages(root="{self.root}", len: {len(self)})'272 273    def __len__(self):274        return len(self.paths)275 276    def __getitem__(self, key):277        path = self.paths[key]278        with open(path, 'rb') as f:279            image = Image.open(f).convert('RGB')280        image = self.transform(image)281        return image,282 283 284class CSVLogger:285    def __init__(self, filename, columns):286        self.filename = Path(filename)287        self.columns = columns288        if self.filename.exists():289            self.file = open(self.filename, 'a')290        else:291            self.file = open(self.filename, 'w')292            self.write(*self.columns)293 294    def write(self, *args):295        print(*args, sep=',', file=self.file, flush=True)296 297 298@contextmanager299def tf32_mode(cudnn=None, matmul=None):300    """A context manager that sets whether TF32 is allowed on cuDNN or matmul."""301    cudnn_old = torch.backends.cudnn.allow_tf32302    matmul_old = torch.backends.cuda.matmul.allow_tf32303    try:304        if cudnn is not None:305            torch.backends.cudnn.allow_tf32 = cudnn306        if matmul is not None:307            torch.backends.cuda.matmul.allow_tf32 = matmul308        yield309    finally:310        if cudnn is not None:311            torch.backends.cudnn.allow_tf32 = cudnn_old312        if matmul is not None:313            torch.backends.cuda.matmul.allow_tf32 = matmul_old314