from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import torch
from diffusers import FluxPipeline
import base64
from io import BytesIO

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load FLUX.1-schnell (Best for 16GB VRAM + Fast Generation)
print("Initializing Flux-Schnell on RTX 5060 Ti (sm_120)...")
try:
    pipe = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-schnell", 
        torch_dtype=torch.bfloat16
    )
    # Optimized for 16GB: Moves parts of the model to CPU when not in use
    # pipe.enable_model_cpu_offload()
    pipe.enable_sequential_cpu_offload() 
    pipe.vae.enable_tiling()
except Exception as e:
    print(f"Error loading model: {e}")

class ImageRequest(BaseModel):
    prompt: str
    model: str = "flux-schnell" # Matches the error fix from earlier

@app.post("/v1/images/generations")
async def generate(request: ImageRequest):
    try:
        # Flux-Schnell generates high quality in just 4 steps
        result = pipe(
            request.prompt,
            guidance_scale=0.0,
            num_inference_steps=4,
            max_sequence_length=256,
        ).images[0]
        
        # Convert to Base64 for the API response
        buffered = BytesIO()
        result.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        return {"data": [{"b64_json": img_str}]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5005)