diffusers-rs：Rust/Torch中的Diffusers API

一个手持火炬的生锈机器人，由Rust和libtorch生成的稳定扩散。

“diffusers”软件包是Huggingface惊人的diffusers Python库的Rust等价物。它基于tch软件包。该实现支持运行Stable Diffusion v1.5和v2.1。

获取权重

可以从HuggingFace模型仓库检索权重文件，并将它们移动到“data/”目录中。

• 对于Stable Diffusion v2.1，从v2.1仓库获取“bpe_simple_vocab_16e6.txt”，“clip_v2.1.safetensors”，“unet_v2.1.safetensors”和“vae_v2.1.safetensors”文件。
• 对于Stable Diffusion v1.5，从v1.5仓库获取“bpe_simple_vocab_16e6.txt”，“pytorch_model.safetensors”，“unet.safetensors”和“vae.safetensors”文件。
• 或者，您可以运行以下Python脚本。

# Add --sd_version 1.5 to get the v1.5 weights rather than the v2.1.
python3 ./scripts/get_weights.py

运行一些示例。

cargo run --example stable-diffusion --features clap -- --prompt "A rusty robot holding a fire torch."

默认情况下，最终图像命名为“sd_final.png”。默认调度器是去噪扩散隐式模型调度器（DDIM）。原始论文和一些代码可以在相关仓库中找到。

这将生成一些手持火炬的生锈机器人的图像！

图像到图像管道

稳定扩散模型也可以用于根据另一张图像生成图像。以下命令运行此图像到图像管道

cargo run --example stable-diffusion-img2img --features clap -- --input-image media/in_img2img.jpg

默认提示语为“一个流行的艺术站幻想风景”，但可以通过-prompt标志进行更改。

修复绘画管道

修复绘画可以根据提示修改现有图像，并修改由遮罩指定的初始图像部分。这需要不同的unet权重unet-inpaint.safetensors，这些权重也可以从该仓库获取，并应放置在data/目录中。

以下命令运行此图像到图像管道

wget https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png -O sd_input.png
wget https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png -O sd_mask.png
cargo run --example stable-diffusion-inpaint --features clap --input-image sd_input.png --mask-image sd_mask.png

默认提示语为“黄色猫的脸，高分辨率，坐在公园长椅上。”，但可以通过-prompt标志进行更改。

ControlNet管道

ControlNet架构可用于控制稳定扩散生成图像的方式。这需要使用稳定扩散1.5的权重（见上方获取方法）。还需要从HuggingFace仓库检索额外的权重，并将其复制到data/controlnet.safetensors。

ControlNet管道以样本图像为输入，在默认模式下，它将使用Canny边缘检测器在此图像上进行边缘检测，并将生成的边缘图像作为引导。

cargo run --example controlnet --features clap,image,imageproc -- \
  --prompt "a rusty robot, lit by a fire torch, hd, very detailed" \
  --input-image media/vermeer.jpg

media/vermeer.jpg图像是左侧著名的画作，执行边缘检测后得到右侧的图像。

仅使用边缘检测图像，ControlNet模型生成了以下样本。

常见问题解答

内存问题

这需要一个大于8GB内存的GPU，作为后备，可以使用CPU版本，但速度较慢。

cargo run --example stable-diffusion --features clap -- --prompt "A very rusty robot holding a fire torch." --cpu all

对于8GB的GPU，可以使用UNet的fp16权重，并将UNet仅放在GPU上。

PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.6,max_split_size_mb:128 RUST_BACKTRACE=1 CARGO_TARGET_DIR=target2 cargo run \
    --example stable-diffusion --features clap -- --cpu vae --cpu clip \
    --unet-weights data/unet-fp16.safetensors