Update app.py

app.py

@@ -74,9 +74,9 @@ def process_input_image(image_source, download_masks=False):
     
     mask_zip, output_name = create_mask_zip(predicted_image)
     if download_masks:
-        return "Predicted Masked Image", rgb_image, mask_zip.getvalue(), output_name
+        return {"img_output": rgb_image, "masks_zip": mask_zip.getvalue()}
     else:
-        return "Predicted Masked Image", rgb_image, None, None
+        return {"img_output": rgb_image}
 
 
 my_app = gr.Blocks()