文本生成三维物体领域,目前基于优化的策略仍然离不开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(