DEV Community

Cover image for Implementing Stable Diffusion in Rust
DhruvThu
DhruvThu

Posted on

Implementing Stable Diffusion in Rust

Recently, we have seen the boom of many machine learning algorithms such as stable diffusion which can be used to create digital artworks or NFTs. One of the primary issues with stable diffusion is the possibility of memory leaks in Python, which can result in excessive use of GPU and CPU resources. To prevent this type of scenario, we at Qolaba chose to explore with Stable Diffusion developed in Rust in order to benefit from Rust’s memory management. This article will walk you through the process of building stable diffusion in Rust.

Table of contents:

  • What is Rust and its advantage?
  • What is Stable-diffusion?
  • Implementation of Rust based Stable-Diffusion.

What is Rust and its advantage?

Rust is a comparatively recent programming language that gained popularity swiftly due to its capacity to build dependable, memory-efficient, and high-performance program. The syntax of this statically typed programming language is comparable to C++. It doesn’t have run time or garbage collection. Rust thus offers solutions to numerous C++ problems, such as concurrency and memory management problems.

What is Stable-diffusion?

A machine learning system called Stable Diffusion uses diffusion to produce visuals from text. It functions as both Text to Image and Image to Image. The Stable Diffusion model is used to generate the majority of contemporary AI art that can be found online. With simply a word prompt and an open-source application, anyone can easily produce amazing art images. We prefer to run the model on the GPU because it is quite time-consuming to run on the CPU. The new version of Stable diffusion offers a number of features, which may be found on the stability ai blog at Stable Diffusion 2.0 Release — Stability AI.

Stable Diffusion Architecture

Implementation of Rust based Stable-Diffusion

  • To execute the Rust-based stable diffusion, we must first clone the Rust-based stable diffusion implementation from GitHub.
git clone https://github.com/LaurentMazare/diffusers-rs.git
Enter fullscreen mode Exit fullscreen mode
  • Next we need to install Rust. We can utilize the Ubuntu commands listed below. But, for Windows, we may follow the procedure outlined in this Link.
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
source $HOME/.cargo/env
Enter fullscreen mode Exit fullscreen mode
  • Following that, we need to obtain the model’s weights from hugging face model hub. Hugging face model hub is based on Git LFS, and we can store massive files (more than 10 GB), such as ML model checkpoints there. It operates on the Git LFS, so we can use it in the same way as a Git repository.

  • The following command may be used to install git LFS on Ubuntu. The command discussed in this article can be used on different platforms. This will require the Git pre installed.

curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt-get install git-lfs
git lfs install
Enter fullscreen mode Exit fullscreen mode
  • Once Git LFS is configured, we must clone the repo from Hugging face model hub. We can accomplish this by using the command listed below.
git clone https://huggingface.co/lmz/rust-stable-diffusion-v2-1
Enter fullscreen mode Exit fullscreen mode
  • After obtaining the weight, we must place the weights folder into the diffusers-rs repo that we already cloned.
cd rust-stable-diffusion-v2-1/
mv weights data
mv data/ ../diffusers-rs/
cd ..
Enter fullscreen mode Exit fullscreen mode
  • To execute the Rust-based stable diffusion with GPU, we must first configure Libtorch. The command may be obtained from the official Pytorch installation page, which is also referenced below. We need to unzip the folder after collecting the Libtorch.
wget https://download.pytorch.org/libtorch/cu117/libtorch-cxx11-abi-shared-with-deps-1.13.1%2Bcu117.zip
unzip libtorch-cxx11-abi-shared-with-deps-1.13.1+cu117.zip

Enter fullscreen mode Exit fullscreen mode
  • Once the preceding step is accomplished, we must add the Llibtorch folder to the path in the environment variable so that it can be used by the rust programme. We may directly write the environment variable using the export command, or we can put it in the bashrc file and it will be run every time we restart the terminal.
vi .bashrc
Enter fullscreen mode Exit fullscreen mode
export LIBTORCH=/path/to/libtorch
export LD_LIBRARY_PATH=${LIBTORCH}/lib:$LD_LIBRARY_PATH
Enter fullscreen mode Exit fullscreen mode
  • When we have finished modifying the bashrc file, we must restart the terminal so that the aforementioned command may be run.
  • Following that, we need execute the rust-based application using the cargo command, which is detailed below.
cd diffusers-rs/
Enter fullscreen mode Exit fullscreen mode
cargo run --example stable-diffusion --features clap -- --prompt "the cutest chibi sloth you'll ever see dancing wearing a 3 piece suit holding a bunch of flowers pixar style"
Enter fullscreen mode Exit fullscreen mode

Output of Rust SD for v2.1 (Prompt : the cutest chibi sloth you’ll ever see dancing wearing a 3 piece suit holding a bunch of flowers pixar style)

The total time needed for generating one image for Rust-based stable diffusion is more than that necessary for Python-based stable diffusion. With more efficient code, there is a likelihood that total response time in Rust-based Stable diffusion will be lowered in the near future.

Thankyou🙏

References

  1. Rust vs Python: Which One Is Best for Your Project?
  2. The Illustrated Stable Diffusion
  3. GitHub — LaurentMazare/diffusers-rs: An implementation of the diffusers api in Rust

Top comments (0)