Update app.py

app.py

@@ -8,6 +8,10 @@ from model import (
 from huggingface_hub import hf_hub_download
 import json
 
+
+print("Gradio version:", gr.__version__)
+
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Download config and checkpoint files from HF Hub
@@ -50,7 +54,6 @@ def generate_image(text_prompt, mask_upload):
     This function returns a generator that yields an intermediate
     decoded image at every timestep from the diffusion process.
     """
-    # sample_ddpm_inference is assumed to be a generator function (using yield)
     return sample_ddpm_inference(unet, vae, text_prompt, mask_upload, device)
 
 
@@ -80,14 +83,12 @@ with gr.Blocks(css=css_str) as demo:
         )
         mask_input = gr.Image(type="pil", label="Optional Mask for Conditioning")
     generate_button = gr.Button("Generate Image")
-    output_image = gr.Image(label="Generated Image", type="pil")
+    # output_image = gr.Image(label="Generated Image", type="pil")
 
-    # Adding stream=True allows Gradio to process the generator output
     generate_button.click(
         fn=generate_image,
         inputs=[text_input, mask_input],
-        outputs=output_image,
-        stream=True,
+        outputs="image",
     )
 
 if __name__ == "__main__":