from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from diffusers import Flux2KleinPipeline
import io
import base64
from optimization import optimize_pipeline_ #

app = FastAPI()

# --- Load Model ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "black-forest-labs/FLUX.2-klein-9B"

print("🚀 Loading Flux.2-klein-9B Engine...")
pipe = Flux2KleinPipeline.from_pretrained(model_id, torch_dtype=dtype)
pipe.to(device)

# --- Run Optimization ---
# Melakukan kompilasi model agar generate gambar lebih cepat
print("⚙️ Running AOT Optimization...")
optimize_pipeline_(pipe, prompt="Mamboro AI initialization", num_inference_steps=4)

class ImageRequest(BaseModel):
    prompt: str

@app.post("/generate")
async def generate(request: ImageRequest):
    try:
        generator = torch.Generator(device).manual_seed(0)
        
        # Generate gambar
        output = pipe(
            prompt=request.prompt, 
            num_inference_steps=4, 
            generator=generator
        ).images[0]
        
        # Konversi ke Base64 agar mudah diterima React Native
        buffered = io.BytesIO()
        output.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        return {"image": f"data:image/jpeg;base64,{img_str}"}
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
def health_check():
    return {"status": "Mamboro AI Docker is Running", "model": model_id}