StyleDrop

StyleDrop-PyTorch 的代码已迁移至 aim-uofa/StyleDrop-PyTorch。欢迎大家尝试并享受其中的乐趣！

这是对 StyleDrop: Text-to-Image Generation in Any Style 的非官方 PyTorch 实现。

与论文中的参数（第一轮）不同，我们设置了 $\lambda_A=2.0$、$\lambda_B=5.0$，以及 d_prj=32 和 is_shared=False，这些超参数经我们测试后表现更好，具体可在 configs/custom.py 中查看。 我们将其公开，以促进社区研究。

最新消息

• [2023年7月11日] StyleDrop-PyTorch 的代码已迁移至 aim-uofa/StyleDrop-PyTorch。欢迎大家尝试并享受其中的乐趣！
• [2023年7月6日] 在线 Gradio 演示现已上线，访问地址为 这里

待办事项

• 发布代码。
• 添加本地运行的 Gradio 推理演示。
• 添加迭代训练（第二轮）。

数据与权重准备

首先，请从该 链接 下载 VQGAN（来自 MAGE，感谢！），并将下载的 VQGAN 放入 assets/vqgan_jax_strongaug.ckpt。

然后，从该 链接 下载预训练检查点，并将其放置在 assets/ckpts 目录下，用于评估或继续训练更多轮次。

最后，通过运行命令 python extract_empty_feature.py 准备空特征。

最终的目录结构如下：

.
├── assets
│   ├── ckpts
│   │   ├── cc3m-285000.ckpt
│   │   │   ├── lr_scheduler.pth
│   │   │   ├── nnet_ema.pth
│   │   │   ├── nnet.pth
│   │   │   ├── optimizer.pth
│   │   │   └── step.pth
│   │   └── imagenet256-450000.ckpt
│   │       ├── lr_scheduler.pth
│   │       ├── nnet_ema.pth
│   │       ├── nnet.pth
│   │       ├── optimizer.pth
│   │       └── step.pth
│   ├── fid_stats
│   │   ├── fid_stats_cc3m_val.npz
│   │   └── fid_stats_imagenet256_guided_diffusion.npz
│   ├── pipeline.png
|   ├── contexts
│   │   └── empty_context.npy
└── └── vqgan_jax_strongaug.ckpt

依赖项

与 MUSE-PyTorch 相同。

conda install pytorch torchvision torchaudio cudatoolkit=11.3
pip install accelerate==0.12.0 absl-py ml_collections einops wandb ftfy==6.1.1 transformers==4.23.1 loguru webdataset==0.2.5 gradio

训练

论文中所有风格数据均放置在 data 目录中。

1. 修改 data/one_style.json 文件（需要注意的是，one_style.json 和风格数据必须位于同一目录），格式为 file_name:[object,style]。

{"image_03_05.jpg":["一只熊","儿童蜡笔画风格"]}

2. 训练脚本如下。

#!/bin/bash
unset EVAL_CKPT
unset ADAPTER
export OUTPUT_DIR="output_dir/for/this/experiment"
accelerate launch --num_processes 8 --mixed_precision fp16 train_t2i_custom_v2.py --config=configs/custom.py

推理

预训练的 style_adapter 权重可从 🤗 Hugging Face 下载。

#!/bin/bash
export EVAL_CKPT="assets/ckpts/cc3m-285000.ckpt" 
export ADAPTER="path/to/your/style_adapter"

export OUTPUT_DIR="output/for/this/experiment"

accelerate launch --num_processes 8 --mixed_precision fp16 train_t2i_custom_v2.py --config=configs/custom.py

Gradio 演示

将 style_adapter 权重 放入 ./style_adapter 文件夹，并运行以下命令即可启动演示：

python gradio_demo.py

该演示也在 HuggingFace 上提供。

引用

