Image Generation and Storage using Stable Diffusion, AWS Bedrock, and TypeScript

In a previous post, we looked at AWS Bedrock and how to generate text using it. However, in this post, we’re going to take that one step further and look at using AWS Bedrock to generate images.

In this tutorial, we’re going to build a simple CDK stack that provisions a Lambda function that will send a request to AWS Bedrock (via the AWS SDK) to generate an image using the Stable Diffusion model. We’ll then take that image and store it in an S3 bucket before returning the URL to the requester so they can download and view the image.

Configuring Bedrock

Before we can jump in and start building our application in the CDK though, we need to make sure AWS Bedrock is configured for our account. At the time of writing Bedrock is only available in a few AWS regions and the models inside of Bedrock aren’t automatically enabled for all accounts.

Instead, we need to manually go and request access to each of the models we want to use in our chosen region. For me, I’ll be using the `us-east-1` region but you can use another region if you’d prefer as long as it has access to the Stable Diffusion model.

Once you have Bedrock configured in your target region and you have access to the Stable Diffusion model in that region, we’re ready to get started with the tutorial.

Setting up Our CDK Stack

To follow along with this tutorial, you’ll need to have a CDK stack, you can use an existing one if you’d like or you can create a new one by running the command `cdk init app --language typescript`. Once you have your CDK stack you need to ensure you’re deploying it to the same region that Bedrock was configured in, this can be configured inside the `bin` directory.

1#!/usr/bin/env node
2import "source-map-support/register";
3import * as cdk from "aws-cdk-lib";
4import { YourStack } from "../lib/your-stack";
5
6const app = new cdk.App();
7new YourStack(app, "YourStack", {
8  env: {
9    region: "YOUR_REGION",
10  },
11});

*NOTE: You don’t necessarily need to set the CDK region to be the same as the Bedrock region as you can change the region for the Bedrock SDK client independently. But, for the sake of simplicity and to avoid using multiple regions, we’re going to change the entire stack’s region to align with Bedrock.*

Defining Resources

Once your CDK stack region and Bedrock are configured, we’re ready to get started with defining the resources we want to provision.

S3 Bucket

To start, we’re going to define our S3 bucket that will be used to store the images that are generated by Bedrock so that we’re able to download them and view them later on. To define the S3 bucket add the below code to your stack definition file in the `lib` directory.

1const s3Bucket = new Bucket(this, "ImageBucket", {
2  removalPolicy: RemovalPolicy.DESTROY,
3  publicReadAccess: true,
4  blockPublicAccess: BlockPublicAccess.BLOCK_ACLS,
5  cors: [
6    {
7      allowedHeaders: ["*"],
8      allowedMethods: [HttpMethods.GET],
9      allowedOrigins: ["*"],
10      exposedHeaders: [],
11      maxAge: 3000,
12    },
13  ],
14});

With this code, we define a new S3 bucket called `ImageBucket` and grant `publicReadAccess` to it so anyone can download the files from it as well as configure CORS to only allow `GET` requests to it.

Lambda

With our S3 bucket now defined, let’s define our lambda function next. To do this, add the below code under the code we just added for the S3 bucket.

1new NodejsFunction(this, "ImageLambda", {
2  entry: "resources/image-lambda.ts",
3  handler: "handler",
4  runtime: Runtime.NODEJS_18_X,
5  timeout: Duration.minutes(3),
6  bundling: {
7    nodeModules: ["@aws-sdk/client-bedrock-runtime", "@aws-sdk/client-s3"],
8  },
9  environment: {
10    S3_BUCKET_NAME: s3Bucket.bucketName,
11  },
12  initialPolicy: [
13    new PolicyStatement({
14      effect: Effect.ALLOW,
15      actions: ["bedrock:*", "s3:*"],
16      resources: [`*`],
17    }),
18  ],
19});

In this code, we define our new `ImageLambda` lambda function as well as configure some other aspects of the function such as its timeout, runtime, and the file where the code is stored.

However, two properties we need to take a closer look at are the `environment` and the `initialPolicy` properties. With the `environment` property, we provide the name of the S3 bucket to the function so that we can use it inside the function to upload files to the bucket. We then use the `initialPolicy` property to configure a basic IAM policy that allows the Lambda function to interact with S3 and Bedrock.

Now, in a production application, we’d make this policy more restrictive so it only grants the necessary permissions to the resources and not all of them. But, as this is just a tutorial, this policy is ideal because it will allow us to experiment with Bedrock without worrying about IAM permissions preventing us from accessing a model or performing an action against one.

Writing the Lambda

To start with creating our lambda function, we need to create a new file to contain the function’s code; to do this, create a new file at `./resources/image-lambda.ts`. Inside this file add the below code to it.

