分数雅可比链：将预训练的2D扩散模型提升用于3D生成

Haochen Wang*， Xiaodan Du*， Jiahao Li*， Raymond A. Yeh†， Greg Shakhnarovich (* 表示共同第一作者)

TTI-芝加哥，†普渡大学

摘要：扩散模型学习预测梯度场。我们提出对学习到的梯度应用链式法则，并通过可微渲染器的雅可比矩阵反向传播扩散模型的分数，这里我们将该渲染器实例化为体素辐射场。这种设置可以将多个相机视角下的2D分数聚合为3D分数，从而重新利用预训练的2D模型进行3D数据生成。我们识别出在这种应用中出现的分布不匹配这一技术挑战，并提出了一种新颖的估计机制来解决它。我们在几款现成的扩散图像生成模型上运行了我们的算法，其中包括最近发布的、基于大规模LAION数据集训练的Stable Diffusion。

非常感谢dvschultz提供的Colab笔记本，以及AmanKishore提供的Hugging Face演示。

SJC现在也已集成到threestudio中。

更新

• 我们添加了用于最终高质量可视化的亚像素渲染脚本。您之前可能看到的抖动视频现在应该会好很多。请在训练完成后，在实验文件夹中运行python /path/to/sjc/highres_final_vis.py。脚本中有一些可调节的选项，但默认设置即可。在A5000显卡上大约需要5分钟和11GB显存，额外的时间主要来自SD解码器。
• 如果您使用DreamBooth微调后的模型运行SJC：该模型的输出分布已经显著收窄。此时可以尝试使用较低的引导尺度，例如--sd.scale 50.0。强烈的模式寻找是导致多脸问题的原因之一。我们内部尝试过结合视图相关提示进行DreamBooth微调，但总体而言，DreamBooth的集成尚未成熟。

待办事项

• 使种子可配置。目前所有种子都硬编码为0。
• 添加脚本以复现图4中的2D实验。一旦与种子绑定，这张图可能需要调整。需要注意的是，对于人脸这样简单的对齐领域，仅使用单一σ=1.5的简单调度就能生成一些不错的图像。但对于卧室场景则不然；其多样性太大，似乎仍然需要退火过程。
• 主论文中的图片未使用亚像素渲染；附录中的图片则使用了。应替换主论文中的图片，使其保持一致。

许可证

由于我们使用了Stable Diffusion，因此我们根据其OpenRAIL许可证发布。除此之外，我们并未发现任何组件或上游代码带有限制性许可要求。

结构

除了SJC之外，该仓库还包含Karras采样器的实现，以及一个定制的简单体素NeRF。我们提供了基于Karras等人工作的抽象父类，并在此包含了若干种扩散模型的实现。详情请参见adapt.py。

安装

根据您的CUDA版本安装PyTorch，例如：

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116

使用pip install -r requirements.txt安装其他依赖项。

手动安装taming-transformers：

git clone --depth 1 git@github.com:CompVis/taming-transformers.git && pip install -e taming-transformers

下载检查点

我们已将您需要下载的最小集合（SD v1.5检查点、LSUN和FFHQ的gddpm检查点）打包在一个tar文件中，并在我们的下载服务器这里提供下载。这是一个12GB的单个文件，您可以使用wget或curl下载。

请务必__更新__env.json，使其指向您解压文件后的新检查点根目录。

使用方法

创建一个新的目录来运行实验（该脚本会生成大量日志文件。请勿在代码仓库的根目录下运行，以免污染。）

mkdir exp
cd exp

运行以下命令以生成一个新的3D资产。在单个A5000显卡上，优化10000步大约需要25分钟和10GB显存。

python /path/to/sjc/run_sjc.py \
--sd.prompt "天坛的高分辨率远景照片" \
--n_steps 10000 \
--lr 0.05 \
--sd.scale 100.0 \
--emptiness_weight 10000 \
--emptiness_step 0.5 \
--emptiness_multiplier 20.0 \
--depth_weight 0 \
--var_red False

sd.prompt是传递给稳定扩散模型的提示词

n_steps是梯度迭代次数

