1、Denoising Diffusion Probabilistic Models（扩散模型）_measurement-conditioned denoising diffusion probab

简介

主页：https://hojonathanho.github.io/diffusion/

扩散模型 （diffusion models）是深度生成模型中新的SOTA。

扩散模型在图片生成任务中超越了原SOTA：GAN，并且在诸多应用领域都有出色的表现，如计算机视觉，NLP、波形信号处理、多模态建模、分子图建模、时间序列建模、对抗性净化等。

GAN要训练两个网络，训练难度大，容易不收敛，而且多样性比较差，毕竟生成器是为了骗过鉴别器，生成器可能学到稀奇古怪的技巧，

此外，扩散模型与其他研究领域有着密切的联系，如稳健学习、表示学习、强化学习。

然而，原始的扩散模型也有缺点，它的采样速度慢，通常需要数千个评估步骤才能抽取一个样本；它的最大似然估计无法和基于似然的模型相比；它泛化到各种数据类型的能力较差。

如今很多研究已经从实际应用的角度解决上述限制做出了许多努力，或从理论角度对模型能力进行了分析。但是，现在仍缺乏对扩散模型从算法到应用的最新进展的系统回顾。

基础知识

实现流程

生成式建模的一个核心问题是模型的灵活性和可计算性之间的权衡。

扩散模型的基本思想是正向扩散过程来系统地扰动数据中的分布，然后通过学习反向扩散过程恢复数据的分布，这样就了产生一个高度灵活且易于计算的生成模型。

前向过程

前向过程概括起来就是从原始图像 $X_0$ 开始，不断往图像中加入高斯噪声，每一个时刻由前一时刻的图像增加噪声得到，最后得到纯噪声的图像。这个过程可以看作是不断构建标签（高斯噪声）的过程。

构建 $X_t$ 时刻是公式如下：

$X_t=\sqrt{{\alpha }_t}{X}_{t-1}+\sqrt{1-{\alpha }_t}Z$

其中 ${\alpha }_t=1-{\beta }_t$

$\beta$ 随着时刻 t 增大而增加，论文从0.0001 增加到 0.002。那么 $\alpha$ 随着时刻 t 增大而减少，这表明了后一时刻的图像对前一时刻的图像的依赖逐渐减少，高斯噪声的权重逐渐增大，最后得到纯噪声的图像

从 $X_0$ 开始往后迭代可以得到每一时刻的分布，但是这个过程过于繁琐且消耗大量计算资源，可不可以根据 $X_0$ 获取任意指定时刻的分布呢？

1. 步骤一

   首先，时刻 t 的图像记为 $X_t$，前一时刻 t-1 的图像记为 $X_{t-1}$，这里使用 Z 表示高斯分布

   已知，$X_t=\sqrt{{\alpha }_t}{X}_{t-1}+\sqrt{1-{\alpha }_t}Z$

   将 $X_{t-1}=\sqrt{{\alpha }_{t-1}}{X}_{t-2}+\sqrt{1-{\alpha }_{t-1}}Z$ 代入上述公式，取代 $X_{t-1}$，得到

   $X_t=\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{X}_{t-2}+\sqrt{1-{\alpha }_{t-1}}Z)+\sqrt{1-{\alpha }_t}Z$

   化简得

   $X_t=\sqrt{{\alpha }_t{\alpha }_{t-1}}{X}_{t-2}+\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}Z+\sqrt{1-{\alpha }_t}Z$

2. 步骤二

   高斯分布 $Z\sim \mathrm{N}(0,I)$

   $\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}Z\sim \mathrm{N}(0,{\alpha }_t-{\alpha }_t{\alpha }_{t-1})$

   $\sqrt{1-{\alpha }_t}Z\sim \mathrm{N}(0,1-{\alpha }_t)$

   由于高斯分布符合以下规律

   $\mathrm{N}(0,{\sigma }_1^{2}I)+\mathrm{N}(0,{\sigma }_2^{2}I)\sim \mathrm{N}(0,({\sigma }_1^2+{\sigma }_2^2)I)$

   所以

   $\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}Z+\sqrt{1-{\alpha }_t}Z\sim \mathrm{N}(0,1-{\alpha }_t{\alpha }_{t-1})$

   从步骤一得到的公式：

   $X_t=\sqrt{{\alpha }_t{\alpha }_{t-1}}{X}_{t-2}+\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}Z+\sqrt{1-{\alpha }_t}Z$

   化简可得

   $X_t=\sqrt{{\alpha }_t{\alpha }_{t-1}}{X}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}Z$

   从而可以推出

   $X_t=\sqrt{\bar{\alpha }}{X}_0+\sqrt{1-\bar{\alpha }}Z$，（$\bar{\alpha }$ 表示连乘）

