Update app.py

app.py

@@ -1,19 +1,22 @@
 import gradio as gr
 from hf_diffusion_service import HFDiffusionService
+from PIL import Image
 
+# Initialize your diffusion model service
 service = HFDiffusionService()
 
-def generate_ct(mask_image):
-    # Return the PIL image directly
-    return service.generate_image(mask_image)
+# Function that takes a mask image and returns generated CT image
+def generate_ct(mask_image: Image.Image):
+    result = service.generate_image(mask_image)
+    return result  # Gradio automatically converts PIL to base64
 
-# ✅ Single-interface Gradio app
+# Create a single Gradio interface
 demo = gr.Interface(
     fn=generate_ct,
     inputs=gr.Image(type="pil", label="Segmentation Mask"),
     outputs=gr.Image(type="pil", label="Generated CT Scan"),
     title="Conditional Diffusion Medical Image Generator",
-    description="Upload or draw a mask to generate a synthetic CT scan"
+    description="Upload or draw a mask to generate a synthetic CT scan."
 )
 
 if __name__ == "__main__":