from fastapi import FastAPI, UploadFile
from fastapi.responses import JSONResponse
from hf_diffusion_service import HFDiffusionService
from PIL import Image
import io, base64
import uvicorn

app = FastAPI()
service = HFDiffusionService()

@app.post("/generate")
async def generate(file: UploadFile):
    # Read file bytes
    image_bytes = await file.read()
    mask = Image.open(io.BytesIO(image_bytes)).convert("L")
    
    # Generate CT
    result = service.generate_image(mask)
    
    # Convert to base64
    buffer = io.BytesIO()
    result.save(buffer, format="PNG")
    base64_img = base64.b64encode(buffer.getvalue()).decode("utf-8")
    
    return JSONResponse(content={"image": f"data:image/png;base64,{base64_img}"})


# ✅ Add this to start FastAPI server for Spaces
if __name__ == "__main__":
    # Run uvicorn explicitly so Hugging Face sees a running server
    uvicorn.run(app, host="0.0.0.0", port=7860)