我们现在可以实现加噪声的过程了，但是目的是去噪生成，也就是接下来的逆向过程

逆向过程

那么我们回到我们的初始目的，如何从 $T_N$ 时刻分布 $X_t$ 一步一步往前推得到生成目标图像 $X_0$ 呢？

回到原始公式

$X_t=\sqrt{{\alpha }_t}{X}_{t-1}+\sqrt{1-{\alpha }_t}Z$

那么我们要如何使用 $X_t$ 表示 $X_{t-1}$呢

1. 步骤一

   这里使用贝叶斯公式

   $q(X_{t-1}|X_t)=q(X_t|X_{t-1})\frac{q(X_{t-1})}{q(X_t)}$

   在前向过程，任意时刻 t 的分布 $X_t$ 可以由 $X_0$ 表示

   $X_t=\sqrt{\bar{\alpha }}{X}_0+\sqrt{1-\bar{\alpha }}Z$，（$\bar{\alpha }$ 表示连乘）

   那么套用贝叶斯的原始公式可以使用初始条件 $X_0$ 表示

   $q(X_{t-1}|X_t,X_0)=q(X_t|X_{t-1},X_0)\frac{q(X_{t-1}|X_0)}{q(X_t|X_0)}$

   右边三项未知数可以表示为：

   $q(X_{t-1}|X_0):\sqrt{{\bar{\alpha }}_{t-1}}{X}_0+\sqrt{1-{\bar{\alpha }}_{t-1}}Z\sim \mathrm{N}(\sqrt{{\bar{\alpha }}_{t-1}}{X}_0,1-{\bar{\alpha }}_{t-1})$

   $q(X_t|X_0):\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}Z\sim \mathrm{N}(\sqrt{{\bar{\alpha }}_t}{X}_0,1-{\bar{\alpha }}_t)$

   $q(X_t|X_{t-1},X_0):\sqrt{{\alpha }_t}{X}_{t-1}+\sqrt{1-{\alpha }_t}Z\sim \mathrm{N}(\sqrt{{\alpha }_t}{X}_{t-1},1-{\alpha }_t)$

   将上面三条公式带入贝叶斯公式

   $q(X_{t-1}|X_t,X_0)=q(X_t|X_{t-1},X_0)\frac{q(X_{t-1}|X_0)}{q(X_t|X_0)}$

   我们知道高斯分布 $Z=e^{-\frac{1}{2}\frac{(x-\mu {)}^2}{{\sigma }^2}}$

   化简得到

   $X_{t-1}=e^{(-\frac{1}{2}(\frac{(x_t-\sqrt{{\alpha }_t}{X}_{t-1}{)}^2}{{\beta }_t}+\frac{(X_{t-1}-\sqrt{{\bar{\alpha }}_{t-1}}{X}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}-\frac{(X_t-\sqrt{{\bar{\alpha }}_t}{X}_0{)}^2}{1-{\bar{\alpha }}_t}))}$

2. 步骤二

   将步骤一的 $X_{t-1}$ 表达式展开后，汇总化简得到

   $e^{(-\frac{1}{2}((\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})X_{t-1}^2-(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{X}_t+\frac{2\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{X}_0)X_{t-1}+C(X_t,X_0)))}$

   $C(X_t,X_0)$ 为常数项，不影响任务，核心是求 $X_t$ 与 $X_{t-1}$ 的关系。

   将高斯分布(Z) 展开后为

   $Z=e^{(-\frac{(x-\mu {)}^2}{2{\sigma }^2})}=e^{(-\frac{1}{2}(\frac{1}{{\sigma }^2}{X}^2-\frac{2\mu }{{\sigma }^2}X+\frac{{\mu }^2}{{\sigma }^2}))}$

   对比 高斯分布(Z) 展开后公式 与 上述得到的 $X_{t-1}$ 表达式，可以得到 均值 和 方差
   $\frac{1}{{\sigma }^2}=(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})$

   $\sigma =\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$

   $\tilde{\mu }(X_t,X_0)=\frac{{\sqrt{\alpha }}_t(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{X}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{X}_0$

   其中 $X_0$ 未知，但是我们知道 $X_t$ 可以由 $X_0$ 得到，那么将原公式逆过来

   $X_t=\sqrt{\bar{\alpha }}{X}_0+\sqrt{1-\bar{\alpha }}Z$，（$\bar{\alpha }$ 表示连乘）

   $X_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(X_t-\sqrt{1-{\bar{\alpha }}_t}Z)$

   再将 $X_0$ 带入均值表达式，化简得

   ${\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}(X_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}Z)$

3. 步骤三

   每一时刻的 $X_t$ 都是一个高斯分布，因此，可以通过高斯分布重采样策略得到 $X_{t-1}$

   我们现在得到了有样本 X 得到的分布 $X\sim N(\mu ,{\sigma }^2)$。采样这个操作本身是不可导的，但是我们可以通过重参数化技巧，将简单分布的采样结果变换到特定分布中，如此一来则可以对变换过程进行求导。具体而言，我们从标准高斯分布中采样，并将其变换到$X\sim N(\mu ,{\sigma }^2)$，过程如下

   $\epsilon \sim \mathrm{N}(0,I)$
   $Z=\mu +\sigma \times \epsilon$

   也就是说，从 $\mathrm{N}(\mu ,{\sigma }^2)$ 采样 $Z$，等同于从 $\epsilon \sim \mathrm{N}(0,I)$ 中采样高斯噪声 $\epsilon$，再将其按 $Z=\mu +\sigma \times \epsilon$ 变换

   $X_{t-1}={\tilde{\mu }}_t+{\sigma }_{t}Z\sim \mathrm{N}({\tilde{\mu }}_t,{\sigma }_t)$

目标数据分布的似然函数

对于两个单一变量的高斯分布 p 和 q而言，KL散度为$KL(p,q)=\log \frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}-\frac{1}{2}$

伪代码

总体网络可以采用了简单的U-net实现

Training

目标：让网络预测不同时刻的高斯分布 ${\epsilon }_{\theta }$

首先从数据集中随机采样图像 $X_0$，选取超参数时刻上限 $T$，在$1,...,T$ 中随机采样时刻（batch size）并为此生成时刻对应的高斯分布 $\epsilon$，根据公式

$X_t=\sqrt{\bar{\alpha }}{X}_0+\sqrt{1-\bar{\alpha }}Z$，（$\bar{\alpha }$ 表示连乘）

将 t 时刻的分布 $X_t$ 和时刻 t 输入网络，其中时刻 t 经过位置编码后与 $X_t$ 拼接，网络预测得到时刻 t 的高斯分布 ${\epsilon }_{\theta }$，将其与对应时刻的高斯分布 $\epsilon$ 作 $L_2$ 损失

Sampling

分布 $X_t$ 由高斯分布给出，进行 $T$ 次循环，从模型${\epsilon }_{\theta }(X_t,t)$中获取时刻 t 的高斯分布预测值 ${\epsilon }_{\theta }$，通过公式：

$X_{t-1}={\tilde{\mu }}_t+{\sigma }_{t}Z\sim \mathrm{N}({\tilde{\mu }}_t,{\sigma }_t)$

预测前一时刻的分布 $X_{t-1}$，循环该过程得到最终图像 $X_0$

扩散模型代码

