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"
    )