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

app = FastAPI()

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

print("Initializing Flux-Schnell on RTX 5060 Ti (sm_120)...")
hf_token = os.environ.get("HF_TOKEN")

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
    token=hf_token
)

# Offload execution stages sequentially to maintain 16GB headroom
pipe.enable_sequential_cpu_offload()
pipe.vae.enable_tiling()

class ImageRequest(BaseModel):
    prompt: str
    model: str = "flux-schnell"
    width: int = 1024
    height: int = 1024

    @model_validator(mode="before")
    @classmethod
    def validate_dimensions(cls, data: Any) -> Any:
        if isinstance(data, dict):
            width = data.get("width", 1024)
            height = data.get("height", 1024)
            if width % 32 != 0 or height % 32 != 0:
                raise ValueError("Width and height must be multiples of 32")
            if width > 2048 or height > 2048:
                raise ValueError("Maximum width and height is 2048")
        return data

@app.post("/v1/images/generations")
async def generate(request: ImageRequest):
    try:
        result = pipe(
            request.prompt,
            guidance_scale=0.0,
            num_inference_steps=4,
            max_sequence_length=256,
            width=request.width,
            height=request.height,
        ).images[0]

        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)