class GaussianDiffusion:
  """
  Contains utilities for the diffusion model.
  """

  def __init__(self, *, betas, loss_type, tf_dtype=tf.float32):
    self.loss_type = loss_type

    assert isinstance(betas, np.ndarray)
    self.np_betas = betas = betas.astype(np.float64)  # computations here in float64 for accuracy
    assert (betas > 0).all() and (betas <= 1).all()
    timesteps, = betas.shape
    self.num_timesteps = int(timesteps)

    alphas = 1. - betas
    alphas_cumprod = np.cumprod(alphas, axis=0)
    alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
    assert alphas_cumprod_prev.shape == (timesteps,)

    self.betas = tf.constant(betas, dtype=tf_dtype)
    self.alphas_cumprod = tf.constant(alphas_cumprod, dtype=tf_dtype)
    self.alphas_cumprod_prev = tf.constant(alphas_cumprod_prev, dtype=tf_dtype)

    # calculations for diffusion q(x_t | x_{t-1}) and others
    self.sqrt_alphas_cumprod = tf.constant(np.sqrt(alphas_cumprod), dtype=tf_dtype)
    self.sqrt_one_minus_alphas_cumprod = tf.constant(np.sqrt(1. - alphas_cumprod), dtype=tf_dtype)
    self.log_one_minus_alphas_cumprod = tf.constant(np.log(1. - alphas_cumprod), dtype=tf_dtype)
    self.sqrt_recip_alphas_cumprod = tf.constant(np.sqrt(1. / alphas_cumprod), dtype=tf_dtype)
    self.sqrt_recipm1_alphas_cumprod = tf.constant(np.sqrt(1. / alphas_cumprod - 1), dtype=tf_dtype)

    # -------------- 方差

    # calculations for posterior q(x_{t-1} | x_t, x_0)
    posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
    # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
    self.posterior_variance = tf.constant(posterior_variance, dtype=tf_dtype)
    # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain

    # sqrt_alphas  sqrt_0.0001 一开始很小 为了避免为0，取对数
    self.posterior_log_variance_clipped = tf.constant(np.log(np.maximum(posterior_variance, 1e-20)), dtype=tf_dtype)

    # -------------- 均值
    # 知道 X_0 的情况下计算 均值 后部分
    self.posterior_mean_coef1 = tf.constant(
      betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod), dtype=tf_dtype)
    # 知道 X_0 的情况下计算 均值 前部分
    self.posterior_mean_coef2 = tf.constant(
      (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod), dtype=tf_dtype)

  @staticmethod
  def _extract(a, t, x_shape):
    """
    从 a 中按t抽取元素
    Extract some coefficients at specified timesteps,
    then reshape to [batch_size, 1, 1, 1, 1, ...] for broadcasting purposes.
    """
    bs, = t.shape
    assert x_shape[0] == bs
    out = tf.gather(a, t)
    assert out.shape == [bs]
    return tf.reshape(out, [bs] + ((len(x_shape) - 1) * [1]))

#  ======================= 加噪过程 q(z | x)  =======================

  def q_mean_variance(self, x_start, t):
    """
    加噪过程 计算 X_t 均值和方差
    """
    mean = self._extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
    variance = self._extract(1. - self.alphas_cumprod, t, x_start.shape)
    log_variance = self._extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
    return mean, variance, log_variance

  def q_sample(self, x_start, t, noise=None):
    """
    Diffuse the data (t == 0 means diffused for 1 step)
    根据 X_0,t 计算 X_t 时刻分布
    """
    if noise is None:
      noise = tf.random_normal(shape=x_start.shape)
    assert noise.shape == x_start.shape
    return (
        self._extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
        self._extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
    )


#  ======================= 模型损失、验证  =======================


  def p_losses(self, denoise_fn, x_start, t, noise=None):
    """
    Training loss calculation
    训练过程，模型 预测噪声 与 标注正态分布
    """
    B, H, W, C = x_start.shape.as_list()
    assert t.shape == [B]

    if noise is None:
      noise = tf.random_normal(shape=x_start.shape, dtype=x_start.dtype)
    assert noise.shape == x_start.shape and noise.dtype == x_start.dtype
    x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
    x_recon = denoise_fn(x_noisy, t)
    assert x_noisy.shape == x_start.shape
    assert x_recon.shape[:3] == [B, H, W] and len(x_recon.shape) == 4

    if self.loss_type == 'noisepred':
      # predict the noise instead of x_start. seems to be weighted naturally like SNR
      assert x_recon.shape == x_start.shape
      losses = nn.meanflat(tf.squared_difference(noise, x_recon))
    else:
      raise NotImplementedError(self.loss_type)

    assert losses.shape == [B]
    return losses

