Terraform and AI Integration for Stable Diffusion
A Serverless Fine-Tuning Pipeline for Personalized AI Images
A consumer photo app startup wanted users to upload their own photos and get back images generated in their likeness, which meant fine-tuning an actual model per user, not just calling a hosted image API. I built a serverless AWS pipeline that takes a batch of a user's photos, runs a SageMaker fine-tuning job in the background, and stores the resulting personal model, plus a separate CloudFormation-provisioned GPU box running the Stable Diffusion web UI for interactive generation and testing. The pipeline runs async end to end so the mobile client gets an immediate response while the GPU work happens off to the side.
The problem
A consumer photo app startup wanted a feature that was popular at the time: a user uploads a batch of their own photos, and the app can then generate new AI images of them. That is not a single prompt to a hosted image API, it needs an actual model fine-tuned on that person's face, the same idea behind Dreambooth-style personalization. A generic image-generation API call cannot produce a personal likeness on demand, the fine-tuning step itself needs a GPU and real compute time, on the order of minutes, not the length of an HTTP request. Running a GPU instance per user around the clock is not something a startup can afford, so the fine-tuning had to run as an on-demand background job that a mobile app could trigger and walk away from, not a live synchronous service.
Constraints
- ·Built end to end by a single engineer, covering both the serverless pipeline and the GPU infrastructure.
- ·The client was a mobile app sending a batch of photos as base64 over a REST API, so the upload path had to work inside typical API Gateway and Lambda payload and timeout limits.
- ·Fine-tuning needs a GPU, but a GPU instance running continuously was not affordable, so the training path had to be triggered on demand rather than always-on.
- ·Primary infrastructure lives in eu-central-1.
Architecture
Two subsystems, traced separately from source. First, a serverless fine-tuning pipeline: API Gateway accepts a batch of user photos, an Upload Lambda stores them in S3 and starts a Step Functions execution, which invokes a SageMaker endpoint for the actual fine-tune, waits a fixed duration, then saves the resulting model back to S3, all keyed by username. Second, an interactive Stable Diffusion web UI on a GPU EC2 instance, provisioned by a CloudFormation template that installs AUTOMATIC1111's stable-diffusion-webui plus the AWS Stable Diffusion extension as a systemd service, with its lifecycle (start, stop, or a fresh relaunch from a saved AMI) scripted through Lambda and SSM Run Command rather than left running.
- ›Amazon API Gateway: accepts the photo batch from the mobile client
- ›Upload_s3 Lambda: decodes and stores photos in S3 under a per-username prefix, starts the Step Functions execution
- ›AWS Step Functions state machine: invoke SageMaker, fixed wait, save model (3 states)
- ›InvokeSageMakerEndpointLambda: calls a named SageMaker endpoint to run the fine-tune
- ›SaveModelToS3Lambda: writes the resulting custom model to S3 under the user's prefix
- ›Amazon S3: single bucket, per-username prefixes for both source photos and the resulting model
- ›AWS CloudFormation (ec2.yaml): provisions the GPU EC2 instance running Stable Diffusion WebUI + the AWS extension
- ›ec2-ssm-lambda.py / model_creation_lambda.py: EC2 lifecycle scripts over SSM Run Command, and AMI-based fresh instance creation
How one request flows
- 01The mobile client captures a batch of the user's photos and POSTs them, base64-encoded, through API Gateway to the Upload_s3 Lambda.
- 02Upload_s3 decodes each photo and writes it to S3 under a per-username prefix (base_image/).
- 03Upload_s3 starts a Step Functions execution, passing the bucket name, image keys, and username.
- 04The state machine invokes InvokeSageMakerEndpointLambda, which sends the batch to a SageMaker endpoint to run the fine-tune.
- 05The state machine waits a fixed 300 seconds, then calls SaveModelToS3Lambda, which writes the resulting model to S3 under the same username's prefix (custom_model/).
- 06Separately, an on-demand GPU EC2 instance running Stable Diffusion WebUI is available for interactive generation, provisioned from a CloudFormation template.
- 07That instance's lifecycle, or a full fresh relaunch from a saved AMI, is scripted through Lambda and SSM Run Command rather than left running around the clock.
Engineering decisions
Async fine-tuning through Step Functions, not a blocking Lambda
What. Moved the SageMaker fine-tune call out of the API-facing Lambda and into a separate Step Functions execution.
Why. API Gateway and Lambda both cap out well under the minutes a fine-tune needs, so the upload Lambda starts the state machine and returns immediately instead of holding the connection open.
Tradeoff. Adds a state machine and two more Lambdas to reason about. Worth it since the alternative is a request that always exceeds the timeout.
A fixed wait instead of a completion check
What. The state machine waits a fixed 300 seconds after invoking SageMaker, then calls the save step unconditionally.
Why. Keeps the state machine to three states with no polling logic to write and maintain.
Tradeoff. This is the pipeline's weakest point. Not chosen was polling the SageMaker job status or waiting on a completion callback, either would confirm the fine-tune actually finished before saving. As built, a job that runs past 300 seconds can hit the save step before a result exists.
Fresh EC2 instances from a saved AMI, not a mutated long-lived box
What. model_creation_lambda.py launches a new EC2 instance from a known-good AMI and EBS snapshot rather than reconfiguring one instance repeatedly.
Why. State does not drift between runs, there is one script that recreates the box instead of manual changes nobody tracked.
Tradeoff. A fresh boot costs the wait for instance status checks to pass, slower than reusing a warm box. Chosen for repeatability over speed.
SSM Run Command instead of SSH for scripted EC2 work
What. Automation scripts drive the EC2 instance through SSM's send_command, not a key pair and an open management port.
Why. No SSH key to distribute or rotate, and no inbound port 22 needed purely for scripted jobs.
Tradeoff. Not chosen was scripting over SSH, which would need a key and an open port just for automation. The one place SSH does stay open is the interactive web UI box, a deliberate, separate exception.
Per-user state addressed by username, no user table
What. Every S3 key, both the source photos and the resulting custom model, is prefixed with the username string coming from the API request.
Why. Simple and fast to build for a single-user-per-run flow, no database or generated ID needed.
Tradeoff. Not chosen was a real identity layer with generated user IDs and access checks. That is a genuine gap if usernames can collide or are not authenticated upstream, acceptable for the scope this was built for.
A stock CloudFormation template adapted, not IaC from scratch
What. The interactive Stable Diffusion box is provisioned from a known AWS workshop CloudFormation template, with the security group and instance user data adapted to install stable-diffusion-webui and the AWS extension as a systemd service.
Why. Reuses a template AWS already validated across every region's AMI mapping, so the work is only the parts specific to this box.
Tradeoff. Less bespoke than a from-scratch build. The right call for an interactive/demo box built by one engineer alongside the main serverless pipeline.
What changed along the way
- →The first version of the upload Lambda hardcoded the bucket name and the state machine ARN directly in code, a later version (save_s3_lambda_v1.py) read the bucket from an environment variable instead.
- →ec2-ssm-lambda.py defines start_ec2_instance and stop_ec2_instance, but the calls to them ended up commented out. The script runs commands against a box that is already on, not a full automated start-stop cycle.
- →Two versions of the same upload flow existed side by side (aws/Upload_s3.py and save_s3_lambda_v1.py) while the approach to bucket naming and payload shape was still being settled.
Security and reliability
- ·SSM Run Command is used for scripted EC2 operations, so there is no standing SSH key or open management port for that path.
- ·The interactive Stable Diffusion box is CloudFormation-provisioned, but its security group opens both SSH and the web UI port to 0.0.0.0/0, acceptable for a short-lived demo box, a real gap if left running.
- ·A dedicated IAM role and instance profile back the automation Lambdas rather than reusing a broad default role.
- ·There is no retry or dead-letter handling around the Step Functions execution, a failed SageMaker call or a slow fine-tune is not automatically retried or flagged.
Metrics
What shipped
- ✓Delivered a serverless pipeline, API Gateway, Lambda, S3, Step Functions, SageMaker, that turns a batch of a user's photos into a per-user fine-tuned Stable Diffusion model.
- ✓Delivered a CloudFormation template that stands up a GPU EC2 instance running Stable Diffusion WebUI with the AWS extension for interactive generation.
- ✓Wrote Lambda scripts to control EC2 lifecycle over SSM Run Command and to launch a fresh instance from a saved AMI, so the GPU box does not need hand configuration on every run.
- ✓Kept all per-user state, source photos and the resulting model, addressable in S3 by username with no separate database.
Lessons learned
- ·A fixed wait state is a simple stand-in for a completion check, but it only holds up while the job reliably finishes inside the window. A production version needs to poll SageMaker or listen for a completion event instead.
- ·Booting a fresh EC2 instance from a saved AMI keeps state from drifting between runs, a better default than repeatedly reconfiguring one long-lived box.
- ·SSM Run Command removes the need for a standing SSH key or open port for scripted work, real interactive access should stay a separate, deliberate exception.
- ·Moving a hardcoded bucket name into an environment variable paid off almost immediately in the next iteration of the same Lambda.