Stable Diffusion is the name of a Deep Learning model created by stability.ai that allows you to generate images from their description. In short, you feed a textual scene description to the model and it returns an image that fits that description. What is cool about it is that it can generate very artistic and arguable beautiful images resembling pieces of art like the following
In the aws-sagemaker-stable-diffusion repo you will find everything needed in order to spin-up your own personal public endpoint with a Stable Diffusion model deployed using AWS SageMaker to show your friends 😎
- Your own Stable Diffusion endpoint with AWS SageMaker
Assuming you already have
- Define the bucket name (where the model will be sent) in
# Here you will need to provide your huggingface credentials in order to confirm # the model license has been accepted INSTALL_TOOLING=true bash setup.sh cd terraform/ && terraform init && terraform apply && cd ../ cd sagemaker/ && pip install -r requirements && bash zip-model.sh && cd ../ cd lambda/sd-public-endpoint/ && pip install -r requirements.txt && bash deploy.sh && cd ../../
The endpoint that will appear in the output can be used for the inferences. And here's how to use it
<endpoint>/-> Generates a random description and feeds it into the model
<endpoint>/default-> Uses the default description "a photo of an astronaut riding a horse on mars"
<endpoint>/<description>-> Uses the
<description>as input for the model. You can use spaces here.
For this repo, you will need these tools:
For python, it's recommended to use
pyenv, which allows you to install several versions of python at the same time with simple commands like this:
pyenv install 3.9.13
For the rest, you can use a tool called
asdf, which allows basically the same but for several other tools. If you have it installed, you can install the rest with just
To be able to clone the repo with the ML model (
stable-diffusion-v1-4/), we will also need
git-lfs (that allows versioning of large files). The model is defined as a submodule of this repo and after attempting to clone the submodule, you will be requested to provide your huggingface account credentials in order to clone it. This it to confirm that you have accepted the license required to have access to the model weights.
pigz is also necessary as it allows parallel gzip compression. This makes the model packaging for SageMaker run much faster.
So, to do all of this, the
setup.sh script can be used for that. It can be used on Mac OS X or Debian-based linux distros. If you already have
pyenv installed and want to install the rest, just run
INSTALL_TOOLING=true bash setup.sh
However, if you already have the tooling, and just want
pigz with the submodule cloned, please run:
To better understand how to use the model and the inference code we'll use later on, we first should try to run the model locally. Based on the official release page you should be able to access the model weights on their huggingface page. Here we will use version
On the page we can see code samples showing how we should be able to run an inference on the model. There you will see examples on how to make a standard inference, however it's recommended to have a NVIDIA GPU with at least 10GB of VRAM. To allow my weaker hardware to make an inference, I choose to test the configuration using
float16 precision instead, which is the code that you can find at
ℹ️ There is also a
local/code-mac.pyavailable for those that want to try it out on a mac. However, bare in mind that one inference can take several minutes (~30 min on a M1 Pro), whereas in a GPU it might take 5 seconds using the low VRAM configuration on a RTX 3070
For my local computer, I used
nvidia-docker to allow myself to run the code inside a container. This way I don't have to worry about installing the right version of CUDA (relative to pytorch) on my machine.
So, going to the fun part, to make an inference, you just need to run
cd local && bash build-run.sh.
All script in this repo assume that they are being executed from their own directories
The whole execution can take several minutes on the first time in order to build the container image with the model inside. And after everything has finished, you should be able to see image generated at
You might see something like this on the default description of
"a photo of an astronaut riding a horse on mars. VFX, octane renderer":
And you can change the description by changing the content of the
prompt variable in
Cool, so now that we have the inference code working, is time to put this in the cloud to, ultimately, make it available to others. Let's start by talking about the resource we'll need and the overall architecture.
Firstly, about how the ML model will be executed, although we could simply spin-up an EC2 instance and attach a web-server to it to receive the images as requests, we will use AWS SageMaker, which allows us to do exactly that and much more in a managed way (it means that several components, e.g. the web-server implementation, will be managed by AWS). SageMaker will manage the EC2 instance with a GPU and will give us a private endpoint to interact with the model. Since we are using huggingface, both have a nice integration and you can learn more about it from the huggingface docs about SageMaker or the AWS docs about huggingface 🤝.
However, one thing to notice is that SageMaker will provide a private endpoint only accessible if you have access to the AWS account associated with the resource. Since we want a public endpoint, we need to put on top of it a lambda that will forward requests from an API Gateway. This combination follows the recommended approach from an official AWS blog post.
⚠️ One import limitation of API Gateway to notice is the hard limit of 30 seconds timeout on the request. Because of this, the same low VRAM configuration (
float16precision) was needed in order to guarantee a lower response time.
Putting it all together, the architecture goes like this:
Showing all resources needed and from which directory each resource is managed. Here is a brief description:
||Has the lambda + API Gateway implementation using a framework called chalice|
||Contains the python code to manage the SageMaker Model, Endpoint, and the custom inference code. It also has the script used to pack and send the model to an S3 bucket|
||Manages the S3 bucket itself and the IAM roles required by the lambda code (to access the SageMaker endpoint) and for the Sagemaker endpoint (to access the model on the S3 bucket)|
Before anything, we need to deploy the roles and the S3 bucket mentioned in order to set the stage to the rest of the resources. You can configure the bucket name used in the file
terraform/variables.tf. After that, you can run (assuming that you already have the access to your AWS account properly configured):
cd terraform terraform init terraform apply
In total 11 resources should be created. In this process, a couple of
.txt will be created. Those files will be used by other parts of this repo to get the value of AWS ARNs and the bucket name used
A brief explanation of the roles is as follows:
|Role||Name on IAM||Description|
|SageMaker Endpoint Access||
||This role basically allows the lambda to execute properly (with the
|SageMaker Full Access||
||This one has the
With the IAM roles and the S3 bucket in place, now it the time to create the SageMaker endpoint itself. For that, we use the python package from AWS which allow us to use Transformers models with the huggingface SDK.
However, since we will use diffusion type of model, and not transformers, the default inference code will need to be overwritten to use the
diffusers package (also from huggingface), just like we do in the
In order to overwrite it, the package readme has some general information about it, and also there is an example in this jupyter notebook. We are doing what is necessary via the files inside
sagemaker/code, which has the inference code following SageMaker requirements, and a
requirements.txt, that has the necessary dependencies that will be installed when the endpoint gets created
With the inference code ready, it's time to ship the model with the code to the created S3 bucket (as the SageMaker endpoint will access the model via this bucket). For this, run:
cd sagemaker/ bash zip-model.sh
Keep in mind that this process will send around 4.2GB of data to S3. Just remember that cost you are willing to pay 😉
With that, we should have the what we need in the bucket. Now, let's create the SageMaker endpoint. To manage the endpoint, as stated before, we use the SageMaker toolkit for huggingface. As such, we have three python scripts that serve as functions to create, use and delete the endpoint.
So let's go ahead and create the endpoint with
pip install -r requirements # To install the sagemaker pkg python sagemaker-create-endpoint.py
After this the endpoint will be created, and also a file
endpoint-name.txt so that the other python scripts can keep its reference. When the endpoint is ready, you will be able to see it in the AWS console like the following image:
Now we should be ready to run an inference. To do it, execute:
And we should get a cool image like this one in
And it's nice that we just this amount of code, now we have a deployment of a ML model server with metrics, logs, health checks, etc. available, as we can see in the following images.
Yes, to get that cool image from Darth Vader, lots of attempts were required 😅
So far so good, we have a python script capable of make inference on a the Stable Diffusion model, but this python script interacts directly to your AWS resources. To achieve public access, it's time to work on the Lambda + API Gateway
For this part we had couple of options to choose from, like serverless and Zappa, both being serverless frameworks that would allow us to deploy a serverless app (fulfilling the lambda + API Gateway combo).
However, for this project we will go with chalice, which is a serverless framework for python designed from AWS itself that resembles the API design of Flask. Since I never heard about it before I'll give it a try.
The code for the serverless app it's all inside
lambda/. To manage the
chalice app, do
cd lambda pip install -r requirements.txt cd sd-public-endpoint
This way, you will have the chalice tool installed and will be able to deploy the app with
pip install -r requirements.txt bash deploy.sh
The script will get the proper IAM role and deploy the app.
ℹ️ You need to install the dependencies of the serverless app because the dependencies are packaged locally and sent with the app code when publishing the lambda
⚠️ Just a reminder about that API Gateway has a non-configurable 30 seconds timeout that might impact your inferences depending on the configuration of the Model. However it seems there is a workaround for API Gateway
And now you should have a public endpoint to access the Model, you should be able to use the app like this:
||The home page, just for testing|
||You can type any text replacing
||And that were the fun actually begins!! If you leave the
To clean up everything (and avoid unwanted AWS costs 💸), you can do:
# Delete the serverless app cd lambda/sd-public-endpoint/ chalice delete # Delete the SageMaker resources cd ../../sagemaker python sagemaker-delete-endpoint.py bash delete-model.sh cd ../terraform/ terraform destroy
Now you will only need to delete the related CloudWatch Log groups via the console. Those will be
So here create our own personal cloud deployment of the popular Stable Diffusion ML model and played a little with several cloud resources in the process.
The architecture works fine with play around, but an asynchronous implementation might become necessary to allow inferences that take more than 30 seconds, and also for batch processing.