Posted inAI

Tự build Image Generation API với Stable Diffusion và FastAPI: Làm chủ AI trên Linux Server

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

Tại sao bạn nên tự xây dựng Image Generation API?

Nếu đang định tích hợp tính năng tạo ảnh AI vào ứng dụng, chắc hẳn bạn sẽ nghĩ ngay đến DALL-E 3 hay Midjourney. Tiện thì có tiện, nhưng ví của bạn sẽ “khóc” sớm thôi. Một request DALL-E 3 tiêu chuẩn tốn khoảng $0.04 – $0.08. Với 10.000 request mỗi tháng, bạn dễ dàng bay mất gần chục triệu đồng tiền phí dịch vụ.

Nếu bạn có sẵn server Linux lắp GPU NVIDIA, việc tự dựng hệ thống riêng là bước đi cực kỳ kinh tế. Bạn không chỉ tiết kiệm hầu bao mà còn tự do tùy chỉnh model. Bạn có thể thoải mái vọc các bản LoRA hay Checkpoint chuyên biệt trên Civitai để tạo ra phong cách hình ảnh độc bản.

Thách thức lớn nhất là biến các script Python chạy offline thành một dịch vụ Web ổn định. Hệ thống phải biết xử lý hàng đợi và không làm tràn bộ nhớ GPU (VRAM). Mình đã áp dụng công thức này cho một dự án tạo nội dung tự động với lưu lượng 2.000 ảnh/ngày và kết quả rất mượt mà.

Khái niệm cốt lõi: Stable Diffusion và FastAPI

Để bắt đầu, chúng ta cần điểm qua ba thành phần then chốt trong bộ stack này:

  • Stable Diffusion (SD): Model khuếch tán mã nguồn mở. Chúng ta sẽ dùng thư viện diffusers từ Hugging Face để điều khiển model chuyên nghiệp hơn.
  • FastAPI: Framework Python tốc độ cao, hỗ trợ async cực tốt. Nó cung cấp sẵn Swagger UI, giúp bạn test API chỉ trong vài cú click.
  • CUDA/PyTorch: Lớp nền tảng cho phép Python tận dụng sức mạnh tính toán kinh khủng từ card đồ họa NVIDIA.

Yêu cầu phần cứng: Đừng để thắt nút cổ chai

Kinh nghiệm thực tế cho thấy, để render ảnh Stable Diffusion v1.5 trong khoảng 3-5 giây, bạn cần cấu hình tối thiểu như sau:

  • OS: Ubuntu 22.04 LTS (ổn định nhất cho driver AI).
  • GPU: NVIDIA với ít nhất 8GB VRAM. Một chiếc RTX 3060 cũ là lựa chọn cực ngon bổ rẻ hiện nay.
  • RAM: Từ 16GB để hệ thống không bị giật lag.
  • Disk: 20GB trống (mỗi model SD thường nặng từ 2GB đến 5GB).

Thực hành: Triển khai từ Zero đến API hoàn chỉnh

Bước 1: Thiết lập môi trường Linux

Hãy chắc chắn bạn đã cài Driver NVIDIA và CUDA Toolkit. Sau đó, chúng ta sẽ nhốt các thư viện vào môi trường ảo để tránh làm rác hệ thống.

# Cập nhật hệ thống
sudo apt update && sudo apt upgrade -y

# Cài đặt Python và venv
sudo apt install python3-venv python3-pip -y

# Tạo môi trường ảo
python3 -m venv sd_env
source sd_env/bin/activate

# Cài đặt các thư viện lõi
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install diffusers transformers accelerate fastapi uvicorn python-multipart

Bước 2: Viết script xử lý hình ảnh (Worker)

