Repository files navigation

A rusty robot holding a fire torch, generated by stable diffusion using Rust and libtorch.

Thediffusers crate is a Rust equivalent to Huggingface's amazingdiffusers Python library.It is based on thetch crate.The implementation supports running Stable Diffusion v1.5 and v2.1.

The weight files can be retrieved from the HuggingFace model repos and should bemoved in thedata/ directory.

• For Stable Diffusion v2.1, get thebpe_simple_vocab_16e6.txt,clip_v2.1.safetensors,unet_v2.1.safetensors, andvae_v2.1.safetensorsfiles from thev2.1 repo.
• For Stable Diffusion v1.5, get thebpe_simple_vocab_16e6.txt,pytorch_model.safetensors,unet.safetensors, andvae.safetensorsfiles from thisv1.5 repo.
• Alternatively, you can run the following python script.

# 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."

The final image is namedsd_final.png by default.The default scheduler is the Denoising Diffusion Implicit Model scheduler (DDIM). Theoriginal paper and some code can be found in theassociated repo.

This generates some images of rusty robots holding some torches!

The stable diffusion model can also be used to generate an image based onanother image. The following command runs this image to image pipeline:

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

The default prompt is "A fantasy landscape, trending on artstation.", but canbe changed via the-prompt flag.

Inpainting can be used to modify an existing image based on a prompt and modifying the part of theinitial image specified by a mask.This requires different unet weightsunet-inpaint.safetensors that could also be retrieved from thisrepo and should also beplaced in thedata/ directory.

The following command runs this image to image pipeline:

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

The default prompt is "Face of a yellow cat, high resolution, sitting on a park bench.", but canbe changed via the-prompt flag.

TheControlNet architecture can beused to control how stable diffusion generate images. This is to be used withthe weights for stable diffusion 1.5 (see how to get these above). Additionalweights have to be retrieved from thisHuggingFacerepoand copied indata/controlnet.safetensors.

The ControlNet pipeline takes as input a sample image, in the default mode itwill perform edge detection on this image using theCanny edgedetector and will use theresulting edge image as a guide.

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

Themedia/vermeer.jpg image is the well known painting on the left hand side,this results in the right hand side image after performing edge detection.

Using only the edge detection image, the ControlNet model generate the followingsamples.

This requires a GPU with more than 8GB of memory, as a fallback the CPU version can be usedbut is slower.

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

For a GPU with 8GB, one can use thefp16 weights for the UNet and put only the UNet on the 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