Anonymous-123/ImageNet-Editing

1"""2Helpers for various likelihood-based losses. These are ported from the original3Ho et al. diffusion models codebase:4https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py5"""6 7import numpy as np8 9import torch as th10 11 12def normal_kl(mean1, logvar1, mean2, logvar2):13    """14    Compute the KL divergence between two gaussians.15 16    Shapes are automatically broadcasted, so batches can be compared to17    scalars, among other use cases.18    """19    tensor = None20    for obj in (mean1, logvar1, mean2, logvar2):21        if isinstance(obj, th.Tensor):22            tensor = obj23            break24    assert tensor is not None, "at least one argument must be a Tensor"25 26    # Force variances to be Tensors. Broadcasting helps convert scalars to27    # Tensors, but it does not work for th.exp().28    logvar1, logvar2 = [29        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)30        for x in (logvar1, logvar2)31    ]32 33    return 0.5 * (34        -1.035        + logvar236        - logvar137        + th.exp(logvar1 - logvar2)38        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)39    )40 41 42def approx_standard_normal_cdf(x):43    """44    A fast approximation of the cumulative distribution function of the45    standard normal.46    """47    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))48 49 50def discretized_gaussian_log_likelihood(x, *, means, log_scales):51    """52    Compute the log-likelihood of a Gaussian distribution discretizing to a53    given image.54 55    :param x: the target images. It is assumed that this was uint8 values,56              rescaled to the range [-1, 1].57    :param means: the Gaussian mean Tensor.58    :param log_scales: the Gaussian log stddev Tensor.59    :return: a tensor like x of log probabilities (in nats).60    """61    assert x.shape == means.shape == log_scales.shape62    centered_x = x - means63    inv_stdv = th.exp(-log_scales)64    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)65    cdf_plus = approx_standard_normal_cdf(plus_in)66    min_in = inv_stdv * (centered_x - 1.0 / 255.0)67    cdf_min = approx_standard_normal_cdf(min_in)68    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))69    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))70    cdf_delta = cdf_plus - cdf_min71    log_probs = th.where(72        x < -0.999,73        log_cdf_plus,74        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),75    )76    assert log_probs.shape == x.shape77    return log_probs78