Cones 2
Cones 2: 可定制多主体图像合成的官方代码库 | 项目主页。
Cones 2允许您通过微调预训练文本到图像扩散模型(如Stable Diffusion)中的文本编码器,将特定主体表示为残差嵌入。微调后,我们只需保存微调后的文本编码器与冻结编码器之间的残差。因此,每个额外主体所需的存储空间仅为5 KB。这一步骤在单个80G A100 GPU上对每个主体只需约20~30分钟。
在采样时,我们的布局引导采样方法进一步允许您使用易于获取的布局作为多个主体排列的指导,如下图所示。
结果
所有结果都是由预训练的Stable Diffusion v2.1模型和我们的自定义残差嵌入合成的。我们展示了各种类别图像的多样化结果,包括场景、宠物、个人玩具、人物等。更多结果请参考我们的论文或网站。
两个主体的结果
三个主体的结果
四个主体的结果
更多结果
方法
(a) 给定自定义主体的少量图像,我们微调文本编码器,在原始主体的基本嵌入之上学习残差嵌入。(b) 基于残差嵌入,我们随后提出将布局作为空间指导,用于在注意力图中排列主体。之后,我们可以加强目标主体的信号并削弱不相关主体的信号。更多详情,请参阅我们的论文。
入门指南
安装依赖项
Cones 2的实现完全基于diffuser。 在运行我们的代码之前,请确保安装库的训练依赖项。要做到这一点,请在新的虚拟环境中执行以下步骤:
git clone https://github.com/damo-vilab/Cones-V2.git
cd Cones-V2
pip install -r requirements.txt
然后使用以下命令初始化🤗—Accelerate环境:
accelerate config
或者使用以下命令进行默认加速配置,无需回答有关环境的问题
accelerate config default
训练(花朵示例)
首先,让我们从这里下载数据集,并将其解压到./data
。现在使用几张花朵图像(下载到./data/flower
)来学习其自定义残差嵌入。
export MODEL_NAME='path-to-stable-diffusion-v2-1'
export INSTANCE_DIR="./data/flower"
export OUTPUT_DIR="path-to-save-model"
accelerate launch train_cones2.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--instance_prompt="flower" \
--token_num=1 \
--output_dir=$OUTPUT_DIR \
--resolution=768 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=4000 \
--loss_rate_first=1e-2 \
--loss_rate_second=1e-3
推理
一旦您使用上述命令训练了几个不同主体的残差嵌入,您可以使用inference.py简单地运行我们的布局引导采样方法。我们提供了几个预训练模型以供快速验证。
残差嵌入 | 原始标记 | 下载 |
---|---|---|
谷仓 | 谷仓 | barn.pt |
白色狗 | 狗 | dog.pt |
花 | 花 | flower.pt |
湖 | 湖 | lake.pt |
杯子 | 杯子 | mug.pt |
太阳镜 | 太阳镜 | sunglasses.pt |
包 | 包 | bag.pt |
机器人 | 机器人玩具 | robot_toy.pt |
请记得提供一个预定义的布局,如layout_example.png,以及一个包含推理设置详细信息的json文件。json文件应包含以下信息:
- "prompt":您想要生成的文本提示。
- "residual_dict":所有所需残差嵌入的路径。
- "color_context":布局中不同区域的颜色信息及其对应的主题,以及用于加强目标主题信号的权重(默认:2.5)。
- "guidance_steps":布局引导的步数。
- "guidance_weight":布局引导的强度(默认:0.08,我们建议 0.05 ~ 0.10)。
- "weight_negative":用于减弱不相关主题信号的权重。
- "layout":用户定义布局图像的路径。
- "subject_list":包含所有要自定义的主题及其在提示中对应位置的列表。
然后您可以简单地运行推理脚本:
python inference.py --pretrained_model_name_or_path /path/to/stable-diffusion-2-1 --inference_config guidance_config_example.json
参考文献
@article{liu2023cones,
title={Cones 2: Customizable Image Synthesis with Multiple Subjects},
author={Liu, Zhiheng and Zhang, Yifei and Shen, Yujun and Zheng, Kecheng and Zhu, Kai and Feng, Ruili and Liu, Yu and Zhao, Deli and Zhou, Jingren and Cao, Yang},
journal={arXiv preprint arXiv:2305.19327},
year={2023}
}
致谢
我们感谢Stable Diffusion v2.1和diffuser提供预训练模型和开源代码库。