Stable Diffusion on a GPU
This example tutorial demonstrates how to use stable diffusion on a GPU and run it on the Bacalhau 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.
TD;LR
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
:::info 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.
:::tip 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.
:::info 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.
Display the Output Image
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 (No 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.
:::info See more information on how to containerize your script/apphere :::
Build the container
We will run docker build
command to build the container;
Before running the command replace;
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
To submit a job, run the following Bacalhau command:
Structure of the command
Let's look closely at the command above:
bacalhau docker run
: call to bacalhau--gpu 1
: No of GPUsghcr.io/bacalhau-project/examples/stable-diffusion-gpu:0.0.1
: the name and the tag of the docker image we are using../outputs
: path to the outputpython main.py
: execute the script
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.
:::tip 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
To view the file, run the following command:
Display image
To display and view your image run the code below: