small qol change - add export to .zip option + make it automatic (#3)

- add export to .zip option (0d7435c6f042cf30aa9205d5dee9be347a743572)
- Update app.py (4255813cd1cd050a7aef1951d8fe17c398635c99)
- make exports automatic (4b81adda4f15186c9b054e9af37e63785997a63f)
- Update app.py (c8e8d7cf878beee394ede97752e47a29c36d33d6)


Co-authored-by: Linoy Tsaban <[email protected]>

app.py

@@ -3,6 +3,7 @@ import uuid
 import numpy as np
 import random
 import tempfile
+import zipfile 
 import spaces
 from PIL import Image
 from diffusers import QwenImageLayeredPipeline
@@ -58,6 +59,19 @@ def export_gallery(images):
     pptx_path = imagelist_to_pptx(images)
     return pptx_path
 
+def export_gallery_zip(images):
+    # images: list of tuples (file_path, caption)
+    images = [e[0] for e in images]
+    
+    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
+        with zipfile.ZipFile(tmp.name, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for i, img_path in enumerate(images):
+                # Get the file extension from original file
+                ext = os.path.splitext(img_path)[1] or '.png'
+                # Add each image to the zip with a numbered filename
+                zipf.write(img_path, f"layer_{i+1}{ext}")
+        return tmp.name
+
 @spaces.GPU(duration=300)
 def infer(input_image,
           seed=777,
@@ -104,9 +118,25 @@ def infer(input_image,
         output_images = output.images[0]
     
     output = []
+    temp_files = []
     for i, image in enumerate(output_images):
         output.append(image)
-    return output
+        # Save to temp file for export
+        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        image.save(tmp.name)
+        temp_files.append(tmp.name)
+    
+    # Generate PPTX
+    pptx_path = imagelist_to_pptx(temp_files)
+    
+    # Generate ZIP
+    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
+        with zipfile.ZipFile(tmp.name, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for i, img_path in enumerate(temp_files):
+                zipf.write(img_path, f"layer_{i+1}.png")
+        zip_path = tmp.name
+    
+    return output, pptx_path, zip_path
 
 ensure_dirname(LOG_DIR)
 examples = [
@@ -187,10 +217,11 @@ with gr.Blocks() as demo:
                 
                 run_button = gr.Button("Decompose!", variant="primary")
 
-            with gr.Column(scale=1):
+            with gr.Column(scale=2):
                 gallery = gr.Gallery(label="Layers", columns=4, rows=1, format="png")
-                export_btn = gr.Button("Export as PPTX")
-                export_file = gr.File(label="Download PPTX")
+                with gr.Row():
+                    export_file = gr.File(label="Download PPTX")
+                    export_zip_file = gr.File(label="Download ZIP")
 
     gr.Examples(examples=examples,
                 inputs=[input_image], 
@@ -201,12 +232,6 @@ with gr.Blocks() as demo:
                 run_on_click=True
     )
 
-    export_btn.click(
-        fn=export_gallery,
-        inputs=gallery,
-        outputs=export_file
-    )
-
     run_button.click(
         fn=infer,
         inputs=[
@@ -221,7 +246,7 @@ with gr.Blocks() as demo:
             cfg_norm,
             use_en_prompt,
         ], 
-        outputs=gallery,
+        outputs=[gallery, export_file, export_zip_file],
     )
 
 if __name__ == "__main__":