Stable Diffusion on a CPU
Introduction
Stable Diffusion is an open-source text-to-image model, which generates images from text. It's a cutting-edge alternative to DALL·E 2 and uses the Diffusion Probabilistic Model for image generation. At the core the model generates graphics from text using a Transformer.
This example demonstrates how to use stable diffusion online on a CPU and run it on the Bacalhau demo network. The first section describes the development of the code and the container. The second section demonstrates how to run the job using Bacalhau.
This model generated the images presented on this page.
TL;DR
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.
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:
A Debian-flavoured Linux (although you might be able to get it working on the newest machines)
Converting Stable Diffusion to a CPU Model Using OpenVINO
First we convert the trained stable diffusion models so that they work efficiently on a CPU with OpenVINO. Choose the fine tuned version of Stable Diffusion you want to use. The example is quite complex, so we have created a separate repository to host the code. This is a fork from this Github repository.
In summary, the code downloads a pre-optimized OpenVINO version of the original pre-trained stable diffusion model. This model leverages OpenAI's CLIP transformer and is wrapped inside an OpenVINO runtime, which executes the model.
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 continue.
Install Dependencies
Note that these dependencies are only known to work on Ubuntu-based x64 machines.
Clone the Repository and Dependencies
The following commands clone the example repository, and other required repositories, and install the Python dependencies.
Generate 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.
When the generation is complete, you can open the generated hello.png and see something like this:
Lets try another prompt and see what we get:
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.
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:
Prerequisites
To run this example you will need Bacalhau installed and running
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.
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.
Structure of the command
Some of the jobs presented in the Examples section may require more resources than are currently available on the demo network. Consider starting your own network or running less resource-intensive jobs on the demo network
export JOB_ID=$( ... ): Export results of a command execution as environment variablebacalhau docker run: Run a job using docker executor.--id-only: Flag to print out only the job idghcr.io/bacalhau-project/examples/stable-diffusion-cpu:0.0.1: The name and the tag of the Docker image.The command to run inference on the model:
python demo.py --prompt "First Humans On Mars" --output ../outputs/mars.png. It consists of:demo.py: The Python script that runs the inference process.--prompt "First Humans On Mars": Specifies the text prompt to be used for the inference.--output ../outputs/mars.png: Specifies the path to the output image.
When a job is submitted, Bacalhau prints out the related job_id. We store that in an environment variable so that we can reuse it later on.
Checking the State of your Jobs
Job status: You can check the status of the job using bacalhau job list:
When it says Completed, that means the job is done, and we can get the results.
Job information: You can find out more information about your job by using bacalhau job describe:
Job download: You can download your job results directly by using bacalhau job get. Alternatively, you can choose to create a directory to store your results. In the command below, we created a directory and downloaded our job output to be stored in that directory.
Viewing your Job Output
After the download has finished we can see the results in the results/outputs folder.
Was this helpful?