lr是优化器的基础学习率

sd.scale是稳定扩散的引导尺度

emptiness_weight是空洞损失的权重因子

emptiness_step表示在经过emptiness_step * n_steps次更新后，emptiness_weight将乘以emptiness_multiplier。

emptiness_multipler如上所述

depth_weight是中心深度损失的权重因子

var_red表示是否使用公式16还是公式15。对于某些提示词，比如奥巴马，我们实际上发现使用公式15的效果更好。

可视化结果存储在当前目录中。在名为test_*的目录下，分别存放不同迭代步骤生成的图像（位于view文件夹）和视频（位于view_seq文件夹）。

复现论文中的结果

首先为您的实验创建一个干净的目录，然后从该目录运行以下脚本之一：

特朗普

python /path/to/sjc/run_sjc.py --sd.prompt "特朗普雕像" --n_steps 30000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

奥巴马

python /path/to/sjc/run_sjc.py --sd.prompt "奥巴马人物" --n_steps 30000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0 --var_red False

拜登

python /path/to/sjc/run_sjc.py --sd.prompt "拜登人物" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

天坛

python /path/to/sjc/run_sjc.py --sd.prompt "天坛的高画质远景照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

汉堡

python /path/to/sjc/run_sjc.py --sd.prompt "美味汉堡的高画质照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

冰淇淋

python /path/to/sjc/run_sjc.py --sd.prompt "巧克力冰淇淋甜筒的高画质照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 10

榕树

python /path/to/sjc/run_sjc.py --sd.prompt "盆栽榕树" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 100

城堡

python /path/to/sjc/run_sjc.py --sd.prompt "小型城堡的远景照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 50

悉尼歌剧院

python /path/to/sjc/run_sjc.py --sd.prompt "悉尼歌剧院的高画质远景照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

玫瑰

python /path/to/sjc/run_sjc.py --sd.prompt "单反相机拍摄的玫瑰照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 50

校车

python /path/to/sjc/run_sjc.py --sd.prompt "黄色校车的高画质照片" --n_steps 30000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0 --var_red False

火箭

python /path/to/sjc/run_sjc.py --sd.prompt "从远处拍摄的土星五号火箭广角远景照片" --n_steps 30000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0 --var_red False

薯条

python /path/to/sjc/run_sjc.py --sd.prompt "麦当劳薯条的高画质照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 10

摩托车

python /path/to/sjc/run_sjc.py --sd.prompt "玩具摩托车的高画质照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

汽车

python /path/to/sjc/run_sjc.py --sd.prompt "经典银色肌肉车的高画质照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

坦克

python /path/to/sjc/run_sjc.py --sd.prompt "玩具坦克的产品照片" --n_steps 20000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

椅子

python /path/to/sjc/run_sjc.py --sd.prompt "维多利亚风格木制天鹅绒软垫椅的高画质照片" --n_steps 50000 --lr 0.01 --sd.scale 100.0 --emptiness_weight 7000

鸭子

python /path/to/sjc/run_sjc.py --sd.prompt "单反相机拍摄的黄鸭照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 10

马

python /path/to/sjc/run_sjc.py --sd.prompt "一匹马在行走的照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

长颈鹿

python /path/to/sjc/run_sjc.py --sd.prompt "长颈鹿的广角远景照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 50

斑马

python /path/to/sjc/run_sjc.py --sd.prompt "斑马在行走的照片" --n_steps 10000 --lr 0.02 --sd.scale 100.0 --emptiness_weight 30000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0 --var_red False

打印机

python /path/to/sjc/run_sjc.py --sd.prompt "佳能家用打印机的产品照片" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0 --var_red False

塞尔达·林克

python /path/to/sjc/run_sjc.py --sd.prompt "塞尔达·林克" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0 --var_red False

猪

python /path/to/sjc/run_sjc.py --sd.prompt "一头猪" --n_steps 10000 --lr 0.05 --sd.scale 100.0 --emptiness_weight 10000 --emptiness_step 0.5 --emptiness_multiplier 20.0 --depth_weight 0

测试体素NeRF

python /path/to/sjc/run_nerf.py

我们的压缩包中包含乐高推土机数据集的tar文件。解压后即可使用。

使用Karras采样器生成2D图像

python /path/to/sjc/run_img_sampling.py

使用-h选项查看可用参数。详细信息将在后续补充。

参考文献

@article{sjc,
      title={Score Jacobian Chaining: Lifting Pretrained 2D Diffusion Models for 3D Generation},
      author={Wang, Haochen and Du, Xiaodan and Li, Jiahao and Yeh, Raymond A. and Shakhnarovich, Greg},
      journal={arXiv preprint arXiv:2212.00774},
      year={2022},
}

SJC (Score Jacobian Chaining) 快速上手指南

SJC 是一种利用预训练的 2D 扩散模型（如 Stable Diffusion）生成 3D 资产的技术。它通过链式法则将 2D 模型的梯度分数反向传播到可微渲染器的雅可比矩阵中，从而优化出 3D 体素辐射场。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• GPU: NVIDIA 显卡，显存建议 10GB 以上 (示例中基于 A5000 测试)
• CUDA: 已安装与 PyTorch 版本匹配的 CUDA 驱动
• Python: 3.8 或更高版本
• 前置依赖: git, wget 或 curl

安装步骤

1. 安装 PyTorch

根据您的 CUDA 版本安装 PyTorch。以下以 CUDA 11.6 为例：

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116

提示：国内用户可使用清华源加速安装：

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116 -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装项目依赖

克隆仓库后（假设已进入项目根目录），安装基础依赖：

pip install -r requirements.txt

3. 手动安装 taming-transformers

该项目需要手动安装 taming-transformers 组件：

git clone --depth 1 git@github.com:CompVis/taming-transformers.git && pip install -e taming-transformers

如果无法访问 GitHub，请尝试配置代理或使用国内镜像源克隆。

4. 下载模型检查点 (Checkpoints)

下载包含 Stable Diffusion v1.5 及所需 GDDPM 检查点的压缩包 (约 12GB)：

wget https://dl.ttic.edu/pals/sjc/release.tar
# 或者使用 curl
curl -O https://dl.ttic.edu/pals/sjc/release.tar

解压文件：

tar -xvf release.tar

重要配置： 解压后，请编辑项目根目录下的 env.json 文件，将检查点路径指向您刚才解压的文件夹位置。

基本使用

1. 创建实验目录

为了避免日志文件污染代码库，请在独立目录下运行实验：

mkdir exp
cd exp

2. 运行生成命令

使用以下命令生成一个简单的 3D 资产（例如“天坛”）。该过程在单张 A5000 GPU 上约需 25 分钟，占用约 10GB 显存。

python /path/to/sjc/run_sjc.py \
--sd.prompt "A zoomed out high quality photo of Temple of Heaven" \
--n_steps 10000 \
--lr 0.05 \
--sd.scale 100.0 \
--emptiness_weight 10000 \
--emptiness_step 0.5 \
--emptiness_multiplier 20.0 \
--depth_weight 0 \
--var_red False

参数说明：

• --sd.prompt: 发送给 Stable Diffusion 的文本提示词。
• --n_steps: 优化迭代的步数。
• --lr: 优化器的基础学习率。
• --sd.scale: Stable Diffusion 的引导尺度 (Guidance Scale)。
• --emptiness_weight / step / multiplier: 控制“空性损失”(emptiness loss) 的参数，用于防止几何体坍塌。
• --depth_weight: 中心深度损失的权重。
• --var_red: 是否使用方差缩减技术 (Eq. 16)，某些特定提示词下设为 False (Eq. 15) 效果更佳。

3. 查看结果

运行结束后，当前目录下会生成结果：

• test_* 文件夹中包含不同迭代阶段渲染的图像 (view) 和视频 (view_seq)。

4. (可选) 高质量可视化

训练完成后，可运行子像素渲染脚本以获取更高质量、无抖动的最终视频：

python /path/to/sjc/highres_final_vis.py

(注：请将 /path/to/sjc/ 替换为您实际的代码路径)