Posted inAI

Build Your Own Image Generation API with Stable Diffusion and FastAPI: Mastering AI on Linux Servers

Artificial Intelligence tutorial - IT technology blog
Artificial Intelligence tutorial - IT technology blog

Why You Should Build Your Own Image Generation API?

If you’re planning to integrate AI image generation into your app, you might immediately think of DALL-E 3 or Midjourney. While convenient, your wallet will soon be crying. A standard DALL-E 3 request costs about $0.04 – $0.08. With 10,000 requests per month, you could easily spend nearly ten million VND in service fees.

If you have a Linux server equipped with an NVIDIA GPU, building your own system is an extremely economical move. You don’t just save money; you also gain the freedom to customize models. You can freely experiment with specialized LoRA or Checkpoint versions from Civitai to create unique image styles.

The biggest challenge is turning offline Python scripts into a stable Web service. The system must handle queues and avoid overflowing GPU memory (VRAM). I applied this formula to an automated content creation project with a traffic of 2,000 images per day, and the results were very smooth.

Core Concepts: Stable Diffusion and FastAPI

To get started, we need to look at the three key components of this stack:

  • Stable Diffusion (SD): An open-source diffusion model. We will use the diffusers library from Hugging Face for more professional model control.
  • FastAPI: A high-speed Python framework with excellent async support. It comes with built-in Swagger UI, allowing you to test your API in just a few clicks.
  • CUDA/PyTorch: The foundational layer that allows Python to leverage the massive computing power of NVIDIA graphics cards.

Hardware Requirements: Avoid the Bottleneck

Practical experience shows that to render a Stable Diffusion v1.5 image in about 3-5 seconds, you need at least the following configuration:

  • OS: Ubuntu 22.04 LTS (the most stable for AI drivers).
  • GPU: NVIDIA with at least 8GB VRAM. An older RTX 3060 is currently a great budget-friendly choice.
  • RAM: At least 16GB to prevent system lag.
  • Disk: 20GB free (each SD model typically weighs between 2GB and 5GB).

Hands-on: Deployment from Zero to a Complete API

Step 1: Setting up the Linux Environment

Make sure you have the NVIDIA Driver and CUDA Toolkit installed. Then, we will isolate the libraries in a virtual environment to avoid cluttering the system.

# Update the system
sudo apt update && sudo apt upgrade -y

# Install Python and venv
sudo apt install python3-venv python3-pip -y

# Create a virtual environment
python3 -m venv sd_env
source sd_env/bin/activate

# Install core libraries
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install diffusers transformers accelerate fastapi uvicorn python-multipart

Step 2: Writing the Image Processing Script (Worker)

Create a worker.py file to manage the model. The secret here is using fp16 (16-bit floating point). It reduces VRAM consumption by 50% while the image quality remains almost identical to the original.

import torch
from diffusers import StableDiffusionPipeline
import io
from PIL import Image

class ImageGenerator:
    def __init__(self):
        self.model_id = "runwayml/stable-diffusion-v1-5"
        # Load model directly to GPU
        self.pipe = StableDiffusionPipeline.from_pretrained(
            self.model_id, 
            torch_dtype=torch.float16
        )
        self.pipe = self.pipe.to("cuda")
        # Optimize memory for low VRAM cards
        self.pipe.enable_attention_slicing()

    def generate(self, prompt: str):
        with torch.autocast("cuda"):
            image = self.pipe(prompt).images[0]
        
        # Export image as bytes for network transmission
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        return img_byte_arr.getvalue()

Step 3: Building the REST API with FastAPI

Now, let’s wrap the logic above into an Endpoint. The main.py file will act as the gateway for receiving user requests.

from fastapi import FastAPI, Response
from worker import ImageGenerator

app = FastAPI(title="ITFromZero SD API")

# Initialize generator once at startup
gen = ImageGenerator()

@app.post("/generate-image")
async def generate_image(prompt: str):
    if not prompt:
        return {"error": "Prompt cannot be empty"}
    
    image_bytes = gen.generate(prompt)
    return Response(content=image_bytes, media_type="image/png")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

Step 4: Handling GPU Bottlenecks (Concurrency)

A rookie mistake I once made was allowing multiple requests to hit the GPU simultaneously. GPUs are not multi-tasking like CPUs; if 2-3 people generate images at once, the server will immediately hit an Out of Memory (OOM) error.

The solution is to use a Lock. It’s like waiting in line for a restroom: first come, first served. Others must wait their turn to ensure the GPU isn’t overloaded.

import asyncio

# GPU Lock
gpu_lock = asyncio.Lock()

@app.post("/generate-image")
async def generate_image(prompt: str):
    async with gpu_lock:
        # Run heavy tasks in a separate thread to avoid blocking the API
        image_bytes = await asyncio.to_thread(gen.generate, prompt)
    return Response(content=image_bytes, media_type="image/png")

Production Operation on Linux

Don’t just run the script manually and close the terminal. Use Systemd to turn it into a background service that automatically restarts if it crashes.

A simple /etc/systemd/system/sd_api.service configuration looks like this:

[Unit]
Description=Stable Diffusion FastAPI Service
After=network.target

[Service]
User=ubuntu
WorkingDirectory=/home/ubuntu/sd-project
ExecStart=/home/ubuntu/sd-project/sd_env/bin/uvicorn main:app --host 0.0.0.0 --port 8000
Restart=always

[Install]
WantedBy=multi-user.target

Activate the service with just three commands:

sudo systemctl daemon-reload
sudo systemctl enable sd_api
sudo systemctl start sd_api

Hard-earned Lessons for Performance Optimization

After months of operation, I’ve gathered three tips that significantly improve system performance:

  1. Install xFormers: This library speeds up image generation and reduces VRAM by another 15%. It’s highly valuable for 8GB cards.
  2. Use Offline Mode: Download the model to your hard drive beforehand. Don’t wait for it to download from Hugging Face every time the server starts; it’s slow and prone to connection errors.
  3. Caching Mechanism: For popular prompts, use Redis to store the results. The GPU shouldn’t waste power recalculating things it has already done.

Conclusion

Building your own Image Generation API isn’t hard; the difficulty lies in managing GPU resources wisely. Instead of depending on tech giants with expensive costs, you now have full control over the AI image generation process. If your application explodes to millions of users, that’s the time to consider Celery and RabbitMQ to coordinate multiple GPUs simultaneously. For now, this solution is more than enough to get you started in the AI game.

Share:
Tags: