Generative AI is a fantastic tool to use to quickly create images based on prompts.
One of the issues with some of these platforms is that they don’t actually store the images in a way that makes them easier to retrieve after they’ve been created. Oftentimes you have to make sure to save it immediately after the process is completed, otherwise, it's gone. Luckily, Stability offers an API that can be used to programmatically generate images, and Tigris is the perfect solution to store those images for retrieval.
In this article, you’ll learn how to deploy an app to Fly.io that will allow you to generate images using the Stability API and automatically store them in a Tigris bucket.
The Stability AI Tigris Database
Let’s take a look at what you’ll be deploying. There are several key components of the project:
- A Next.js app that the user interacts with
- An API endpoint (which is part of the Next app) that processes jobs.
- A background job that periodically polls for new jobs.
- A Postgres database to store jobs
- A Tigris bucket to store the generated images.
The Next.js app
The first part of the project is a Next.js project that contains a single page
that users will interact with. There is a simple form that accepts a prompt and
image dimensions. This form uses server actions to store the request in the
jobs
table of a Fly Postgres database. Each grid item will periodically poll
the table to check on the execution status of each job.
The API processing endpoint
The Next project also contains a single API endpoint that is used to execute jobs against the Stability API before storing the results in a Tigris bucket. This allows for a queue-like structure where jobs can be processed asynchronously.
This endpoint does much of the heavy lifting to make this app possible. Let’s step through what happens when it’s called.
It will start by checking to see if there are any jobs with a status of
pending (0)
:
let res = await db
.select()
.from(jobs)
.where(eq(jobs.status, 0))
.limit(1)
.execute();
If a job is found, the status is set to in progress (1)
. This prevents other
executions from processing a job twice.
await db.update(jobs).set({ status: 1 }).where(eq(jobs.id, job.id)).execute();
Next, the prompt and image dimensions are sent to the Stability API for generating an image. The base64 encoded image is returned in the response from Stability.
const engineId = "stable-diffusion-v1-6";
const apiHost = process.env.API_HOST ?? "https://api.stability.ai";
// Request an image from Stability
const stabilityRes = await fetch(
`${apiHost}/v1/generation/${engineId}/text-to-image`,
{
method: "POST",
headers: {
"Content-Type": "application/json",
Accept: "application/json",
Authorization: `Bearer ${process.env.STABILITY_API_KEY}`,
},
body: JSON.stringify({
text_prompts: [
{
text: job.prompt,
},
],
cfg_scale: 7,
height: job.height,
width: job.width,
steps: 30,
samples: 1,
}),
}
);
let rb = await stabilityRes.json();
if (!rb.artifacts) {
throw new Error(`${rb.name} -- ${rb.message} -- ${rb.details}`);
}
Then we can take that image and upload it to Tigris using the AWS SDK before
setting the job to done (2
.
let artifact = rb.artifacts[0];
if (artifact.finishReason == "SUCCESS") {
let imgdata = artifact.base64;
var buf = Buffer.from(imgdata, "base64");
const upload = new Upload({
params: {
Bucket: process.env.BUCKET_NAME,
Key: `${job.id}.png`,
Body: buf,
},
client: S3,
queueSize: 3,
});
// Upload the file to Tigris
await upload.done();
await db.update(jobs).set({ status: 2 }).where(eq(jobs.id, job.id)).execute();
}
The background job
Using node-cron
, a simple background job is used to poll the API endpoint
from the Next.js app. When polled, that endpoint will handle the next job in the
list. This is run as a secondary process in Fly using concurrently
to avoid
needing unnecessary infrastructure, keeping the project isolated to a single
container. The following diagram demonstrates what’s performed in the background
job:
- The background job polls the API endpoint in Next regularly.
- When a job is detected, the API will set the status to
in progress
. - Next will then dispatch a message to the Stability API, which will respond with a base64 encoded image when processing is complete.
- That image will be stored in a Tigris bucket.
- The database record is set to
complete
.