#  ======================= 去噪过程 p(z | x)  =======================

  def q_posterior(self, x_start, x_t, t):
    """
    Compute the mean and variance of the diffusion posterior q(x_{t-1} | x_t, x_0)
    根据 x_0 x_t t 计算方差 和 均值
    """
    assert x_start.shape == x_t.shape
    posterior_mean = (
        self._extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
        self._extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
    )
    posterior_variance = self._extract(self.posterior_variance, t, x_t.shape)
    posterior_log_variance_clipped = self._extract(self.posterior_log_variance_clipped, t, x_t.shape)
    assert (posterior_mean.shape[0] == posterior_variance.shape[0] == posterior_log_variance_clipped.shape[0] ==
            x_start.shape[0])
    return posterior_mean, posterior_variance, posterior_log_variance_clipped

  def predict_start_from_noise(self, x_t, t, noise):
    """
    根据 X_t t 计算 X_0
    """
    assert x_t.shape == noise.shape
    return (
        self._extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
        self._extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
    )

  def p_mean_variance(self, denoise_fn, *, x, t, clip_denoised: bool):
    """
      去噪过程获取 均值 和 方差
    """
    if self.loss_type == 'noisepred':
      x_recon = self.predict_start_from_noise(x, t=t, noise=denoise_fn(x, t))
    else:
      raise NotImplementedError(self.loss_type)

    if clip_denoised:
      x_recon = tf.clip_by_value(x_recon, -1., 1.)

    model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
    assert model_mean.shape == x_recon.shape == x.shape
    assert posterior_variance.shape == posterior_log_variance.shape == [x.shape[0], 1, 1, 1]
    return model_mean, posterior_variance, posterior_log_variance

  def p_sample(self, denoise_fn, *, x, t, noise_fn, clip_denoised=True, repeat_noise=False):
    """
    Sample from the model
    计算 X_t
    """
    model_mean, _, model_log_variance = self.p_mean_variance(denoise_fn, x=x, t=t, clip_denoised=clip_denoised)
    noise = noise_like(x.shape, noise_fn, repeat_noise)
    assert noise.shape == x.shape
    # no noise when t == 0
    nonzero_mask = tf.reshape(1 - tf.cast(tf.equal(t, 0), tf.float32), [x.shape[0]] + [1] * (len(x.shape) - 1))

    # tf.exp(0.5 * model_log_variance) 去掉 log
    return model_mean + nonzero_mask * tf.exp(0.5 * model_log_variance) * noise

  def p_sample_loop(self, denoise_fn, *, shape, noise_fn=tf.random_normal):
    """
    Generate samples
    循环 去噪过程
    """
    i_0 = tf.constant(self.num_timesteps - 1, dtype=tf.int32)
    assert isinstance(shape, (tuple, list))
    img_0 = noise_fn(shape=shape, dtype=tf.float32)
    _, img_final = tf.while_loop(
      cond=lambda i_, _: tf.greater_equal(i_, 0),
      body=lambda i_, img_: [
        i_ - 1,
        self.p_sample(denoise_fn=denoise_fn, x=img_, t=tf.fill([shape[0]], i_), noise_fn=noise_fn)
      ],
      loop_vars=[i_0, img_0],
      shape_invariants=[i_0.shape, img_0.shape],
      back_prop=False
    )
    assert img_final.shape == shape
    return img_final

