Blame for diffusionGS/models/diffusion/gaussian_diffusion.py

# Modified from OpenAI's diffusion repos
#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py


import math

import numpy as np
import torch as th
import enum

from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl
from pdb import set_trace as stx
import copy
from tqdm import tqdm
import os
from einops import rearrange
from PIL import Image


def mean_flat(tensor):
    """
    Take the mean over all non-batch dimensions.
    """
    return tensor.mean(dim=list(range(1, len(tensor.shape))))


'''
    ModelMeanType 是一个枚举类, 里面有三个枚举值, 可以通过属性的方式来访问, 如 ModelMeanType.PREVIOUS_X
    enum.auto() 会自动分配一个值, 从 1 开始, 依次递增
'''
class ModelMeanType(enum.Enum):
    """
    Which type of output the model predicts.
    """

    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}, 从 x_t 预测 x_{t-1}
    START_X = enum.auto()  # the model predicts x_0
    EPSILON = enum.auto()  # the model predicts epsilon


'''
    Gaussian 的 variance 不应该是 schedule 的吗?
'''
class ModelVarType(enum.Enum):
    """
    What is used as the model's output variance.
    The LEARNED_RANGE option has been added to allow the model to predict
    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
    """

    LEARNED = enum.auto()
    FIXED_SMALL = enum.auto()
    FIXED_LARGE = enum.auto()
    LEARNED_RANGE = enum.auto()


class LossType(enum.Enum):
    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
    RESCALED_MSE = (
        enum.auto()
    )  # use raw MSE loss (with RESCALED_KL when learning variances)
    KL = enum.auto()  # use the variational lower-bound
    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB

    def is_vb(self):
        return self == LossType.KL or self == LossType.RESCALED_KL


# num_diffusion_timesteps = T
# warmup_frac - 指的是 warmup 的比例
# beta 逐步变化，最后到一个 beta_end 后保持不变
def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac):
    betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
    warmup_time = int(num_diffusion_timesteps * warmup_frac)
    betas[:warmup_time] = np.linspace(
        beta_start, beta_end, warmup_time, dtype=np.float64
    )
    return betas



# * 之前位置参数, * 之后必须是关键字参数
# 不同的 beta_schedule 策略
def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
    """
    This is the deprecated API for creating beta schedules.
    See get_named_beta_schedule() for the new library of schedules.
    """
    if beta_schedule == "quad":
        betas = (
            np.linspace(
                beta_start**0.5,
                beta_end**0.5,
                num_diffusion_timesteps,
                dtype=np.float64,
            )
            ** 2
        )
    elif beta_schedule == "linear":
        betas = np.linspace(
            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
        )
    elif beta_schedule == "warmup10":
        betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1)
    elif beta_schedule == "warmup50":
        betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5)
    elif beta_schedule == "const":
        betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
    elif beta_schedule == "jsd":  # 1/T, 1/(T-1), 1/(T-2), ..., 1
        betas = 1.0 / np.linspace(
            num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64
        )
    else:
        raise NotImplementedError(beta_schedule)
    assert betas.shape == (num_diffusion_timesteps,)
    return betas


# 另一些之前被命名过的 beta scheduler
def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
    """
    Get a pre-defined beta schedule for the given name.
    The beta schedule library consists of beta schedules which remain similar
    in the limit of num_diffusion_timesteps.
    Beta schedules may be added, but should not be removed or changed once
    they are committed to maintain backwards compatibility.
    """
    if schedule_name == "linear":
        # Linear schedule from Ho et al, extended to work for any number of
        # diffusion steps.
        scale = 1000 / num_diffusion_timesteps
        return get_beta_schedule(
            "linear",
            beta_start=scale * 0.0001,
            beta_end=scale * 0.02,
            num_diffusion_timesteps=num_diffusion_timesteps,
        )
    elif schedule_name == "squaredcos_cap_v2":
        return betas_for_alpha_bar(
            num_diffusion_timesteps,
            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
        )
    else:
        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")


