Doubiiu/DynamiCrafter_interp_loop

1import math2import numpy as np3import torch4import torch.nn.functional as F5from einops import repeat6 7 8def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):9    """10    Create sinusoidal timestep embeddings.11    :param timesteps: a 1-D Tensor of N indices, one per batch element.12                      These may be fractional.13    :param dim: the dimension of the output.14    :param max_period: controls the minimum frequency of the embeddings.15    :return: an [N x dim] Tensor of positional embeddings.16    """17    if not repeat_only:18        half = dim // 219        freqs = torch.exp(20            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half21        ).to(device=timesteps.device)22        args = timesteps[:, None].float() * freqs[None]23        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)24        if dim % 2:25            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)26    else:27        embedding = repeat(timesteps, 'b -> b d', d=dim)28    return embedding29 30 31def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):32    if schedule == "linear":33        betas = (34                torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 235        )36 37    elif schedule == "cosine":38        timesteps = (39                torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s40        )41        alphas = timesteps / (1 + cosine_s) * np.pi / 242        alphas = torch.cos(alphas).pow(2)43        alphas = alphas / alphas[0]44        betas = 1 - alphas[1:] / alphas[:-1]45        betas = np.clip(betas, a_min=0, a_max=0.999)46 47    elif schedule == "sqrt_linear":48        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)49    elif schedule == "sqrt":50        betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.551    else:52        raise ValueError(f"schedule '{schedule}' unknown.")53    return betas.numpy()54 55 56def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):57    if ddim_discr_method == 'uniform':58        c = num_ddpm_timesteps // num_ddim_timesteps59        ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))60        steps_out = ddim_timesteps + 161    elif ddim_discr_method == 'uniform_trailing':62        c = num_ddpm_timesteps / num_ddim_timesteps63        ddim_timesteps = np.flip(np.round(np.arange(num_ddpm_timesteps, 0, -c))).astype(np.int64)64        steps_out = ddim_timesteps - 165    elif ddim_discr_method == 'quad':66        ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)67        steps_out = ddim_timesteps + 168    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    # print(f'ddim_timesteps={ddim_timesteps}, len_alphacums={len(alphacums)}')82    alphas = alphacums[ddim_timesteps]83    alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())84 85    # according the the formula provided in https://arxiv.org/abs/2010.0250286    sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))87    if verbose:88        print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')89        print(f'For the chosen value of eta, which is {eta}, '90              f'this results in the following sigma_t schedule for ddim sampler {sigmas}')91    return sigmas, alphas, alphas_prev92 93 94def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):95    """96    Create a beta schedule that discretizes the given alpha_t_bar function,97    which defines the cumulative product of (1-beta) over time from t = [0,1].98    :param num_diffusion_timesteps: the number of betas to produce.99    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and100                      produces the cumulative product of (1-beta) up to that101                      part of the diffusion process.102    :param max_beta: the maximum beta to use; use values lower than 1 to103                     prevent singularities.104    """105    betas = []106    for i in range(num_diffusion_timesteps):107        t1 = i / num_diffusion_timesteps108        t2 = (i + 1) / num_diffusion_timesteps109        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))110    return np.array(betas)111 112def rescale_zero_terminal_snr(betas):113    """114    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)115 116    Args:117        betas (`numpy.ndarray`):118            the betas that the scheduler is being initialized with.119 120    Returns:121        `numpy.ndarray`: rescaled betas with zero terminal SNR122    """123    # Convert betas to alphas_bar_sqrt124    alphas = 1.0 - betas125    alphas_cumprod = np.cumprod(alphas, axis=0)126    alphas_bar_sqrt = np.sqrt(alphas_cumprod)127 128    # Store old values.129    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].copy()130    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].copy()131 132    # Shift so the last timestep is zero.133    alphas_bar_sqrt -= alphas_bar_sqrt_T134 135    # Scale so the first timestep is back to the old value.136    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)137 138    # Convert alphas_bar_sqrt to betas139    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt140    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod141    alphas = np.concatenate([alphas_bar[0:1], alphas])142    betas = 1 - alphas143 144    return betas145 146 147def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):148    """149    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and150    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4151    """152    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)153    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)154    # rescale the results from guidance (fixes overexposure)155    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)156    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images157    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg158    return noise_cfg