xdecoder/Instruct-X-Decoder

1# adopted from2# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py3# and4# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py5# and6# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py7#8# thanks!9import importlib10 11import os12import math13import torch14import torch.nn as nn15import numpy as np16from einops import repeat17 18 19def instantiate_from_config(config):20    if not "target" in config:21        if config == '__is_first_stage__':22            return None23        elif config == "__is_unconditional__":24            return None25        raise KeyError("Expected key `target` to instantiate.")26    return get_obj_from_str(config["target"])(**config.get("params", dict()))27 28 29def get_obj_from_str(string, reload=False):30    module, cls = string.rsplit(".", 1)31    if reload:32        module_imp = importlib.import_module(module)33        importlib.reload(module_imp)34    return getattr(importlib.import_module(module, package=None), cls)35 36 37def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):38    if schedule == "linear":39        betas = (40                torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 241        )42 43    elif schedule == "cosine":44        timesteps = (45                torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s46        )47        alphas = timesteps / (1 + cosine_s) * np.pi / 248        alphas = torch.cos(alphas).pow(2)49        alphas = alphas / alphas[0]50        betas = 1 - alphas[1:] / alphas[:-1]51        betas = np.clip(betas, a_min=0, a_max=0.999)52 53    elif schedule == "sqrt_linear":54        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)55    elif schedule == "sqrt":56        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.557    else:58        raise ValueError(f"schedule '{schedule}' unknown.")59    return betas.numpy()60 61 62def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):63    if ddim_discr_method == 'uniform':64        c = num_ddpm_timesteps // num_ddim_timesteps65        ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))66    elif ddim_discr_method == 'quad':67        ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)68    else:69        raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')70 71    # assert ddim_timesteps.shape[0] == num_ddim_timesteps72    # add one to get the final alpha values right (the ones from first scale to data during sampling)73    steps_out = ddim_timesteps + 174    if verbose:75        print(f'Selected timesteps for ddim sampler: {steps_out}')76    return steps_out77 78 79def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):80    # select alphas for computing the variance schedule81    alphas = alphacums[ddim_timesteps]82    alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())83 84    # according the the formula provided in https://arxiv.org/abs/2010.0250285    sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))86    if verbose:87        print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')88        print(f'For the chosen value of eta, which is {eta}, '89              f'this results in the following sigma_t schedule for ddim sampler {sigmas}')90    return sigmas, alphas, alphas_prev91 92 93def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):94    """95    Create a beta schedule that discretizes the given alpha_t_bar function,96    which defines the cumulative product of (1-beta) over time from t = [0,1].97    :param num_diffusion_timesteps: the number of betas to produce.98    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and99                      produces the cumulative product of (1-beta) up to that100                      part of the diffusion process.101    :param max_beta: the maximum beta to use; use values lower than 1 to102                     prevent singularities.103    """104    betas = []105    for i in range(num_diffusion_timesteps):106        t1 = i / num_diffusion_timesteps107        t2 = (i + 1) / num_diffusion_timesteps108        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))109    return np.array(betas)110 111 112def extract_into_tensor(a, t, x_shape):113    b, *_ = t.shape114    out = a.gather(-1, t)115    return out.reshape(b, *((1,) * (len(x_shape) - 1)))116 117 118def checkpoint(func, inputs, params, flag):119    """120    Evaluate a function without caching intermediate activations, allowing for121    reduced memory at the expense of extra compute in the backward pass.122    :param func: the function to evaluate.123    :param inputs: the argument sequence to pass to `func`.124    :param params: a sequence of parameters `func` depends on but does not125                   explicitly take as arguments.126    :param flag: if False, disable gradient checkpointing.127    """128    if flag:129        args = tuple(inputs) + tuple(params)130        return CheckpointFunction.apply(func, len(inputs), *args)131    else:132        return func(*inputs)133 134 135class CheckpointFunction(torch.autograd.Function):136    @staticmethod137    def forward(ctx, run_function, length, *args):138        ctx.run_function = run_function139        ctx.input_tensors = list(args[:length])140        ctx.input_params = list(args[length:])141 142        with torch.no_grad():143            output_tensors = ctx.run_function(*ctx.input_tensors)144        return output_tensors145 146    @staticmethod147    def backward(ctx, *output_grads):148        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]149        with torch.enable_grad():150            # Fixes a bug where the first op in run_function modifies the151            # Tensor storage in place, which is not allowed for detach()'d152            # Tensors.153            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]154            output_tensors = ctx.run_function(*shallow_copies)155        input_grads = torch.autograd.grad(156            output_tensors,157            ctx.input_tensors + ctx.input_params,158            output_grads,159            allow_unused=True,160        )161        del ctx.input_tensors162        del ctx.input_params163        del output_tensors164        return (None, None) + input_grads165 166 167def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):168    """169    Create sinusoidal timestep embeddings.170    :param timesteps: a 1-D Tensor of N indices, one per batch element.171                      These may be fractional.172    :param dim: the dimension of the output.173    :param max_period: controls the minimum frequency of the embeddings.174    :return: an [N x dim] Tensor of positional embeddings.175    """176    if not repeat_only:177        half = dim // 2178        freqs = torch.exp(179            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half180        ).to(device=timesteps.device)181        args = timesteps[:, None].float() * freqs[None]182        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)183        if dim % 2:184            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)185    else:186        embedding = repeat(timesteps, 'b -> b d', d=dim)187    return embedding188 189 190def zero_module(module):191    """192    Zero out the parameters of a module and return it.193    """194    for p in module.parameters():195        p.detach().zero_()196    return module197 198 199def scale_module(module, scale):200    """201    Scale the parameters of a module and return it.202    """203    for p in module.parameters():204        p.detach().mul_(scale)205    return module206 207 208def mean_flat(tensor):209    """210    Take the mean over all non-batch dimensions.211    """212    return tensor.mean(dim=list(range(1, len(tensor.shape))))213 214 215def normalization(channels):216    """217    Make a standard normalization layer.218    :param channels: number of input channels.219    :return: an nn.Module for normalization.220    """221    return GroupNorm32(32, channels)222 223 224# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.225class SiLU(nn.Module):226    def forward(self, x):227        return x * torch.sigmoid(x)228 229 230class GroupNorm32(nn.GroupNorm):231    def forward(self, x):232        return super().forward(x.float()).type(x.dtype)233 234def conv_nd(dims, *args, **kwargs):235    """236    Create a 1D, 2D, or 3D convolution module.237    """238    if dims == 1:239        return nn.Conv1d(*args, **kwargs)240    elif dims == 2:241        return nn.Conv2d(*args, **kwargs)242    elif dims == 3:243        return nn.Conv3d(*args, **kwargs)244    raise ValueError(f"unsupported dimensions: {dims}")245 246 247def linear(*args, **kwargs):248    """249    Create a linear module.250    """251    return nn.Linear(*args, **kwargs)252 253 254def avg_pool_nd(dims, *args, **kwargs):255    """256    Create a 1D, 2D, or 3D average pooling module.257    """258    if dims == 1:259        return nn.AvgPool1d(*args, **kwargs)260    elif dims == 2:261        return nn.AvgPool2d(*args, **kwargs)262    elif dims == 3:263        return nn.AvgPool3d(*args, **kwargs)264    raise ValueError(f"unsupported dimensions: {dims}")265 266 267class HybridConditioner(nn.Module):268 269    def __init__(self, c_concat_config, c_crossattn_config):270        super().__init__()271        self.concat_conditioner = instantiate_from_config(c_concat_config)272        self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)273 274    def forward(self, c_concat, c_crossattn):275        c_concat = self.concat_conditioner(c_concat)276        c_crossattn = self.crossattn_conditioner(c_crossattn)277        return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}278 279 280def noise_like(shape, device, repeat=False):281    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))282    noise = lambda: torch.randn(shape, device=device)283    return repeat_noise() if repeat else noise()