add export to .zip option

app.py

@@ -3,6 +3,7 @@ import uuid
 import numpy as np
 import random
 import tempfile
+import zipfile  # Added import
 import spaces
 from PIL import Image
 from diffusers import QwenImageLayeredPipeline
@@ -15,7 +16,7 @@ LOG_DIR = "/tmp/local"
 MAX_SEED = np.iinfo(np.int32).max
 
 from huggingface_hub import login
-login(token=os.environ.get('hf'))
+login(token=os.environ.get('HF_TOKEN'))
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -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,
@@ -189,8 +203,12 @@ with gr.Blocks() as demo:
 
             with gr.Column(scale=1):
                 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_btn = gr.Button("Export as PPTX")
+                    export_zip_btn = gr.Button("Export as ZIP")
+                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], 
@@ -207,6 +225,12 @@ with gr.Blocks() as demo:
         outputs=export_file
     )
 
+    export_zip_btn.click(
+        fn=export_gallery_zip,
+        inputs=gallery,
+        outputs=export_zip_file
+    )
+
     run_button.click(
         fn=infer,
         inputs=[