#  ======================= 验证过程  =======================

  def p_sample_loop_trajectory(self, denoise_fn, *, shape, noise_fn=tf.random_normal, repeat_noise_steps=-1):
    """
    生成样本，返回中间图像用于可视化去噪图像如何随时间演变


    Generate samples, returning intermediate images
    Useful for visualizing how denoised images evolve over time
    Args:
      repeat_noise_steps (int): Number of denoising timesteps in which the same noise
        is used across the batch. If >= 0, the initial noise is the same for all batch elemements.
    """
    i_0 = tf.constant(self.num_timesteps - 1, dtype=tf.int32)
    assert isinstance(shape, (tuple, list))
    img_0 = noise_like(shape, noise_fn, repeat_noise_steps >= 0)
    times = tf.Variable([i_0])
    imgs = tf.Variable([img_0])

    # 相同噪声的步骤
    # Steps with repeated noise
    times, imgs = tf.while_loop(
      cond=lambda times_, _: tf.less_equal(self.num_timesteps - times_[-1], repeat_noise_steps),
      body=lambda times_, imgs_: [
        tf.concat([times_, [times_[-1] - 1]], 0),
        tf.concat([imgs_, [self.p_sample(denoise_fn=denoise_fn,
                                         x=imgs_[-1],
                                         t=tf.fill([shape[0]], times_[-1]),
                                         noise_fn=noise_fn,
                                         repeat_noise=True)]], 0)
      ],
      loop_vars=[times, imgs],
      shape_invariants=[tf.TensorShape([None, *i_0.shape]),
                        tf.TensorShape([None, *img_0.shape])],
      back_prop=False
    )

    # 每个批处理元素具有不同噪声的步骤
    # Steps with different noise for each batch element
    times, imgs = tf.while_loop(
      cond=lambda times_, _: tf.greater_equal(times_[-1], 0),
      body=lambda times_, imgs_: [
        tf.concat([times_, [times_[-1] - 1]], 0),
        tf.concat([imgs_, [self.p_sample(denoise_fn=denoise_fn,
                                         x=imgs_[-1],
                                         t=tf.fill([shape[0]], times_[-1]),
                                         noise_fn=noise_fn,
                                         repeat_noise=False)]], 0)
      ],
      loop_vars=[times, imgs],
      shape_invariants=[tf.TensorShape([None, *i_0.shape]),
                        tf.TensorShape([None, *img_0.shape])],
      back_prop=False
    )
    assert imgs[-1].shape == shape
    return times, imgs

  def interpolate(self, denoise_fn, *, shape, noise_fn=tf.random_normal):
    """
    在图像之间插入。
    Interpolate between images.
    t == 0 means diffuse images for 1 timestep before mixing.
    """
    assert isinstance(shape, (tuple, list))

    # 占位符，用于插入真实样本
    # Placeholders for real samples to interpolate
    x1 = tf.placeholder(tf.float32, shape)
    x2 = tf.placeholder(tf.float32, shape)
    # lam == 0.5 averages diffused images.
    lam = tf.placeholder(tf.float32, shape=())
    t = tf.placeholder(tf.int32, shape=())

    # 通过前向扩散添加噪声
    # Add noise via forward diffusion
    # TODO: use the same noise for both endpoints?
    # t_batched = tf.constant([t] * x1.shape[0], dtype=tf.int32)
    t_batched = tf.stack([t] * x1.shape[0])
    xt1 = self.q_sample(x1, t=t_batched)
    xt2 = self.q_sample(x2, t=t_batched)

    # Mix latents
    # Linear interpolation
    xt_interp = (1 - lam) * xt1 + lam * xt2
    # Constant variance interpolation
    # xt_interp = tf.sqrt(1 - lam * lam) * xt1 + lam * xt2

    # Reverse diffusion (similar to self.p_sample_loop)
    # t = tf.constant(t, dtype=tf.int32)
    _, x_interp = tf.while_loop(
      cond=lambda i_, _: tf.greater_equal(i_, 0),
      body=lambda i_, img_: [
        i_ - 1,
        self.p_sample(denoise_fn=denoise_fn, x=img_, t=tf.fill([shape[0]], i_), noise_fn=noise_fn)
      ],
      loop_vars=[t, xt_interp],
      shape_invariants=[t.shape, xt_interp.shape],
      back_prop=False
    )
    assert x_interp.shape == shape

    return x1, x2, lam, x_interp, t

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299