# 还是一个 scheduler
def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
    """
    Create a beta schedule that discretizes the given alpha_t_bar function,
    which defines the cumulative product of (1-beta) over time from t = [0,1].
    :param num_diffusion_timesteps: the number of betas to produce.
    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
                      produces the cumulative product of (1-beta) up to that
                      part of the diffusion process.
    :param max_beta: the maximum beta to use; use values lower than 1 to
                     prevent singularities.
    """
    betas = []
    for i in range(num_diffusion_timesteps):
        t1 = i / num_diffusion_timesteps
        t2 = (i + 1) / num_diffusion_timesteps
        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
    return np.array(betas)


# 注意，这只是个 general 的框架, model 是灵活的
'''
    如果想要 x_0 output 的话，好像就只用在调用 GaussianDiffusion 时的时候，设置 model_mean_type = ModelMeanType.START_X 即可
    因为 x_{t-1} 一直是拿 x_t 和 x_0 来算的, q(x_{t-1} | x_t, x_0)
'''
class GaussianDiffusion:
    """
    Utilities for training and sampling diffusion models.
    Original ported from this codebase:
    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
    :param betas: a 1-D numpy array of betas for each diffusion timestep,
                  starting at T and going to 1.
    """

    def __init__(self, *, betas, model_mean_type, model_var_type, loss_type):

        self.model_mean_type = model_mean_type
        self.model_var_type = model_var_type
        self.loss_type = loss_type

        # Use float64 for accuracy.
        '''
            beta 控制的是前向传播过程中逐步所加的噪声的方差
            q(x_t | x_{t-1}) = N(x_t; sqrt(1-beta_t) x_{t-1}, beta_t)
            alpha_t = 1 - beta_t
        '''
        betas = np.array(betas, dtype=np.float64)
        self.betas = betas
        assert len(betas.shape) == 1, "betas must be 1-D, [B,]"
        assert (betas > 0).all() and (betas <= 1).all()

        self.num_timesteps = int(betas.shape[0])

        alphas = 1.0 - betas
        # cumulative product of alphas over time
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        # 在 self.alphas_cumprod 数组的开始处添加了一个元素 1.0，并且移除了最后一个元素
        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
        # 在 self.alphas_cumprod 数组的末尾添加了一个元素 0.0，并且移除了第一个元素
        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)

        # calculations for diffusion q(x_t | x_{t-1}) and others
        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)         # reciprocal - 倒数
        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        self.posterior_variance = (
            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
        # 起始值为 0 不可计算, 只可计算后续的值
        self.posterior_log_variance_clipped = (
            np.log(np.append(self.posterior_variance[1], self.posterior_variance[1:]))
            if len(self.posterior_variance) > 1
            else np.array([])
        )

        '''
            下边两个是用来计算 q(x_{t-1} | x_t, x_0) 均值的两个系数
            coef1 是 x_0 的系数, coef2 是 x_t 的系数
            然后 x_0 是由 x_t 和 predict 的 噪声 eps 线性组合得到的, 是模型这个阶段猜出的 x_0, 并不是真正的 x_0
        '''
        self.posterior_mean_coef1 = (
            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )

        self.posterior_mean_coef2 = (
            (1.0 - self.alphas_cumprod_prev)
            * np.sqrt(alphas)
            / (1.0 - self.alphas_cumprod)
        )

    '''
        q(x_t | x_0) = N(x_t; sqrt(alpha_t_bar) x_0, sqrt(1 - alpha_t_bar))
    '''
    def q_mean_variance(self, x_start, t):
        """
        Get the distribution q(x_t | x_0).
        :param x_start: the [N x C x ...] tensor of noiseless inputs. 真实的 x_0, 非模型所猜
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
        """
        mean = (
            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start   # sqrt(alpha_t_bar) x_0
        )
        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)    # sqrt(1 - alpha_t_bar)
        log_variance = _extract_into_tensor(
            self.log_one_minus_alphas_cumprod, t, x_start.shape
        )
        return mean, variance, log_variance

    # x_0, t -> x_t
    # q(x_t | x_0) = N(x_t; sqrt(alpha_t_bar) x_0, sqrt(1 - alpha_t_bar))
    # 注意：此处不是重参数化采样 ！！ 因为此处不需要求导了, x_t 是U-Net 的输入
    # 重参数化采样发生在 q(x_{t-1} | x_t, x_0), 因为此时就需要求导，loss在 eps 上计算，直接关联 x_{t-1} 的 μ 和 σ, 需要梯度回传
    def q_sample(self, x_start, t, noise=None):
        """
        Diffuse the data for a given number of diffusion steps.
        In other words, sample from q(x_t | x_0).
        :param x_start: the initial data batch.
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :param noise: if specified, the split-out normal noise.
        :return: A noisy version of x_start.
        """
        if noise is None:
            noise = th.randn_like(x_start)
        assert noise.shape == x_start.shape
        return (
            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
            * noise
        )


    # prior - 啥也没观测到时，事件A的概率 p(A)
    # posterior - 观测到一些数据或者结果 B 后, A事件发生的概率 p(A|B)
    # 此处的 x_start 是模型猜的 x_0
    # 由于此时 U-Net 已经完成了 eps 的预测，所以这里没有 model
    def q_posterior_mean_variance(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        # stx()
        assert x_start.shape == x_t.shape
        posterior_mean = (
            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
        )
        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
        posterior_log_variance_clipped = _extract_into_tensor(
            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


    # 此处调用 model
    def p_mean_variance(self, model, input_batch, t, clip_denoised=True, model_kwargs=None):
        """
        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
        the initial x, x_0.
        :param model: the model, which takes a signal and a batch of timesteps
                      as input.
        :param x: the [N x C x ...] tensor at time t.
        :param t: a 1-D Tensor of timesteps.
        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.
        :return: a dict with the following keys:
                 - 'mean': the model mean output.
                 - 'variance': the model variance output.
                 - 'log_variance': the log of 'variance'.
                 - 'pred_xstart': the prediction for x_0.
        """
        if model_kwargs is None:
            model_kwargs = {}

        x = input_batch["image_noisy"]

        B, V, C = x.shape[:3]
        assert t.shape == (B,)
        '''
            input_batch: dict_keys(['image', 'c2w', 'fxfycxcy', 'index', 'depth', 'normal', 'image_noisy'])
            image: [8, 8, 3, 256, 256]
            image_noisy: [8, 3, 3, 256, 256]
            model_output: [8, 3, 3, 256, 256]
            这里存一组图片来看一下
        '''
        # stx()
        if t[0] > 0:
            input_batch['image'] = th.cat([input_batch['image'][:,0:1],input_batch['image_noisy']],dim=1) # select first images and update input
            render_imgs, pred_guassians = model(input_batch, t)
            #render_imgs, pred_guassians = model(**input_batch)
            model_output = render_imgs[:,1:]
            output_dict = {
                'render_images':render_imgs,
                'pred_gaussians':pred_guassians,
            }
        else:
            input_batch['image'] = th.cat([input_batch['image'][:,0:1],input_batch['image_noisy']],dim=1) # select first images and update input
            render_imgs, pred_guassians = model(input_batch, t)
            model_output = render_imgs[:,1:]
            output_dict = {
                'render_images':render_imgs,
                'pred_gaussians':pred_guassians,
            }
        # variance 的两种方式，一是模型直接输出
        # 二是预先设定好的 variance (beta) 直接按照 timestep 取值
        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
            assert model_output.shape == (B, C * 2, *x.shape[2:])
            model_output, model_var_values = th.split(model_output, C, dim=1)
            min_log = _extract_into_tensor(
                self.posterior_log_variance_clipped, t, x.shape
            )
            max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
            # The model_var_values is [-1, 1] for [min_var, max_var].
            frac = (model_var_values + 1) / 2
            model_log_variance = frac * max_log + (1 - frac) * min_log
            model_variance = th.exp(model_log_variance)
        else:
            model_variance, model_log_variance = {
                # for fixedlarge, we set the initial (log-)variance like so
                # to get a better decoder log likelihood.
                ModelVarType.FIXED_LARGE: (
                    np.append(self.posterior_variance[1], self.betas[1:]),
                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
                ),
                ModelVarType.FIXED_SMALL: (
                    self.posterior_variance,
                    self.posterior_log_variance_clipped,
                ),
            }[self.model_var_type]
            model_variance = _extract_into_tensor(model_variance, t, x.shape)
            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)


        # 为什么要 clamp 到 [-1, 1]?
        def process_xstart(x):
            if clip_denoised:
                return x.clamp(-1, 1)
            return x

        if self.model_mean_type == ModelMeanType.START_X:
            # stx()
            pred_xstart = process_xstart(model_output)
        else:
            pred_xstart = process_xstart(
                self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
            )
        
        # q(x_{t-1} | x_t, x_0), 这里的 x_0 是模型猜的
        model_mean, _, _ = self.q_posterior_mean_variance(
            x_start=pred_xstart, x_t=x, t=t
        )


        ######################################################################################################################
        ###########################################    sampling 存图 debug     ###############################################
        ######################################################################################################################

        # sampling_debug_path = os.path.join("/sensei-fs/users/yuanhaoc/gslrm_dit/minLRM/debug_sampling", "gaussian_diffusion")
        # os.makedirs(sampling_debug_path, exist_ok=True)

        # im = model_output
        # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
        # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
        # Image.fromarray(im).save(os.path.join(sampling_debug_path, "render_input.png"))

        # im = input_batch["image_noisy"]
        # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
        # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
        # Image.fromarray(im).save(os.path.join(sampling_debug_path, "xt_noisy.png"))

        # im = model_mean
        # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
        # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
        # Image.fromarray(im).save(os.path.join(sampling_debug_path, "model_mean_next_t.png"))

        # im = model_variance
        # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
        # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
        # Image.fromarray(im).save(os.path.join(sampling_debug_path, "model_variance_next_t.png"))

        # im = pred_xstart
        # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
        # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
        # Image.fromarray(im).save(os.path.join(sampling_debug_path, "pred_xstart_after_clipping.png"))

        # stx()


        assert (
            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
        )
        return {
            "mean": model_mean,                     # μθ of x_{t-1}
            "variance": model_variance,             # σ^2
            "log_variance": model_log_variance,     # log(σ^2)
            "pred_xstart": pred_xstart,             # x_0
            "denoiser_output_dict": output_dict,
        }

    # 从 eps 预测 x_0, 正常走的是这一步
    # x_0 = 1/sqrt(alpha_t_bar) x_t - sqrt((1-alpha_t_bar) / alpha_t_bar) eps
    def _predict_xstart_from_eps(self, x_t, t, eps):
        assert x_t.shape == eps.shape
        return (
            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
        )

    # 由 x_0 预测 eps, 这里假设模型输出 x_0
    # eps = (x_t - sqrt(alpha_t_bar) x_0) / sqrt(1-alpha_t_bar)
    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
        return (
            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - pred_xstart
        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)

    # predict x_{t-1}
    def p_sample(
        self,
        model,
        input_batch,
        t,
        clip_denoised=True,
        model_kwargs=None,
    ):
        """
        Sample x_{t-1} from the model at the given timestep.
        :param model: the model to sample from.
        :param x: the current tensor at x_t.
        :param t: the value of t, starting at 0 for the first diffusion step.
        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.
        :return: a dict containing the following keys:
                 - 'sample': a random sample from the model.
                 - 'pred_xstart': a prediction of x_0.
        """
        # stx()
        x = input_batch["image_noisy"]
        out = self.p_mean_variance(
            model,
            input_batch,
            t,
            clip_denoised=clip_denoised,
            model_kwargs=model_kwargs,
        )

        noise = th.randn_like(x)
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        # 此处是重参数化采样 x_{t-1} = μ + σ * eps
        # [1] 为何要 nonzero_mask, 因为底 0 步可以直接输出 x_0, 不需要加噪声
        # [2] 为何要有一个 torch.exp(0.5), 给方差开根号就是协方差
        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
        input_batch["image_noisy"] = sample
        return {"sample": sample, "pred_xstart": out["pred_xstart"], "input_batch": input_batch, "denoiser_output_dict": out["denoiser_output_dict"]} # 返回 x_{t-1}, x_0, 和 input_batch

    def p_sample_loop(
        self,
        model,
        shape,
        input_batch=None,
        clip_denoised=True,
        model_kwargs=None,
        device=None,
        progress=True,
    ):
        """
        Generate samples from the model.
        :param model: the model module.
        :param shape: the shape of the samples, (N, C, H, W).
        :param noise: if specified, the noise from the encoder to sample.
                      Should be of the same shape as `shape`.
        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
        :param model_kwargs: if not None, a dict of extra keyword arguments to
            pass to the model. This can be used for conditioning.
        :param device: if specified, the device to create the samples on.
                       If not specified, use a model parameter's device.
        :param progress: if True, show a tqdm progress bar.
        :return: a non-differentiable batch of samples.
        """
        final = None
        # 对 self.p_sample_loop_progressive 函数返回的对象进行 forloop
        # stx()
        for sample in self.p_sample_loop_progressive(
            model,
            shape,
            input_batch=input_batch,
            clip_denoised=clip_denoised,
            model_kwargs=model_kwargs,
            device=device,
            progress=progress,
        ):
            final = sample          # 此处对应的 sample 应该是 p_sample 函数的 return 值，但这个函数好像就只需要里面的 "sample"
        return final      # 这个 "sample" 对应 x_{t-1}, 即每一步的状态, 到最后一步会把 x_0 输出

    # 生成一系列的 x_{t-1}，从 t = T 逐渐到 t = 0
    def p_sample_loop_progressive(
        self,
        model,
        shape,
        input_batch=None,
        clip_denoised=True,
        model_kwargs=None,
        device=None,
        progress=False,
    ):
        """
        Generate samples from the model and yield intermediate samples from
        each timestep of diffusion.
        Arguments are the same as p_sample_loop().
        Returns a generator over dicts, where each dict is the return value of
        p_sample().
        """
        if device is None:
            device = next(model.parameters()).device
        assert isinstance(shape, (tuple, list))
        indices = list(range(self.num_timesteps))[::-1]     # 表示的是 t 的序号，indices_t 不一定等于 t

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm  # 自动根据环境选择进度条

            indices = tqdm(indices)

        # stx()
        for i in indices:
            t = th.tensor([i] * shape[0], device=device)
            with th.no_grad():
                # stx()
                output = self.p_sample(
                    model,
                    input_batch,                            # img = x_t, p_sample 会返回 x_{t-1}. img 的初始值是 noise
                    t,
                    clip_denoised=clip_denoised,
                    model_kwargs=model_kwargs,
                )
                yield output
                # 修改 input_batch 里面的迭代字典
                # 应该在 p_sample 里面修改 input_batch 的关键字, 然后这边就直接传就好 
                input_batch = output["input_batch"]

    def _vb_terms_bpd(
        self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
    ):
        """
        Get a term for the variational lower-bound.
        The resulting units are bits (rather than nats, as one might expect).
        This allows for comparison to other papers.
        :return: a dict with the following keys:
                 - 'output': a shape [N] tensor of NLLs or KLs.
                 - 'pred_xstart': the x_0 predictions.
        """
        # 计算 q(x_{t-1} | x_t, x_0) - 此处的 x_0 是真实值
        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
            x_start=x_start, x_t=x_t, t=t
        )

        # 调用 model 计算 p(x_{t-1} | x_t) - 此处的 x_0 是模型猜的
        out = self.p_mean_variance(
            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
        )

        # 计算 KL(q(x_{t-1} | x_t, x_0) || p(x_{t-1} | x_t)), true 和 out 的 KL 散度
        kl = normal_kl(
            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
        )
        kl = mean_flat(kl) / np.log(2.0)

        # 第一步输出对数似然，后面输出 KL 散度
        decoder_nll = -discretized_gaussian_log_likelihood(
            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
        )
        assert decoder_nll.shape == x_start.shape
        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)

        # At the first timestep return the decoder NLL,
        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
        output = th.where((t == 0), decoder_nll, kl)
        return {"output": output, "pred_xstart": out["pred_xstart"]}

    def training_losses(self, model, batch, t, model_kwargs=None, noise=None, mode='training', create_visual = False, x0_loss_weight = 1.0):
        """
        被调用形式: loss_metrics = self.diffusion.training_losses(self.denoiser, batch, t)
        Compute training losses for a single timestep.
        :param model: the model to evaluate loss on.
        :param batch: data batch - 一个 dict: dict_keys(['image', 'c2w', 'fxfycxcy', 'index', 'depth', 'normal'])
        :param t: a batch of timestep indices.

        :return: a dict with the key "loss" containing a tensor of shape [N].
                 Some mean or variance settings may also have other keys.
                 返回一个 loss 的 dict
        """
        # 取 batch 的第二个 view 开始作为 x_0
        # 这里的 batch 还是一整个的 batch
        # (b, v, c, h, w)
        img_condition = batch["image"][:,0:1, ...]          # [8, 1, 3, 256, 256]
        x_start = batch["image"][:,1:4, ...].clone()                 # [8, 3, 3, 256, 256]

        if model_kwargs is None:
            model_kwargs = {}
        if noise is None:
            noise = th.randn_like(x_start)                  # [8, 3, 3, 256, 256]

        # # x_0, t -> x_t
        # q(x_t | x_0) = N(x_t; sqrt(alpha_t_bar) x_0, sqrt(1 - alpha_t_bar))
        # x_t: (b, v-1, c, h, w)
        x_t = self.q_sample(x_start, t, noise=noise)


        batch['image_noisy'] = x_t          # 这个值不会引起别的改变, 因为是新开辟的 key

        # stx()

        terms = {}

        # 计算 KL(q(x_{t-1} | x_t, x_0) || p(x_{t-1} | x_t)) 散度 Loss, KL 散度 Loss 用得其实不多
        # 这里的 if 逻辑有点奇怪，就是说如果是 KL 散度 Loss 的话，就直接返回 KL 散度 Loss
        # 如果是 MSE Loss 的话，不仅要算 MSE Loss，还要算 KL 散度 Loss，返回的 loss 是两者之和
        # 此处是 False, 不执行
        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
            terms["loss"] = self._vb_terms_bpd(
                model=model,
                x_start=x_start,
                x_t=x_t,
                t=t,
                clip_denoised=False,
                model_kwargs=model_kwargs,
            )["output"]
            if self.loss_type == LossType.RESCALED_KL:
                terms["loss"] *= self.num_timesteps

        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:

            denoiser_output = model(batch, t, **model_kwargs, mode = mode, create_visual=create_visual)   # 过 denoiser
            model_output = denoiser_output['render_input']
           
           #######################################################################################################
           #############################  可视化 debug - 计算 loss 前把所有的图片存下来进行调试  #######################
           #######################################################################################################

            # debug_path = os.path.join("/sensei-fs/users/yuanhaoc/gslrm_dit/minLRM/debug", "gaussian_diffusion")
            # os.makedirs(debug_path, exist_ok=True)

            # im = model_output
            # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im).save(os.path.join(debug_path, "model_render_input.png"))

            # im = denoiser_output['render']
            # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im).save(os.path.join(debug_path, "model_render_target.png"))

            # im = denoiser_output['input']['image']
            # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im).save(os.path.join(debug_path, "input_image.png"))

            # im = denoiser_output['target']['image']
            # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im).save(os.path.join(debug_path, "target_image.png"))

            # im_x_start = x_start
            # im_x_start = rearrange(im_x_start, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im_x_start = (im_x_start * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im_x_start).save(os.path.join(debug_path, "x_start.png"))

            # im = x_t
            # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im).save(os.path.join(debug_path, "x_t.png"))


            # combo_image_condition = th.cat([img_condition, x_t], dim=1)
            # im = combo_image_condition
            # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im).save(os.path.join(debug_path, "combo_input_image.png"))


            # stx()

            ################################################################################################################
            ############################################    可视化 debug 结束    #############################################
            ################################################################################################################


            '''
                model_output 理论上也是一个字典
                dict_keys(['input', 'target', 'gaussians', 'pixelalign_xyz', 'img_tokens', 'loss_metrics', 'render'])
                ——> dict_keys(['input', 'target', 'gaussians', 'pixelalign_xyz', 'img_tokens', 'loss_metrics', 'render', 'render_input'])
                要从中抽出 multi-view 的 x0 然后变成 x_{t-1}
            '''

            # learn_sigma = True, 所以模型不仅要 predict mean, 还要 predict variance
            # 我已经设置成 False 了
            if self.model_var_type in [
                ModelVarType.LEARNED,
                ModelVarType.LEARNED_RANGE,
            ]:
                B, C = x_t.shape[:2]
                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
                model_output, model_var_values = th.split(model_output, C, dim=1)
                # Learn the variance using the variational bound, but don't let
                # it affect our mean prediction.
                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
                terms["vb_loss"] = self._vb_terms_bpd(
                    model=lambda *args, r=frozen_out: r,
                    x_start=x_start,
                    x_t=x_t,
                    t=t,
                    clip_denoised=False,
                )["output"]
                if self.loss_type == LossType.RESCALED_MSE:
                    # Divide by 1000 for equivalence with initial implementation.
                    # Without a factor of 1/1000, the VB term hurts the MSE term.
                    terms["vb_loss"] *= self.num_timesteps / 1000.0

            target = {
                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
                    x_start=x_start, x_t=x_t, t=t
                )[0],                                                       # x_{t-1}
                ModelMeanType.START_X: x_start,                             # x_0
                ModelMeanType.EPSILON: noise,                               # eps
            }[self.model_mean_type]  # mean type 决定了在什么上边算 MSE loss
            # stx()
            assert model_output.shape == target.shape == x_start.shape



            #######################################################################################################
            ###################################  x0_loss 会导致掉点 - 再多渲染一次看看  ###############################
            #######################################################################################################

            # debug_path = os.path.join("/sensei-fs/users/yuanhaoc/gslrm_dit/minLRM/debug_2object", "gaussian_diffusion")
            # os.makedirs(debug_path, exist_ok=True)

            # im = target
            # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im).save(os.path.join(debug_path, "input_gt.png"))

            # im = model_output
            # im = rearrange(im, "b v c h w -> (b h) (v w) c").detach().cpu().numpy()
            # im = (im * 255).clip(0.0, 255.0).astype(np.uint8)
            # Image.fromarray(im).save(os.path.join(debug_path, "render_input.png"))

            # stx()

            #######################################################################################################
            ##################################################  调试结束  ##########################################
            #######################################################################################################


            eps_loss = False
            x0_loss = False
            if self.model_mean_type == ModelMeanType.START_X and eps_loss:
                predict_eps = self._predict_eps_from_xstart(x_t=x_t, t=t, pred_xstart=model_output)
                terms["mse_x0_loss"] = mean_flat((noise - predict_eps) ** 2)
            elif self.model_mean_type == ModelMeanType.EPSILON and x0_loss:
                predict_x0 = self._predict_xstart_from_eps(x_t=x_t, t=t, eps=model_output)
                terms["mse_x0_loss"] = mean_flat((x_start - predict_x0) ** 2)
            else:
                terms["mse_x0_loss"] = mean_flat((target - model_output) ** 2) * x0_loss_weight
            

            # 前面的 predict variance 已经被设置成 False
            # terms 字典里没有 vb
            output_loss_metrics = denoiser_output['loss_metrics']
            if "vb_loss" in terms:
                terms["loss"] = terms["mse_x0_loss"] + terms["vb_loss"]      
            else:
                terms["loss"] = terms["mse_x0_loss"]
            
            # 遍历 output_loss_metrics 的键值对
            for key, value in output_loss_metrics.items():
                if key == 'loss':
                    terms[key] = terms[key] + value     # term 里面的 loss 是全部的 loss 之和
                else:
                    terms[key] = value
        else:
            raise NotImplementedError(self.loss_type)

        return terms, denoiser_output


# 从一个一维的 numpy 数组中提取出一批索引对应的值，并将这些值转换为 PyTorch 的 Tensor 格式
# 按照选取的 timestep 去取预先计算好的数值
def _extract_into_tensor(arr, timesteps, broadcast_shape):
    """
    Extract values from a 1-D numpy array for a batch of indices.
    :param arr: the 1-D numpy array.
    :param timesteps: a tensor of indices into the array to extract.
    :param broadcast_shape: a larger shape of K dimensions with the batch
                            dimension equal to the length of timesteps.
    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
    """
    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
    while len(res.shape) < len(broadcast_shape):
        res = res[..., None]
    return res + th.zeros(broadcast_shape, device=timesteps.device)