Stable Diffusion on a GPU
This example tutorial demonstrates how to use Stable Diffusion on a GPU and run it on the Bacalhau demo network. 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.
TL;DR
Prerequisite
To get started, you need to install the Bacalhau client, see more information here
Quick Test
Here is an example of an image generated by this model.
Development
This stable diffusion example is based on the Keras/Tensorflow implementation. You might also be interested in the Pytorch oriented diffusers library.
Installing dependencies
When you run this code for the first time, it will download the pre-trained weights, which may add a short delay.
Based on the requirements here, we will install the following:
Testing the Code
We have a sample code from this the Stable Diffusion in TensorFlow/Keras repo which we will use to check if the code is working as expected. Our output for this code will be a DSLR photograph of an astronaut riding a horse.
When you run this code for the first time, it will download the pre-trained weights, which may add a short delay.
When running this code, if you check the GPU RAM usage, you'll see that it's sucked up many GBs, and depending on what GPU you're running, it may OOM (Out of memory) if you run this again.
You can try and reduce RAM usage by playing with batch sizes (although it is only set to 1 above!) or more carefully controlling the TensorFlow session.
To clear the GPU memory we will use numba. This won't be required when running in a single-shot manner.
Write the Script
You need a script to execute when we submit jobs. The code below is a slightly modified version of the code we ran above which we got from here, however, this includes more things such as argument parsing argument parsing to be able to customize the generator.
For a full list of arguments that you can pass to the script, see more information here
Run the Script
After writing the code the next step is to run the script.
As a result, you will get something like this:
The following presents additional parameters you can try:
python main.py --p "cat with three eyes
- to set promptpython main.py --p "cat with three eyes" --n 100
- to set the number of iterations to 100python stable-diffusion.py --p "cat with three eyes" --b 2
to set batch size to 2 (№ of images to generate)
Containerize Script using Docker
Docker is the easiest way to run TensorFlow on a GPU since the host machine only requires the NVIDIA® driver. To containerize the inference code, we will create a Dockerfile
. The Dockerfile is a text document that contains the commands that specify how the image will be built.
The Dockerfile leverages the latest official TensorFlow GPU image and then installs other dependencies like git
, CUDA
packages, and other image-related necessities. See the original repository for the expected requirements.
See more information on how to containerize your script/app here
Build the container
We will run docker build
command to build the container;
Before running the command replace following:
hub-user with your docker hub username, If you don’t have a docker hub account follow these instructions to create a Docker account, and use the username of the account you created
repo-name with the name of the container, you can name it anything you want
tag this is not required but you can use the latest tag
In our case:
Push the container
Next, upload the image to the registry. This can be done by using the Docker hub username, repo name or tag.
In our case:
Running a Bacalhau Job
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
To submit a job run the Bacalhau command with following structure:
export JOB_ID=$( ... )
exports the job ID as environment variableThe
--gpu 1
flag is set to specify hardware requirements, a GPU is needed to run such a jobThe
--id-only
flag is set to print only job idghcr.io/bacalhau-project/examples/stable-diffusion-gpu:0.0.1
: the name and the tag of the docker image we are using-- python main.py --o ./outputs --p "meme about tensorflow"
: The command to run inference on the model. It consists of:main.py
path to the script--o ./outputs
specifies the output directory--p "meme about tensorflow"
specifies the prompt
The Bacalhau command passes a prompt to the model and generates an image in the outputs directory. The main difference in the example below compared to all the other examples is the addition of the --gpu X
flag, which tells Bacalhau to only schedule the job on nodes that have X
GPUs free. You can read more about GPU support in the documentation.
This will take about 5 minutes to complete and is mainly due to the cold-start GPU setup time. This is faster than the CPU version, but you might still want to grab some fruit or plan your lunchtime run.
Furthermore, the container itself is about 10GB, so it might take a while to download on the node if it isn't cached.
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 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 describe
.
Job download: You can download your job results directly by using bacalhau 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.
After the download has finished you should see the following contents in results directory
Viewing your Job Output
Now you can find the file in the results/outputs
folder:
Last updated