Intermediate

Step 1: Image Generation API

In this lesson you will build the core backend — a FastAPI service that accepts text prompts and returns AI-generated images. You will integrate both Stability AI and Replicate as providers, handle generation parameters, and store images locally.

The Image Service

Create services/image_service.py. This module abstracts the image generation providers behind a single interface so you can switch between them easily.

# services/image_service.py
import os
import uuid
import base64
from datetime import datetime
from pathlib import Path
from PIL import Image
import io
import httpx
from dotenv import load_dotenv

load_dotenv()

IMAGES_DIR = Path("static/images")
IMAGES_DIR.mkdir(parents=True, exist_ok=True)


class ImageService:
    """Handles image generation via Stability AI or Replicate."""

    def __init__(self):
        self.stability_key = os.getenv("STABILITY_API_KEY")
        self.replicate_token = os.getenv("REPLICATE_API_TOKEN")

    async def generate_stability(
        self,
        prompt: str,
        negative_prompt: str = "",
        width: int = 1024,
        height: int = 1024,
        steps: int = 30,
        cfg_scale: float = 7.0,
        seed: int = 0,
    ) -> dict:
        """Generate an image using Stability AI REST API."""
        url = "https://api.stability.ai/v2beta/stable-image/generate/sd3"
        headers = {
            "Authorization": f"Bearer {self.stability_key}",
            "Accept": "application/json",
        }
        payload = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "output_format": "png",
            "width": width,
            "height": height,
            "steps": steps,
            "cfg_scale": cfg_scale,
        }
        if seed:
            payload["seed"] = seed

        async with httpx.AsyncClient(timeout=120) as client:
            response = await client.post(url, headers=headers, json=payload)
            response.raise_for_status()
            data = response.json()

        # Decode base64 image and save
        image_data = base64.b64decode(data["image"])
        return self._save_image(image_data, prompt, "stability", seed)

    async def generate_replicate(
        self,
        prompt: str,
        negative_prompt: str = "",
        width: int = 1024,
        height: int = 1024,
        steps: int = 30,
        guidance_scale: float = 7.0,
        seed: int = 0,
    ) -> dict:
        """Generate an image using Replicate API."""
        import replicate

        model = "stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b"
        input_params = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "width": width,
            "height": height,
            "num_inference_steps": steps,
            "guidance_scale": guidance_scale,
        }
        if seed:
            input_params["seed"] = seed

        output = replicate.run(model, input=input_params)
        # Replicate returns a list of URLs
        image_url = output[0] if isinstance(output, list) else str(output)

        # Download the generated image
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.get(image_url)
            response.raise_for_status()
            image_data = response.content

        return self._save_image(image_data, prompt, "replicate", seed)

    def _save_image(
        self, image_data: bytes, prompt: str, provider: str, seed: int
    ) -> dict:
        """Save image to disk and return metadata."""
        image_id = str(uuid.uuid4())[:12]
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{timestamp}_{image_id}.png"
        filepath = IMAGES_DIR / filename

        # Validate and save with Pillow
        image = Image.open(io.BytesIO(image_data))
        image.save(filepath, "PNG")

        return {
            "id": image_id,
            "filename": filename,
            "url": f"/static/images/{filename}",
            "prompt": prompt,
            "provider": provider,
            "seed": seed,
            "width": image.width,
            "height": image.height,
            "created_at": datetime.now().isoformat(),
        }

The Generation Router

Create routers/generate.py. This handles incoming HTTP requests and delegates to the image service.

# routers/generate.py
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field
from services.image_service import ImageService

router = APIRouter(prefix="/api", tags=["generation"])
image_service = ImageService()

# In-memory store for generated images (replaced with DB later)
generated_images: list[dict] = []


class GenerateRequest(BaseModel):
    prompt: str = Field(..., min_length=1, max_length=2000)
    negative_prompt: str = Field(default="", max_length=1000)
    provider: str = Field(default="stability", pattern="^(stability|replicate)$")
    width: int = Field(default=1024, ge=512, le=2048)
    height: int = Field(default=1024, ge=512, le=2048)
    steps: int = Field(default=30, ge=10, le=50)
    cfg_scale: float = Field(default=7.0, ge=1.0, le=20.0)
    seed: int = Field(default=0, ge=0)


class GenerateResponse(BaseModel):
    id: str
    filename: str
    url: str
    prompt: str
    provider: str
    seed: int
    width: int
    height: int
    created_at: str


@router.post("/generate", response_model=GenerateResponse)
async def generate_image(request: GenerateRequest):
    """Generate an image from a text prompt."""
    try:
        if request.provider == "stability":
            result = await image_service.generate_stability(
                prompt=request.prompt,
                negative_prompt=request.negative_prompt,
                width=request.width,
                height=request.height,
                steps=request.steps,
                cfg_scale=request.cfg_scale,
                seed=request.seed,
            )
        else:
            result = await image_service.generate_replicate(
                prompt=request.prompt,
                negative_prompt=request.negative_prompt,
                width=request.width,
                height=request.height,
                steps=request.steps,
                guidance_scale=request.cfg_scale,
                seed=request.seed,
            )

        generated_images.append(result)
        return result

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/images")
async def list_images():
    """Return all generated images, newest first."""
    return sorted(generated_images, key=lambda x: x["created_at"], reverse=True)

Wire Up the Router

Update main.py to include the generation router:

# main.py (updated)
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.middleware.cors import CORSMiddleware
from dotenv import load_dotenv
from routers import generate
import os

load_dotenv()

app = FastAPI(title="AI Image Generator")

# CORS for frontend development
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

# Register routers
app.include_router(generate.router)


@app.get("/health")
async def health_check():
    return {
        "status": "ok",
        "stability_key": bool(os.getenv("STABILITY_API_KEY")),
        "replicate_key": bool(os.getenv("REPLICATE_API_TOKEN")),
        "openai_key": bool(os.getenv("OPENAI_API_KEY")),
    }

Testing the API

Start the server and test with curl:

# Start the server
uvicorn main:app --reload

# Generate an image with Stability AI
curl -X POST http://localhost:8000/api/generate \
  -H "Content-Type: application/json" \
  -d '{
    "prompt": "A serene mountain landscape at sunset, photorealistic",
    "negative_prompt": "blurry, low quality, distorted",
    "provider": "stability",
    "width": 1024,
    "height": 1024,
    "steps": 30
  }'

You should get a response like this:

{
  "id": "a1b2c3d4e5f6",
  "filename": "20260321_143022_a1b2c3d4e5f6.png",
  "url": "/static/images/20260321_143022_a1b2c3d4e5f6.png",
  "prompt": "A serene mountain landscape at sunset, photorealistic",
  "provider": "stability",
  "seed": 0,
  "width": 1024,
  "height": 1024,
  "created_at": "2026-03-21T14:30:22.456789"
}

Understanding Generation Parameters

Each parameter affects the output in specific ways:

Parameter	Range	Effect
steps	10-50	More steps = higher quality but slower. 25-30 is the sweet spot for most prompts.
cfg_scale	1-20	How closely the image follows the prompt. 7-9 is ideal. Too high creates artifacts.
seed	0+	0 means random. A fixed seed reproduces the exact same image with the same prompt.
width / height	512-2048	Must be multiples of 64. SDXL works best at 1024x1024.
negative_prompt	text	Tells the model what to avoid: "blurry, watermark, low quality" etc.

Error Handling

Add robust error handling for common API failures. Update the generate endpoint:

@router.post("/generate", response_model=GenerateResponse)
async def generate_image(request: GenerateRequest):
    """Generate an image from a text prompt."""
    try:
        if request.provider == "stability":
            if not image_service.stability_key:
                raise HTTPException(
                    status_code=400,
                    detail="Stability AI API key not configured"
                )
            result = await image_service.generate_stability(
                prompt=request.prompt,
                negative_prompt=request.negative_prompt,
                width=request.width,
                height=request.height,
                steps=request.steps,
                cfg_scale=request.cfg_scale,
                seed=request.seed,
            )
        else:
            if not image_service.replicate_token:
                raise HTTPException(
                    status_code=400,
                    detail="Replicate API token not configured"
                )
            result = await image_service.generate_replicate(
                prompt=request.prompt,
                negative_prompt=request.negative_prompt,
                width=request.width,
                height=request.height,
                steps=request.steps,
                guidance_scale=request.cfg_scale,
                seed=request.seed,
            )

        generated_images.append(result)
        return result

    except HTTPException:
        raise
    except httpx.TimeoutException:
        raise HTTPException(status_code=504, detail="Image generation timed out")
    except httpx.HTTPStatusError as e:
        if e.response.status_code == 402:
            raise HTTPException(status_code=402, detail="API credits exhausted")
        raise HTTPException(status_code=502, detail=f"API error: {e.response.text}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")