Design an AI image generation service

GPU workers should pull work, not be pushed work: a worker asks for a job only when it has free capacity, so slow jobs never pile up on one GPU. Each worker serves one model (or one family) that stays loaded in GPU memory, because loading 7-15 GB of weights takes 3-15 s and would double the latency of every job. So queues are per model pool: q:sdxl:paid, q:sdxl:free, q:upscaler, and the autoscaler sizes each pool from its queue depth and wait time, not from CPU. Batching improves throughput: four images of the same prompt run as one batch of 4, and the worker can also merge small jobs with the same model and resolution, raising throughput per GPU by roughly 30-50% at the cost of a little latency. Mixed resolutions do not batch well, so standard aspect ratios are bucketed. Style LoRAs (50-300 MB) are cached in GPU or host memory with LRU and loaded per job, and the scheduler prefers a worker that already holds the LoRA. Rare models get a small warm pool or scale to zero with a visible 'warming up' state. Cold start is cut by baking weights into the node image or local NVMe cache instead of pulling from S3 on every start.

On a festival evening demand can exceed capacity for hours, and you cannot add 300 GPUs in minutes. The design must decide who waits. Paid users go into a fast lane that workers always check first, but with a weighted share (for example 80/20) so the free lane is never fully starved. Per-user concurrency caps (free: 1 job in flight, paid: 4-12) stop one user with a script from filling the queue; jobs above the cap wait in a per-user holding list. Inside a lane, ordering by enqueue time is fine, but a user's 10th job should not jump ahead of another user's 1st, so the score adds a small penalty per in-flight job. Admission control protects latency: if the estimated wait for the free lane passes ~3 minutes, new free requests get a clear 'busy, try later' or are offered a cheaper, faster distilled model, instead of joining a queue they will abandon. Abandoned jobs waste GPUs, so when an SSE client disconnects and does not return within 30 s, queued jobs are cancelled. Showing an honest queue position keeps users calm and reduces retries that would add load.

Check the prompt first, because blocking before generation saves GPU money. Use layers: exact and fuzzy blocklists, a fast text classifier trained for each policy area (sexual content involving minors, real-person sexual content, extremist symbols, graphic violence), and for ambiguous prompts a small LLM that also catches obfuscation ('n.u.d.e', transliterated Hindi). The prompt check cannot catch everything, so every output image goes through an image classifier for nudity, gore and minors, plus hash matching (perceptual hashes such as PDQ) against industry lists of known abusive images, before anything is stored or shown. Thresholds are set per policy: for child safety, block at low confidence and accept false positives; for mild violence, allow more. Blocked outputs are never written to user storage; a hash and labels are logged for review, and confirmed CSAM is reported as the law requires. Add provenance: C2PA manifest in the file metadata and an invisible watermark that survives resizing and compression, because metadata is stripped by most social apps. Regulators, including in India, are moving toward requiring labels on synthetic media, so build it in from day one.

Money and retries interact badly. The client may resend POST /generations on a flaky Jio or Airtel connection, workers may crash and retry, and a job may be blocked after generation. Use a reservation pattern with an append-only ledger. The POST carries an Idempotency-Key; a unique index on (user_id, idempotency_key) makes the second request return the first job. Enqueue writes a 'reserve' entry in the same transaction as the job row, so the balance shown already excludes in-flight work and a user cannot start 50 jobs on 10 credits. When images are delivered, a 'commit' entry is written; on failure, cancel or a safety block, a 'release' entry. A unique constraint on (job_id, type) makes each transition happen at most once even if two workers race after a lease expiry. A nightly reconciliation finds reservations older than an hour without a commit or release and releases them. The GPU retry itself is safe because generation has no external side effect until the final write, and the same seed reproduces the same image.

A synchronous request that holds for 10-60 s ties up connections, breaks on mobile network switches and makes retries dangerous. An async job with SSE progress survives reconnects (Last-Event-ID or polling), lets you show queue position and supports cancellation. The cost is more moving parts: a queue, a job table and a progress channel.

Kafka is built for ordered logs, not per-job priority, cancellation and visibility timeouts. SQS gives leases but no priority within a queue and no fairness logic. Redis sorted sets give priority and custom fairness scoring with sub-millisecond ops; durability comes from Postgres as the source of truth plus a sweeper that rebuilds the queue.

A hosted image API is fastest to launch and has no fleet to run, but cost per image is several times higher and you cannot customise models. Reserved cloud GPUs give the lowest cost at a steady base load; on-demand or spot capacity covers peaks. Image jobs are short and restartable, so spot interruptions only cost a retry.

Distilled or few-step models generate in under a second but with less detail. Offering 'fast draft' (4-8 steps) and 'final' (30 steps, plus upscale) lets users explore cheaply and spend GPU time only on images they like.

At 10× load the queue grows without bound. Fix with admission control and honest wait estimates, cheaper fast models as a fallback, scheduled pre-scaling before known peaks (Diwali, New Year), and spot capacity across several regions.

Hundreds of TB per year and heavy sharing traffic. Expire unsaved free-tier images, store WebP not PNG, generate thumbnails for the gallery, and keep popular shared images hot in the CDN.

Leases expire and the sweeper re-enqueues the job; the same seed reproduces the image; credits stay reserved until delivery. Workers handle SIGTERM by finishing the current batch if it can complete within the 2-minute spot warning, otherwise they release the lease immediately.

Fail closed: no image is delivered without an output check. Jobs wait in the post-processing queue and users see 'finishing up'. This is a conscious choice: a few minutes of delay is far cheaper than serving one illegal image.

Postgres has every queued job; the sweeper rebuilds the sorted sets from jobs with status queued, ordered by created_at. Duplicate pops are harmless because the lease update in Postgres is a conditional write that only one worker can win.