import os
import base64
from io import BytesIO
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
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"

@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,
        ).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)