详细解读DreamFusion

文本生成三维物体领域，目前基于优化的策略仍然离不开DreamFusion提出的SDS Loss，这也是DreamFusion中最核心的部分。承接，这里就以经典的Stable DreamStable DreamfusionStable Dream为例子，在这篇文章中简单地聊聊SDS的具体代码实现。

首先，让我们回顾一下SDS的基本作用和运作机理。对于一个2D的扩散模型，比如Stable Diffusion，我们想利用其2D的信息，监督三维物体的生成，这个时候就会用到SDS loss。SDS不是直接用Diffusion模型生成的最终图片做三维的监督，而是利用diffusion生成图片的过程提取知识。所以我们必须对NeRF渲染出的结果先用高斯噪声加噪，再用diffusion模型去噪。这样“加噪—去噪”的过程中，我们就可以用上Diffusion生成图片时候的先验知识了。

因此，对于一个NeRF渲染出来的结果[B, 3, H, W]（这里的B是batch size，3代表RGB这三个颜色通道，H和W代表宽和高），我们给它赋个值，以便大家有个直观的认识。假定[B, 3, H, W]具体的形状是[1, 3, 64, 64]。这就代表我们用NeRF在每个step里渲染出一张分辨率为64*64的图片，用来给diffusion监督。

一，加噪

第一步，我们先将这张64*64的小图片插值成一张512*512的图片，这是为了迎合Stable Diffusion的分辨率。这是因为Stable Diffusion 1.5的训练分辨率就是512*512，所以当然是插值到这个分辨率下加噪去噪最好：

def train_step(self, text_embeddings, pred_rgb, guidance_scale=100):
          # 输入--pred_rgb：[1,3,64,64]
          # 将输入插值到512*512，用来喂给stable diffusion的vae
          pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False)
          # 输出--pred_rgb_512: [1,3,512,512]

这是有人可能会心里有一些疑惑：既然最后要插值到512*512，那为什么不一开始我们就用NeRF渲染一张512*512的图片来作为输入呢？这其实是受限于早期对diffusion和NeRF的优化并没有那么好，所以要渲染高分辨率的图片会非常的慢，而且对显存的消耗很严重，因此大家普遍采用先渲染低分辨率图片再插值的策略。

然后，我们进行“加噪—去噪”中的加噪过程。这里和一个普通的Stable Diffusion生成图片过程的加噪一样，整个过程需要在隐空间(latent space)上进行。因此，我们利用SD预训练好的vae将[1,3,512,512]的图片encode到隐空间上，隐空间上图片的形状就转变为[1,4,64,64]。这里的4就是隐空间上抽象的向量维度了，和原来的RGB已经不再是一个东西了[1]。

           # 输入--pred_rgb_512: [1,3,512,512]
           latents = self.encode_imgs(pred_rgb_512)
           # 输出--latents: [1,4,64,64]

至此，我们就得到了一个Stable Diffusion的隐向量(latents)，可供我们在“加噪—去噪”的过程中提取出扩散模型对于这个世界的认识。然后，我们正式给这个隐向量(latents)加上噪声，具体操作如下：

        # predict the noise residual with unet, NO grad!
        with torch.no_grad(