stable-diffusion-pytorch - Stable Diffusion PyTorch实现，支持自定义参数

项目介绍：stable-diffusion-pytorch

stable-diffusion-pytorch 是一个基于 PyTorch 的 Stable Diffusion 实现。项目作者致力于使代码库简洁，自包含，易于阅读和修改。对于使用不需要的功能，代码中进行了精简，比如不必要的注意力掩码等。此外，配置都基于 Stable Diffusion v1.x 进行了硬编码，循环拆开以提高代码的可读性和逻辑清晰度。

尽管作者努力地优化了代码，他还是觉得这像是做了一盘“意大利面”，即代码可能显得纷繁复杂。不过，别担心，这个项目还是很有启发性的。

项目依赖

要运行 stable-diffusion-pytorch，你需要以下软件包：

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 可作为子模块进行导入。以下是一些示例代码：

文本生成图像

使用文本提示生成图像：

from stable_diffusion_pytorch import pipeline

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts)
images[0].save('output.jpg')

支持多个文本提示：

prompts = [
    "a photograph of an astronaut riding a horse",
    ""
]
images = pipeline.generate(prompts)

使用无条件（负面）提示生成：

prompts = ["a photograph of an astronaut riding a horse"]
uncond_prompts = ["low quality"]
images = pipeline.generate(prompts, uncond_prompts)

设置随机种子：

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts, uncond_prompts, seed=42)

预加载模型

如果有足够的显存，可以预加载模型：

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

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts, models=models)

如果显存不足但系统内存足够，可视需要将模型移至GPU或CPU：

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

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts, models=models, device='cuda', idle_device='cpu')

图像生成图像

从输入图像生成新图像：

from PIL import Image

prompts = ["a photograph of an astronaut riding a horse"]
input_images = [Image.open('space.jpg')]
images = pipeline.generate(prompts, input_images=input_images)

设置自定义强度参数：

prompts = ["a photograph of an astronaut riding a horse"]
input_images = [Image.open('space.jpg')]
images = pipeline.generate(prompts, input_images=input_images, strength=0.6)

修改其它设置

调整无分类指引的比例：

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts, cfg_scale=11)

禁用无分类指导：

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts, do_cfg=False)

减少生成步骤以加快速度（可能会降低质量）：

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts, n_inference_steps=28)

使用不同的采样器：

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts, sampler="k_euler")
# 还可选用 "k_lms"（默认）和 "k_euler_ancestral"

生成自定义尺寸的图像：

prompts = ["a photograph of an astronaut riding a horse"]
images = pipeline.generate(prompts, height=512, width=768)

许可证

此仓库中的所有代码均采用 MIT License 许可证。请参见 LICENSE 文件。需要注意的是，Stable Diffusion 的检查点文件使用 CreativeML Open RAIL-M 许可证，在使用前请阅读许可证中的使用限制条款。