Doubiiu/DynamiCrafter_interp_loop

1import math2from inspect import isfunction3import torch4from torch import nn5import torch.distributed as dist6 7 8def gather_data(data, return_np=True):9    ''' gather data from multiple processes to one list '''10    data_list = [torch.zeros_like(data) for _ in range(dist.get_world_size())]11    dist.all_gather(data_list, data)  # gather not supported with NCCL12    if return_np:13        data_list = [data.cpu().numpy() for data in data_list]14    return data_list15 16def autocast(f):17    def do_autocast(*args, **kwargs):18        with torch.cuda.amp.autocast(enabled=True,19                                     dtype=torch.get_autocast_gpu_dtype(),20                                     cache_enabled=torch.is_autocast_cache_enabled()):21            return f(*args, **kwargs)22    return do_autocast23 24 25def extract_into_tensor(a, t, x_shape):26    b, *_ = t.shape27    out = a.gather(-1, t)28    return out.reshape(b, *((1,) * (len(x_shape) - 1)))29 30 31def noise_like(shape, device, repeat=False):32    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))33    noise = lambda: torch.randn(shape, device=device)34    return repeat_noise() if repeat else noise()35 36 37def default(val, d):38    if exists(val):39        return val40    return d() if isfunction(d) else d41 42def exists(val):43    return val is not None44 45def identity(*args, **kwargs):46    return nn.Identity()47 48def uniq(arr):49    return{el: True for el in arr}.keys()50 51def mean_flat(tensor):52    """53    Take the mean over all non-batch dimensions.54    """55    return tensor.mean(dim=list(range(1, len(tensor.shape))))56 57def ismap(x):58    if not isinstance(x, torch.Tensor):59        return False60    return (len(x.shape) == 4) and (x.shape[1] > 3)61 62def isimage(x):63    if not isinstance(x,torch.Tensor):64        return False65    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)66 67def max_neg_value(t):68    return -torch.finfo(t.dtype).max69 70def shape_to_str(x):71    shape_str = "x".join([str(x) for x in x.shape])72    return shape_str73 74def init_(tensor):75    dim = tensor.shape[-1]76    std = 1 / math.sqrt(dim)77    tensor.uniform_(-std, std)78    return tensor79 80ckpt = torch.utils.checkpoint.checkpoint81def checkpoint(func, inputs, params, flag):82    """83    Evaluate a function without caching intermediate activations, allowing for84    reduced memory at the expense of extra compute in the backward pass.85    :param func: the function to evaluate.86    :param inputs: the argument sequence to pass to `func`.87    :param params: a sequence of parameters `func` depends on but does not88                   explicitly take as arguments.89    :param flag: if False, disable gradient checkpointing.90    """91    if flag:92        return ckpt(func, *inputs, use_reentrant=False)93    else:94        return func(*inputs)