Update app.py

app.py

@@ -74,14 +74,13 @@ def process_input_image(image_source, download_masks=False):
     
     mask_zip, output_name = create_mask_zip(predicted_image)
     if download_masks:
-        masks_zip.set_attributes(label="Download Masks", filename="masks.zip", hidden=False)
         return "Predicted Masked Image", rgb_image, mask_zip, output_name
     else:
-        masks_zip.set_attributes(hidden=True)
         return "Predicted Masked Image", rgb_image, None, None
 
 
 
+
 my_app = gr.Blocks()
 
 
@@ -103,6 +102,33 @@ def create_mask_zip(mask, output_name='masks.zip'):
     buffer.seek(0)
     return buffer, output_name
 
+class MaskZipOutput(gr.outputs.OutputComponent):
+    def __init__(self):
+        self.label = "Download Masks"
+        self.filename = "masks.zip"
+        self.mime_type = "application/zip"
+
+    def get_template_context(self):
+        return {
+            "label": self.label,
+            "filename": self.filename,
+            "mime_type": self.mime_type
+        }
+
+    def component_type(self):
+        return "output.file"
+
+    def to_response(self, value):
+        if value is None:
+            return {}
+        return {
+            "file": {
+                "name": self.filename,
+                "content": value.read(),
+                "mime_type": self.mime_type
+            }
+        }
+
 
 # Define the custom legend HTML
 legend_html = '''