Lib.rs

---

pyke Diffusers 是一个模块化 Rust 库，用于预训练扩散模型推理以生成图像、视频或音频，使用 ONNX Runtime 作为后端，在 CPU 和 GPU 上实现高度优化的生成。

• 功能
• 先决条件
  • LMS笔记
• 安装
• 使用
  • 示例
  • 模型转换
  • ONNX Runtime 二进制文件
  • CUDA和其他执行提供者
  • 低内存使用
    • 量化

功能

• 稳定扩散v1 & v2的文本到图像
• 优化CPU和GPU推理
• 内存高效的管道，可在小于2GB的RAM上运行！
• >77个令牌提示
• 提示加权，例如 a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).
• 实现DDIM、DDPM、DPM/DPM++、Euler & Euler a、LMS调度器

先决条件

使用pyke Diffusers需要 Rust v1.62.1+

• 如果使用CPU：最近（不早于Haswell/Zen）的x86-64 CPU以获得最佳效果。ARM64受支持，但不太推荐。有关加速的说明，请参阅OpenVINO、oneDNN、ACL、SNPE的说明。
• 如果使用CUDA: CUDA v11.x, cuDNN v8.2.x ^更多信息
• 如果使用TensorRT: CUDA v11.x, TensorRT v8.4 ^更多信息
• 如果使用ROCm: ROCm v5.2 ^更多信息
• 如果使用DirectML: DirectX 12兼容GPU, Windows 10 v1903+ ^更多信息

只有CPU、CUDA和TensorRT有预构建的二进制文件可用。其他执行提供程序需要您手动构建它们；请参阅ONNX Runtime文档以获取更多信息。另外，您还需要将ort链接到您自定义构建的二进制文件。

LMS笔记

注意: 默认情况下，LMS调度器未启用，本节可以简单地跳过。

如果您计划启用all-schedulers或scheduler-lms功能，您需要安装GNU科学库的二进制文件。请参阅rust-GSL的安装说明以设置GSL。

安装

[dependencies]
pyke-diffusers = "0.1"
# if you'd like to use CUDA:
pyke-diffusers = { version = "0.1", features = [ "ort-cuda" ] }

默认功能启用了某些常用的调度器和管道。

使用

use pyke_diffusers::{
    Environment, EulerDiscreteScheduler, SchedulerOptimizedDefaults, StableDiffusionOptions, StableDiffusionPipeline,
    StableDiffusionTxt2ImgOptions
};

let environment = Arc::new(Environment::builder().build()?);
let mut scheduler = EulerDiscreteScheduler::stable_diffusion_v1_optimized_default()?;
let pipeline = StableDiffusionPipeline::new(&environment, "./stable-diffusion-v1-5", &StableDiffusionOptions::default())?;

let imgs = pipeline.txt2img("photo of a red fox", &mut scheduler, &StableDiffusionTxt2ImgOptions::default())?;
imgs[0].clone().into_rgb8().save("result.png")?;

示例

pyke-diffusers包含一个交互式Stable Diffusion演示。使用以下命令运行它：

$ cargo run --example stable-diffusion-interactive --features ort-cuda -- ~/path/to/stable-diffusion/

请参阅examples/获取更多示例和文档以获取更详细的信息。

模型转换

pyke Diffusers目前支持Stable Diffusion v1、v2及其衍生品。

要将模型从HuggingFace diffusers模型转换

1. 创建并激活一个虚拟环境。
2. 安装Python需求
   • 使用CUDA安装torch: python3 -m pip install torch --extra-index-url https://download.pytorch.org/whl/cu116
   • 安装依赖项: python3 -m pip install --r requirements.txt
3. 如果您直接从HuggingFace转换模型，请使用huggingface-cli login登录到HuggingFace Hub - 如果您已在磁盘上有模型，则可以跳过此步骤
4. 使用scripts/hf2pyke.py转换您的模型
   • 要从HF转换float32模型（推荐用于CPU推理）: python3 scripts/hf2pyke.py runwayml/stable-diffusion-v1-5 ~/pyke-diffusers-sd15/
   • 要从磁盘转换float32模型: python3 scripts/hf2pyke.py ~/stable-diffusion-v1-5/ ~/pyke-diffusers-sd15/
   • 将 float16 模型从 HF 转换（推荐用于 GPU 推理）：python3 scripts/hf2pyke.py --fp16 runwayml/stable-diffusion-v1-5@fp16 ~/pyke-diffusers-sd15-fp16/
   • 从磁盘转换 float16 模型：python3 scripts/hf2pyke.py --fp16 ~/stable-diffusion-v1-5-fp16/ ~/pyke-diffusers-sd15-fp16/

float16 模型在某些 GPU 上运行更快且内存使用更少。然而，需要注意的是，如果您在 GPU 推理中使用 float16 模型，由于 ONNX Runtime 的一个错误，它们必须在将要运行的硬件上转换。然而，使用 float16 模型的 CPU 应该没有这个问题。

hf2pyke 支持一些选项来提高性能或 ORT 执行提供者兼容性。请参阅 python3 scripts/hf2pyke.py --help。

ONNX Runtime 二进制文件

在 Windows 上运行此存储库中的示例时，您需要在首次运行时将 onnxruntime* 动态库从 target/debug/ 复制到 target/debug/examples/。如果您的项目在 Cargo 测试中使用 pyke Diffusers，您还需要将 dylibs 复制到 target/debug/deps/。

CUDA和其他执行提供者

CUDA 是唯一无需设置的替代执行提供者。只需启用 pyke Diffusers 的 ort-cuda 功能，并启用 DiffusionDevice::CUDA；有关更多信息，请参阅文档或stable-diffusion 示例。您可能需要重新构建项目以复制库。

对于其他 EP，如 DirectML 或 oneDNN，您需要从源代码构建 ONNX Runtime。请参阅 ort 的执行提供者说明。

低内存使用

分辨率较低的生成需要更少的内存使用。

存在一个用于内存低的环境的 StableDiffusionMemoryOptimizedPipeline。此管道 移除安全检查器，并且只有在需要时才会加载模型，并在之后立即卸载它们。这将严重影响性能，并且仅在极端情况下使用。

量化

在极其受限的环境中（例如 <= 4GB RAM），还可以生成量化 int8 模型。int8 模型的质量会受到严重影响，但在 CPU 上更快且内存占用更少。

要转换 int8 模型

$ python3 scripts/hf2pyke.py --quantize=ut ~/stable-diffusion-v1-5/ ~/pyke-diffusers-sd15-quantized/

--quantize=ut 将仅使用 uint8 模式对 UNet 和文本编码器进行量化以获得最佳质量和性能。您可以选择使用以下格式转换其他模型

• 每个模型都分配了一个字母：u 表示 UNet，v 表示 VAE，t 表示文本编码器。
• 小写字母表示模型将量化为 uint8
• 大写字母表示模型将量化为 int8

通常，uint8 质量更高且速度更快，但您可以尝试不同的设置，看看质量或速度是否有所改善。

通过 StableDiffusionMemoryOptimizedPipeline 生成 256x256 图像与一个 uint8 UNet 结合，仅需 1.3 GB 的内存使用。