Update app.py (#7)

- Update app.py (c99afd26793bc3db68a8a487e67f11c11999cb92)


Co-authored-by: Shengming Yin <[email protected]>

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import random
 import tempfile
 import spaces
+import zipfile 
 from PIL import Image
 from diffusers import QwenImageLayeredPipeline
 import torch
@@ -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 = [
@@ -136,14 +166,14 @@ with gr.Blocks() as demo:
             with gr.Column(scale=1):
                 input_image = gr.Image(label="Input Image", image_mode="RGBA")
                 
-                prompt = gr.Textbox(
-                    label="Prompt (Optional)",
-                    placeholder="Please enter the prompt to descibe the image. （Optional）",
-                    value="",
-                    lines=2,
-                )
                 
                 with gr.Accordion("Advanced Settings", open=False):
+                    prompt = gr.Textbox(
+                        label="Prompt (Optional)",
+                        placeholder="Please enter the prompt to descibe the image. （Optional）",
+                        value="",
+                        lines=2,
+                    )
                     neg_prompt = gr.Textbox(
                         label="Negative Prompt (Optional)",
                         placeholder="Please enter the negative prompt",
@@ -189,26 +219,21 @@ 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], 
-                outputs=[gallery], 
+                outputs=[gallery, export_file, export_zip_file],
                 fn=infer, 
                 examples_per_page=14,
                 cache_examples=False,
                 run_on_click=True
     )
 
-    export_btn.click(
-        fn=export_gallery,
-        inputs=gallery,
-        outputs=export_file
-    )
-
     run_button.click(
         fn=infer,
         inputs=[
@@ -223,7 +248,7 @@ with gr.Blocks() as demo:
             cfg_norm,
             use_en_prompt,
         ], 
-        outputs=gallery,
+        outputs=[gallery, export_file, export_zip_file],
     )
 
 if __name__ == "__main__":