stable-diffusion-pytorch

另一个 Stable Diffusion 的 PyTorch 实现。

我尽力使代码库尽可能简洁、自包含、一致、可修改且易于阅读。如果不需要的功能在 Stable Diffusion 中被修剪掉（例如 CLIP tokenizer/encoder 中的 Attention mask）。配置是硬编码的（基于 Stable Diffusion v1.x）。当这种形状更合理时，循环会被展开。

尽管我做了很多努力，我感觉好像我又做了一盘意大利面。好吧，随你欣赏吧！

大量参考了以下仓库。非常感谢！

依赖

PyTorch
Numpy
Pillow
regex
tqdm

如何安装

克隆或下载这个仓库。
安装依赖: 运行 pip install torch numpy Pillow regex 或 pip install -r requirements.txt。
从这里下载 data.v20221029.tar 并在 stable_diffusion_pytorch 的父文件夹中解压。你的文件夹结构应该如下：

stable-diffusion-pytorch(-main)/
├─ data/
│  ├─ ckpt/
│  ├─ ...
├─ stable_diffusion_pytorch/
│  ├─ samplers/
└  ┴─ ...

注意 data.zip 中包含的检查点文件有不同的许可证——您需要同意许可证才能使用检查点文件。

如何使用

将 stable_diffusion_pytorch 作为子模块导入。

以下是一些示例脚本。您也可以查看 stable_diffusion_pytorch.pipeline.generate 的文档字符串。

文本到图像生成：

from stable_diffusion_pytorch import pipeline

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts)
images[0].save('output.jpg')

...多个提示：

prompts = [
    "一张宇航员骑马的照片",
    ""]
images = pipeline.generate(prompts)

...无条件（负面）提示：

prompts = ["一张宇航员骑马的照片"]
uncond_prompts = ["低质量"]
images = pipeline.generate(prompts, uncond_prompts)

...带种子：

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts, uncond_prompts, seed=42)

预加载模型（您需要足够的显存）：

from stable_diffusion_pytorch import model_loader
models = model_loader.preload_models('cuda')

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts, models=models)

如果使用上述代码时出现 OOM 但有足够的 RAM（而不是显存），您可以在需要时将模型移到 GPU 上，闲置时移回 CPU 上：

from stable_diffusion_pytorch import model_loader
models = model_loader.preload_models('cpu')

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts, models=models, device='cuda', idle_device='cpu')

图像到图像生成：

from PIL import Image

prompts = ["一张宇航员骑马的照片"]
input_images = [Image.open('space.jpg')]
images = pipeline.generate(prompts, input_images=images)

...自定义强度：

prompts = ["一张宇航员骑马的照片"]
input_images = [Image.open('space.jpg')]
images = pipeline.generate(prompts, input_images=images, strength=0.6)

更改classifier-free guidance比例：

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts, cfg_scale=11)

...或禁用 classifier-free guidance：

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts, do_cfg=False)

减少步骤（更快生成、质量较低）：

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts, n_inference_steps=28)

使用不同的采样器：

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts, sampler="k_euler")
# 可用 "k_lms"（默认）、"k_euler" 或 "k_euler_ancestral"

生成自定义尺寸的图像：

prompts = ["一张宇航员骑马的照片"]
images = pipeline.generate(prompts, height=512, width=768)

许可证

此仓库中的所有代码均使用 MIT 许可证。请参阅 LICENSE 文件。

注意 Stable Diffusion 的检查点文件使用 CreativeML Open RAIL-M 许可证。它有基于用途的限制条款，所以你最好阅读一下。