Deploy a Serverless Stable Diffusion API with SST Ion, AWS Lambda, and Amazon Bedrock

In this post, I’ll walk you through how to deploy a Stable Diffusion 3 API using SST Ion, AWS Lambda, and Amazon Bedrock. This API will enable you to generate images from prompts using a fully managed serverless infrastructure, without the need for a dedicated server.

Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs) from leading AI companies like AI21 Labs, Anthropic, Cohere, Meta, Mistral AI, Stability AI, and Amazon through a single API, along with a broad set of capabilities you need to build generative AI applications.

Prerequisites

You will need an AWS account before getting started. You will also need to configure the AWS CLI and your AWS credentials to follow along.

Disclaimer

This example deploys a bare bones public API. For production environments, consider securing your API with IAM principals and Cognito User Pools for web and mobile clients, or AWS SigV4 signed requests for server clients.

Make sure to review the Amazon Bedrock on-demand pricing before deploying this example to production.

Image Generation Demo

Step 1: Set Up Your Project

Start by initializing a new directory and setting up the project using Yarn and SST Ion.

% mkdir serverless-stable-diffusion && cd serverless-stable-diffusion && yarn init -y

yarn init v1.22.19
warning The yes flag has been set. This will automatically answer yes to all questions, which may have security implications.
success Saved package.json
✨  Done in 0.03s.

Now, initialize SST:

% npx sst@latest init

   ███████╗███████╗████████╗
   ██╔════╝██╔════╝╚══██╔══╝
   ███████╗███████╗   ██║
   ╚════██║╚════██║   ██║
   ███████║███████║   ██║
   ╚══════╝╚══════╝   ╚═╝

>  JS project detected. This will...
   - use the JS template
   - create an sst.config.ts

✓  Template: js

✓  Using: aws

✓  Done 🎉

Step 2: Configure Your SST Project

Add the following code to your SST configuration file to set up the API Gateway. Note that this function requires the bedrock:InvokeModel permission to use Amazon Bedrock. Note that we set a higher timeout of 60 seconds to accommodate the longer processing time of the Stable Diffusion model:

/// <reference path="./.sst/platform/config.d.ts" />

export default $config({
  app(input) {
    return {
      name: "serverless-stable-diffusion",
      removal: "remove",
      home: "aws",
    };
  },
  async run() {
    const api = new sst.aws.ApiGatewayV2("StableDiffusionApi");
    api.route("POST /stableDiffusion", {
      handler: "index.stableDiffusion",
      timeout: "60 seconds",
      permissions: [{ actions: ["bedrock:InvokeModel"], resources: ["*"] }],
    });
  },
});

Step 3: Request Model Access

To use Stable Diffusion 3 with Amazon Bedrock, you will need to request access to the model by visiting the Amazon Bedrock Console. Note that this model is only available in the us-west-2 region at the time of writing.

Step 4: Install Dependencies

yarn add -D @aws-sdk/client-bedrock-runtime

Step 5: Add the Image Generation Handler Function

Add an index.ts file to the root directory with the code below. This example uses Amazon Bedrock to generate images from the user's prompt.

import {
  BedrockRuntimeClient,
  InvokeModelCommand,
} from "@aws-sdk/client-bedrock-runtime";

const bedrockClient = new BedrockRuntimeClient({
  region: "us-west-2",
});

export async function stableDiffusion(_evt: any) {
  try {
    const { prompt } = JSON.parse(_evt.body);
    if (!prompt) {
      return {
        statusCode: 500,
        body: JSON.stringify({ message: "No prompt provided" }),
      };
    }
    // Input for the Bedrock Stable Diffusion model
    const command = new InvokeModelCommand({
      modelId: "stability.sd3-large-v1:0",
      contentType: "application/json",
      body: new TextEncoder().encode(JSON.stringify({ prompt: prompt.slice(0, 100) })),
    });

    // Call the Bedrock service and parse the response
    const response = await bedrockClient.send(command);
    // Parse the output body
    const outputBody = JSON.parse(new TextDecoder().decode(response.body));
    // Extract the base64 image string
    const imageBase64 = outputBody.images[0];
    return {
      statusCode: 200,
      headers: {
        "Content-Type": "image/jpeg",
        "Access-Control-Allow-Origin": "*",
      },
      body: imageBase64,
      isBase64Encoded: true,
    };
  } catch (error: any) {
    return {
      statusCode: 500,
      body: JSON.stringify({ message: "Error generating image", error }),
    };
  }
}

Step 6: Deploy Your API

Run the following command to deploy your API:

% npx sst deploy --stage production

SST 3.0.118  ready!

➜  App:        serverless-stable-diffusion
   Stage:      production

~  Deploy

✓  Complete
   StableDiffusionApi: https://9gc2wjo7ld.execute-api.us-west-1.amazonaws.com

Step 7: Test Your API

You can now test your API by sending a POST request to the endpoint provided in the output.

% curl 'https://9gc2wjo7ld.execute-api.us-west-1.amazonaws.com/stableDiffusion' \
  --data-raw '{"prompt":"a dog sitting on a beach"}'

/ns30+XMcsTfuiesbf3fbnI/Ol8QJAkUfiTRo/kcgXMYB+VvXHb/...........(truncated base64 image data)...

And that's it! You've successfully deployed a serverless SD Image Generation API using SST Ion, AWS Lambda, and Amazon Bedrock! 🚀🚀🚀