@article{sohn2023styledrop,
  title={StyleDrop: Text-to-Image Generation in Any Style},
  author={Sohn, Kihyuk and Ruiz, Nataniel and Lee, Kimin and Chin, Daniel Castro and Blok, Irina and Chang, Huiwen and Barber, Jarred and Jiang, Lu and Entis, Glenn and Li, Yuanzhen and others},
  journal={arXiv preprint arXiv:2306.00983},
  year={2023}
}

致谢

• 本实现基于 MUSE-PyTorch。
• 非常感谢 Zanlin Ni 的慷慨帮助。

星标历史

StyleDrop-PyTorch 快速上手指南

StyleDrop 是一个基于 PyTorch 的非官方实现，旨在实现“任意风格的文本到图像生成”。本指南将帮助你快速搭建环境并运行模型。

1. 环境准备

系统要求

• 操作系统: Linux (推荐)
• Python: 3.8+
• GPU: 支持 CUDA 的 NVIDIA 显卡 (建议显存 16GB 以上以进行训练)
• CUDA Toolkit: 11.3

前置依赖

本项目依赖 MUSE-PyTorch 的环境配置。请确保已安装 conda。

2. 安装步骤

2.1 创建虚拟环境并安装基础包

conda create -n styledrop python=3.9 -y
conda activate styledrop

# 安装 PyTorch (CUDA 11.3)
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

# 安装其他 Python 依赖
pip install accelerate==0.12.0 absl-py ml_collections einops wandb ftfy==6.1.1 transformers==4.23.1 loguru webdataset==0.2.5 gradio

2.2 下载模型权重与数据

项目需要 VQGAN 权重、预训练检查点以及空特征文件。请按以下步骤操作：

1. 下载 VQGAN 权重： 从 Google Drive 下载 vqgan_jax_strongaug.ckpt，并将其放入 assets/ 目录下。

2. 下载预训练检查点： 从 Hugging Face 下载所需检查点（如 cc3m-285000.ckpt 或 imagenet256-450000.ckpt），放入 assets/ckpts/ 目录。 注：国内用户若访问 Hugging Face 困难，可使用镜像站 hf-mirror.com 或寻找国内网盘搬运资源。

3. 生成空特征文件： 在项目根目录运行以下命令：

   python extract_empty_feature.py
   

2.3 目录结构确认

完成后，你的目录结构应如下所示：

.
├── assets
│   ├── ckpts
│   │   ├── cc3m-285000.ckpt/...
│   │   └── imagenet256-450000.ckpt/...
│   ├── fid_stats/...
│   ├── contexts
│   │   └── empty_context.npy
│   └── vqgan_jax_strongaug.ckpt
└── ... (其他代码文件)

3. 基本使用

3.1 准备风格数据

在 data/ 目录下创建或修改 one_style.json 文件。 注意：该 JSON 文件必须与风格图片位于同一目录。 格式示例 (file_name: [物体描述, 风格描述]`)：

{"image_03_05.jpg":["A bear","in kid crayon drawing style"]}

请将你的风格参考图（如 image_03_05.jpg）放在 data/ 目录中。

3.2 开始训练

修改 configs/custom.py 可调整超参数（默认已优化为 $\lambda_A=2.0$, $\lambda_B=5.0$ 等）。 运行以下脚本启动训练（示例使用 8 卡混合精度训练）：

#!/bin/bash
unset EVAL_CKPT
unset ADAPTER
export OUTPUT_DIR="output_dir/for/this/experiment"
accelerate launch --num_processes 8 --mixed_precision fp16 train_t2i_custom_v2.py --config=configs/custom.py

3.3 推理与演示 (Gradio)

如果你已有训练好的 style_adapter 权重（可从 Hugging Face 下载），可将其放入 ./style_adapter 文件夹，然后启动本地 Gradio 演示界面：

python gradio_demo.py

启动后，浏览器会自动打开界面，你可以输入提示词并选择风格进行图像生成测试。

---

注：最新代码已迁移至 aim-uofa/StyleDrop-PyTorch，如需迭代训练（Round 2）等高级功能，建议前往新仓库查看。