Update app.py

app.py

@@ -1,20 +1,24 @@
-import gradio as gr
+from fastapi import FastAPI, UploadFile
+from fastapi.responses import JSONResponse
 from hf_diffusion_service import HFDiffusionService
 from PIL import Image
+import io, base64
 
+app = FastAPI()
 service = HFDiffusionService()
 
-def generate_ct(mask_image: Image.Image):
-    return service.generate_image_base64(mask_image)
-
-demo = gr.Interface(
-    fn=generate_ct,
-    inputs=gr.Image(type="pil", label="Segmentation Mask"),
-    outputs=gr.Image(type="pil", label="Generated CT Scan"),
-    title="Synthetic CT Generator",
-    description="Upload or draw a mask to generate a CT scan.",
-    flagging_mode="never"
-)
-
-if __name__ == "__main__":
-    demo.launch()
+@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}"})