fastapi.py

!pip install fastapi uvicorn diffusers transformers torch accelerate safetensors --quiet
!pip install nest_asyncio pyngrok --quiet

from pyngrok import ngrok
import base64
from io import BytesIO
import torch
import uvicorn
import nest_asyncio
from fastapi import FastAPI
from diffusers import DiffusionPipeline

# Đặt authtoken của bạn vào đây
authtoken = "2ufMyhX6auLeWfr87bkTdZHjTvj_5VijMy2cXE411GbRnJ13n"
ngrok.set_auth_token(authtoken)

# Khởi tạo FastAPI
app = FastAPI()

# Load mô hình Animagine-XL-3.1 trên GPU
model_id = "cagliostrolab/animagine-xl-3.1"
pipe = DiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16  # Sử dụng float16 cho GPU để tăng tốc và tiết kiệm bộ nhớ
)
# Đảm bảo đang chạy trên GPU
pipe = pipe.to("cuda")

@app.get("/")
def home():
    return {"message": "API Animagine-XL-3.1 is running on GPU"}

@app.get("/generate")
def generate_image(prompt: str = "1girl, anime style, masterpiece, best quality"):
    # Tạo ảnh từ prompt trên GPU
    image = pipe(
        prompt,
        num_inference_steps=30,  # Số bước mặc định có thể tăng trên GPU
        guidance_scale=7.0,      # Giá trị guidance chuẩn
    ).images[0]
    
    # Chuyển đổi ảnh thành base64
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    
    # Trả về chuỗi base64 trong response
    return {
        "status": "success",
        "base64_image": img_str
    }

# Chạy Uvicorn với Nest Asyncio
if __name__ == "__main__":
    nest_asyncio.apply()
    ngrok_tunnel = ngrok.connect(8000)
    print(f"API đang chạy tại: {ngrok_tunnel.public_url}")
    
    uvicorn.run(app, host="0.0.0.0", port=8000)