import io, os
from typing import Optional
from fastapi import FastAPI
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from huggingface_hub import InferenceClient

app = FastAPI(title="ImageGen API")

# Choose a model that supports text->image via Inference API.
# Using a model that works well with InferenceClient
MODEL_ID = os.getenv("MODEL", "stabilityai/stable-diffusion-xl-base-1.0")
HF_TOKEN = os.getenv("HF_TOKEN")  # set in Space Secrets

# Initialize client with or without token
if HF_TOKEN:
    client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
else:
    # Use public inference API (may have rate limits)
    client = InferenceClient(model=MODEL_ID)

class GenReq(BaseModel):
    prompt: str
    width: int = 768
    height: int = 768
    seed: Optional[int] = None
    num_inference_steps: Optional[int] = None  # some models ignore this

@app.get("/")
def root():
    return {"ok": True, "message": "API is live", "model": MODEL_ID}

@app.post("/generate")
def generate(req: GenReq):
    try:
        # Different models have different parameter requirements
        if "sd-turbo" in MODEL_ID.lower():
            # SD Turbo works with minimal parameters
            img = client.text_to_image(
                prompt=req.prompt,
                num_inference_steps=1,  # SD Turbo works best with 1 step
                guidance_scale=0.0,  # SD Turbo doesn't use guidance scale
                width=req.width,
                height=req.height
            )
        else:
            # Standard SDXL parameters
            img = client.text_to_image(
                prompt=req.prompt,
                width=req.width,
                height=req.height,
                seed=req.seed,
                num_inference_steps=req.num_inference_steps or 20
            )
        
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        buf.seek(0)
        return StreamingResponse(buf, media_type="image/png")
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)

# Keep your original echo for testing
class Prompt(BaseModel):
    text: str

@app.post("/echo")
def echo(p: Prompt):
    return {"length": len(p.text), "upper": p.text.upper()}