Spaces:
Sleeping
Sleeping
File size: 2,102 Bytes
c769f36 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | from fastapi import FastAPI
from fastapi.responses import FileResponse
from pydantic import BaseModel
from diffusers import AutoPipelineForText2Image
import torch
import uuid
import os
# --------------------------------------------------
# 1. Create FastAPI app
# --------------------------------------------------
app = FastAPI(title="Text to Image API (SDXL Turbo)")
# --------------------------------------------------
# 2. Device & dtype
# --------------------------------------------------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
MODEL_ID = "stabilityai/sdxl-turbo"
# --------------------------------------------------
# 3. Load model once
# --------------------------------------------------
pipe = AutoPipelineForText2Image.from_pretrained(
MODEL_ID,
torch_dtype=DTYPE,
variant="fp16" if DEVICE == "cuda" else None
)
pipe = pipe.to(DEVICE)
# --------------------------------------------------
# 4. Output directory
# --------------------------------------------------
OUTPUT_DIR = "outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# --------------------------------------------------
# 5. Request schema
# --------------------------------------------------
class PromptRequest(BaseModel):
prompt: str
# --------------------------------------------------
# 6. Health check
# --------------------------------------------------
@app.get("/")
def root():
return {"status": "ok", "message": "SDXL Turbo API is running"}
# --------------------------------------------------
# 7. Image generation endpoint
# --------------------------------------------------
@app.post("/generate")
def generate_image(request: PromptRequest):
result = pipe(
prompt=request.prompt,
num_inference_steps=4,
guidance_scale=0.0
)
image = result.images[0]
filename = f"{uuid.uuid4()}.png"
filepath = os.path.join(OUTPUT_DIR, filename)
image.save(filepath)
return FileResponse(
filepath,
media_type="image/png",
filename="generated.png"
)
|