import io
import os
import base64
import asyncio
import random
from concurrent.futures import ThreadPoolExecutor
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse
from PIL import Image
import torch
from diffusers import DiffusionPipeline
# -------------------------------------------------------------
# HuggingFace Token (auto)
# -------------------------------------------------------------
HF_TOKEN = os.getenv("HF_TOKEN")
# -------------------------------------------------------------
# Model Settings
# -------------------------------------------------------------
MODEL_REPO = "stabilityai/sdxl-turbo"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
print(f"Loading {MODEL_REPO} on {device}...")
pipe = DiffusionPipeline.from_pretrained(
MODEL_REPO,
torch_dtype=dtype,
use_safetensors=True,
token=HF_TOKEN if HF_TOKEN else None, # auto use token
)
pipe.to(device)
if device == "cpu":
try:
pipe.enable_model_cpu_offload()
except:
pass
print("Model ready.")
# -------------------------------------------------------------
# Core Generation Function
# -------------------------------------------------------------
def generate_image(prompt, negative_prompt, seed, width, height, steps, guidance):
generator = torch.Generator(device=device).manual_seed(seed)
result = pipe(
prompt=prompt,
negative_prompt=negative_prompt if negative_prompt else None,
guidance_scale=guidance,
num_inference_steps=steps,
width=width,
height=height,
generator=generator,
)
return result.images[0]
# -------------------------------------------------------------
# Async Queue
# -------------------------------------------------------------
executor = ThreadPoolExecutor(max_workers=2)
semaphore = asyncio.Semaphore(2)
async def run_generate(prompt, negative_prompt, seed, width, height, steps, guidance):
async with semaphore:
loop = asyncio.get_running_loop()
return await loop.run_in_executor(
executor,
generate_image,
prompt,
negative_prompt,
seed,
width,
height,
steps,
guidance,
)
# -------------------------------------------------------------
# FastAPI App
# -------------------------------------------------------------
app = FastAPI(title="SDXL Turbo Generator", version="2.0")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# -------------------------------------------------------------
# POST-only HTML UI
# -------------------------------------------------------------
@app.get("/", response_class=HTMLResponse)
def home():
return """
SDXL Turbo CPU Generator
SDXL Turbo (POST-only UI)
![]()
"""
# -------------------------------------------------------------
# API Endpoint (POST only)
# -------------------------------------------------------------
@app.post("/api/generate")
async def api_generate(request: Request):
try:
body = await request.json()
prompt = body.get("prompt", "").strip()
negative_prompt = body.get("negative_prompt", "").strip()
except:
return JSONResponse({"status": "error", "message": "Invalid JSON"}, 400)
if not prompt:
return JSONResponse({"status": "error", "message": "Prompt required"}, 400)
width = 768
height = 432
steps = 2
guidance = 0.0
seed = random.randint(0, 2**31 - 1)
try:
img = await run_generate(prompt, negative_prompt, seed, width, height, steps, guidance)
buf = io.BytesIO()
img.save(buf, format="PNG")
b64 = base64.b64encode(buf.getvalue()).decode()
return JSONResponse({
"status": "success",
"image_base64": b64,
"seed": seed,
"width": width,
"height": height
})
except Exception as e:
return JSONResponse({"status": "error", "message": str(e)}, 500)
# -------------------------------------------------------------
# Local run
# -------------------------------------------------------------
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)