1import {
2  BedrockRuntimeClient,
3  InvokeModelCommand,
4} from "@aws-sdk/client-bedrock-runtime";
5import { PutObjectCommand, S3Client } from "@aws-sdk/client-s3";
6
7const bedrockClient = new BedrockRuntimeClient();
8const s3Client = new S3Client();
9
10export const handler = async () => {
11  const { S3_BUCKET_NAME = "" } = process.env;
12
13  const prompt = "Developer making an application";
14  const fileName = `${prompt.toLowerCase().replaceAll(" ", "_")}.jpg`;
15
16  const input = {
17    body: JSON.stringify({
18      text_prompts: [
19        {
20          text: prompt,
21        },
22      ],
23      cfg_scale: 10,
24      seed: 0,
25      steps: 50,
26    }),
27    accept: "application/json",
28    contentType: "application/json",
29    modelId: "stability.stable-diffusion-xl-v0",
30  };
31
32  const command = new InvokeModelCommand(input);
33
34  try {
35    const response = await bedrockClient.send(command);
36
37    // 1. Get our Base 64 encoded image from the response
38    const { artifacts } = JSON.parse(
39      new TextDecoder().decode(response.body)
40    ) as {
41      result: string;
42      artifacts: {
43        seed: string;
44        base64: string;
45        finishReason: string;
46      }[];
47    };
48
49    const imageBase64 = artifacts[0].base64;
50
51    // 2. Decode the Base 64 encoded image
52    const decodedImage: Buffer = Buffer.from(imageBase64, "base64");
53
54    // 3. Upload the image to S3
55    await s3Client.send(
56      new PutObjectCommand({
57        Bucket: S3_BUCKET_NAME,
58        Key: fileName,
59        Body: decodedImage,
60      })
61    );
62
63    // 4. Create the URL of the image
64    const imageUrl = `https://${S3_BUCKET_NAME}.s3.amazonaws.com/${fileName}`;
65
66    // 5. Return the URL of the image
67    return {
68      statusCode: 200,
69      body: imageUrl,
70    };
71  } catch (e) {
72    return {
73      statusCode: 500,
74      body: e,
75    };
76  }
77};

Let’s now walk through what this code does. To start with, we import the required dependencies for S3 and Bedrock from the AWS SDK before creating S3 and Bedrock clients that we’ll use to send commands to AWS.

Then inside our `handler` function, we define the `prompt` we want to send to the Stable Diffusion model on Bedrock before then creating the `fileName` we’re going to be storing the image as on S3.

We then define an `input` object that contains all of the necessary information that Stable Diffusion and Bedrock need to carry out the request and generate an image for us.

At this point, it’s worth noting a couple of things, first of all, you can see inside the `input` object we define a property called `modelId` with a value of `stability.stable-diffusion-xl-v0`, this is how we tell Bedrock what model we want to use. Secondly, the `body` property inside the `input` object is unique to each model so this is the one we require for Stable Diffusion but for different models on Bedrock, you’d need to update this `body` property with different properties and values.

With the `input` object now defined, we then send the request to Bedrock and transform the response so that we have access to the `base64` image data returned to us. We do this with the below code.

1const { artifacts } = JSON.parse(new TextDecoder().decode(response.body)) as {
2  result: string;
3  artifacts: {
4    seed: string;
5    base64: string;
6    finishReason: string;
7  }[];
8};
9
10const imageBase64 = artifacts[0].base64;
11
12// 2. Decode the Base 64 encoded image
13const decodedImage: Buffer = Buffer.from(imageBase64, "base64");

Then once we have the decoded image buffer, we can upload the image to S3 using the below code and the `S3_BUCKET_NAME` environment variable we passed in when we defined the Lambda function.

1// 3. Upload the image to S3
2await s3Client.send(
3  new PutObjectCommand({
4    Bucket: S3_BUCKET_NAME,
5    Key: fileName,
6    Body: decodedImage,
7  })
8);

Finally, we then define the URL where the image will be available and return the image to the requester using the below code.

1// 4. Create the URL of the image
2const imageUrl = `https://${S3_BUCKET_NAME}.s3.amazonaws.com/${fileName}`;
3
4// 5. Return the URL of the image
5return {
6  statusCode: 200,
7  body: imageUrl,
8};

At this point, we’ve finished writing the lambda function and our CDK stack overall so in the next section let’s look at deploying and testing our CDK stack by generating an image using the prompt we defined and then downloading it from S3.

Testing the Image Generation

Before we can test our image generation, we first need to deploy our CDK stack, we can do this by running the command `cdk deploy` in our terminal. Once our CDK stack has finished deploying, we can move on to testing our code.

To do this, we need to invoke our new Lambda function which can be done in various ways such as using the AWS CLI or dashboard. However, I will be using the VSCode extension “AWS Toolkit” which allows you to manage and invoke your AWS resources from inside VSCode.

After you’ve invoked your Lambda function, you should receive a response with a status code of `200` as well as a URL to the image that was generated by Stable Diffusion and Bedrock. After visiting this URL the generated image should be downloaded to your computer and you should be able to view the image.

So, if you can do all of that then congratulations you have a working project where we can prompt Stable Diffusion to generate an image inside AWS Bedrock and have the generated image uploaded to S3 for us to download.

Closing Thoughts

To recap, in this post, we’ve looked at AWS Bedrock and how we can use it to generate images using the Stable Diffusion model before having those images uploaded to an S3 bucket for us to download and view on our local machines.

If you’re interested in checking out the full example project and CDK stack for this tutorial, you can see it over on my GitHub along with all of my other example CDK projects and AWS tutorials.

I hope you found this post helpful and thanks for reading.