Tạo file worker.py để quản lý model. Bí kíp ở đây là dùng fp16 (số thực dấu phẩy động 16-bit). Nó giúp giảm 50% lượng VRAM tiêu thụ nhưng chất lượng ảnh tạo ra hầu như không thay đổi so với bản gốc.

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 lên thẳng GPU
        self.pipe = StableDiffusionPipeline.from_pretrained(
            self.model_id, 
            torch_dtype=torch.float16
        )
        self.pipe = self.pipe.to("cuda")
        # Tối ưu hóa bộ nhớ cho card ít VRAM
        self.pipe.enable_attention_slicing()

    def generate(self, prompt: str):
        with torch.autocast("cuda"):
            image = self.pipe(prompt).images[0]
        
        # Xuất ảnh dưới dạng byte để truyền qua mạng
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        return img_byte_arr.getvalue()

Bước 3: Xây dựng REST API với FastAPI

Bây giờ, hãy bọc logic trên vào một Endpoint. File main.py sẽ đóng vai trò là cửa ngõ tiếp nhận yêu cầu từ người dùng.

from fastapi import FastAPI, Response
from worker import ImageGenerator

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

# Khởi tạo generator một lần duy nhất khi khởi động
gen = ImageGenerator()

@app.post("/generate-image")
async def generate_image(prompt: str):
    if not prompt:
        return {"error": "Prompt không được để trống"}
    
    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)

Bước 4: Xử lý nghẽn cổ chai GPU (Concurrency)

Một sai lầm sơ đẳng mình từng gặp là cho phép nhiều request “đâm” vào GPU cùng lúc. GPU không đa nhiệm như CPU; nếu 2-3 người cùng tạo ảnh, server sẽ dính lỗi Out of Memory (OOM) ngay lập tức.

Giải pháp là dùng Lock. Nó giống như việc xếp hàng đi vệ sinh: ai đến trước dùng trước, người sau phải đợi đến lượt để đảm bảo GPU không bị quá tải.

import asyncio

# Khóa chặn GPU
gpu_lock = asyncio.Lock()

@app.post("/generate-image")
async def generate_image(prompt: str):
    async with gpu_lock:
        # Chạy tác vụ nặng trong luồng riêng để không block API
        image_bytes = await asyncio.to_thread(gen.generate, prompt)
    return Response(content=image_bytes, media_type="image/png")

Vận hành thực tế trên Linux

Đừng chạy script bằng tay rồi tắt terminal. Hãy dùng Systemd để biến nó thành một dịch vụ ngầm, tự động hồi sinh nếu chẳng may bị crash.

Cấu hình file /etc/systemd/system/sd_api.service đơn giản như sau:

[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

Kích hoạt dịch vụ chỉ với 3 dòng lệnh:

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

Kinh nghiệm “xương máu” để tối ưu hiệu suất

Sau nhiều tháng vận hành, mình rút ra 3 mẹo giúp hệ thống chạy mượt hơn đáng kể:

  1. Cài đặt xFormers: Thư viện này giúp tăng tốc độ tạo ảnh và giảm thêm khoảng 15% VRAM. Rất đáng giá cho các dòng card 8GB.
  2. Dùng Offline Mode: Hãy tải sẵn model về đĩa cứng. Đừng để mỗi lần khởi động server lại phải chờ download từ Hugging Face, vừa chậm vừa dễ lỗi kết nối.
  3. Cơ chế Caching: Với các prompt phổ biến, hãy dùng Redis để lưu kết quả. GPU không nên lãng phí sức mạnh để tính toán lại những thứ đã làm rồi.

Lời kết

Tự xây Image Generation API không khó, cái khó là biết cách quản lý tài nguyên GPU sao cho khéo. Thay vì phụ thuộc vào các ông lớn với chi phí đắt đỏ, giờ đây bạn đã làm chủ hoàn toàn quy trình tạo ảnh AI. Nếu ứng dụng của bạn bùng nổ lên hàng triệu user, lúc đó hãy nghĩ đến Celery và RabbitMQ để điều phối nhiều GPU cùng lúc. Còn hiện tại, giải pháp này là quá đủ để bạn bắt đầu cuộc chơi AI rồi.

Share:
Tags: