Skip to main content

Stable Diffusion on a CPU

stars - badge-generator

Stable Diffusion is a state of the art text-to-image model that generates images from text and was developed as an open-source alternative to DALL·E 2. It is based on a Diffusion Probabilistic Model and uses a Transformer to generate images from text.

This example demonstrates how to use stable diffusion on a CPU and run it on the Bacalhau network. The first section describes the development of the code and the container. The section demonstrates how to run the job using Bacalhau.

The following image is an example generated by this model.

TL;DR​

bacalhau docker run ghcr.io/bacalhau-project/examples/stable-diffusion-cpu:0.0.1 -- python demo.py --prompt "cod in space" --output ../outputs/cod.png

1. Development​

The original text-to-image stable diffusion model was trained on a fleet of GPU machines, at great cost. To use this trained model for inference, you also need to run it on a GPU.

However, this isn't always desired or possible. One alternative is to use a project called OpenVINO from Intel that allows you to convert and optimize models from a variety of frameworks (and ONNX if your framework isn't directly supported) to run on a supported Intel CPU. This is what we will do in this example.

tip

Heads up! This example takes about 10 minutes to generate an image on an average CPU. Whilst this demonstrates it is possible, it might not be practical.

Prerequisites​

In order to run this example you need:

Converting Stable Diffusion to a CPU Model Using OpenVINO​

The first step is to convert the trained stable diffusion models so that they work efficiently on a CPU using OpenVINO. The example is quite complex, so we have created a separate repository (which is a fork from Github user Sergei Belousov) to host the code. In summary, the code:

The core code representing these tasks can be found in the stable_diffusion_engine.py file. This is a mashup that creates a pipeline necessary to tokenize the text and run the stable diffusion model. This boilerplate could be simplified by leveraging the more recent version of the diffusers library. But let's crack on.

Install Dependencies​

Note that these dependencies are only known to work on Ubuntu-based x64 machines.

%%bash
sudo apt-get update
sudo apt-get install -y libgl1 libglib2.0-0 git-lfs

Clone the Repository and Dependencies​

The following commands clone the example repository, and other required repositories, and install the Python dependencies.

%%bash
git clone https://github.com/js-ts/stable_diffusion.openvino
cd stable_diffusion.openvino
git lfs install
git clone https://huggingface.co/openai/clip-vit-large-patch14
git clone https://huggingface.co/bes-dev/stable-diffusion-v1-4-openvino
pip3 install -r requirements.txt

Generating an Image​

Now that we have all the dependencies installed, we can call the demo.py wrapper, which is a simple CLI, to generate an image from a prompt.

!cd stable_diffusion.openvino && \
  python3 demo.py --prompt "hello" --output hello.png
import IPython.display as display
display.Image("stable_diffusion.openvino/hello.png")

png

!cd stable_diffusion.openvino && \
  python3 demo.py --prompt "cat driving a car" --output cat.png
import IPython.display as display
display.Image("stable_diffusion.openvino/cat.png")

png

2. Running Stable Diffusion (CPU) on Bacalhau​

Now we have a working example, we can convert it into a format that allows us to perform inference in a distributed environment.

First we will create a Dockerfile to containerize the inference code. The Dockerfile can be found in the repository, but is presented here to aid understanding.

FROM python:3.9.9-bullseye

WORKDIR /src

RUN apt-get update && \
    apt-get install -y \
    libgl1 libglib2.0-0 git-lfs

RUN git lfs install

COPY requirements.txt /src/

RUN pip3 install -r requirements.txt

COPY stable_diffusion_engine.py demo.py demo_web.py /src/
COPY data/ /src/data/

RUN git clone https://huggingface.co/openai/clip-vit-large-patch14
RUN git clone https://huggingface.co/bes-dev/stable-diffusion-v1-4-openvino

# download models
RUN python3 demo.py --num-inference-steps 1 --prompt "test" --output /tmp/test.jpg

This container is using the python:3.9.9-bullseye image and the working directory is set. Next, the Dockerfile installs the same dependencies from earlier in this notebook. Then we add our custom code and pull the dependent repositories.

We've already pushed this image to GHCR, but for posterity, you'd use a command like this to update it:

docker buildx build --platform linux/amd64 --push -t ghcr.io/bacalhau-project/examples/stable-diffusion-cpu:0.0.1 .

Prerequisites​

To run this example you will need:

!command -v bacalhau >/dev/null 2>&1 || (export BACALHAU_INSTALL_DIR=.; curl -sL https://get.bacalhau.org/install.sh | bash)
path=!echo $PATH
%env PATH=./:{path[0]}

Generating an Image Using Stable Diffusion on Bacalhau​

Bacalhau is a distributed computing platform that allows you to run jobs on a network of computers. It is designed to be easy to use and to run on a variety of hardware. In this example, we will use it to run the stable diffusion model on a CPU.

To submit a job, you can use the Bacalhau CLI. The following command passes a prompt to the model and generates an image in the outputs directory.

tip

This will take about 10 minutes to complete. Go grab a coffee. Or a beer. Or both. If you want to block and wait for the job to complete, add the --wait flag.

Furthermore, the container itself is about 15GB, so it might take a while to download on the node if it isn't cached.

%%bash --out job_id
bacalhau docker run ghcr.io/bacalhau-project/examples/stable-diffusion-cpu:0.0.1 --id-only -- python demo.py --prompt "First Humans On Mars" --output ../outputs/mars.png
%env JOB_ID={job_id}

Running the commands will output a UUID that represents the job that was created. You can check the status of the job with the following command:

%%bash
bacalhau list --id-filter ${JOB_ID}

Wait until it says Completed and then get the results.

To find out more information about your job, run the following command:

%%bash
bacalhau describe ${JOB_ID}

If you see that the job has completed and there are no errors, then you can download the results with the following command:

%%bash
rm -rf results && mkdir -p results
bacalhau get $JOB_ID --output-dir results

After the download has finished you should see the following contents in the results directory:

%%bash
ls results/outputs

mars.png

import IPython.display as display
display.Image("results/outputs/mars.png")

png