Update app.py

app.py

@@ -1,45 +1,49 @@
 import gradio as gr
 import numpy as np
-import tempfile
-import shutil
 import os
+from PIL import Image
+import io
 from zipfile import ZipFile
 
-def sepia(input_img, num_copies):
+def sepia_and_zip(input_img, num_copies):
+    # Apply the sepia filter to the image
     sepia_filter = np.array([
-        [0.393, 0.769, 0.189], 
-        [0.349, 0.686, 0.168], 
+        [0.393, 0.769, 0.189],
+        [0.349, 0.686, 0.168],
         [0.272, 0.534, 0.131]
     ])
     input_img = input_img / 255.0
     sepia_img = np.dot(input_img[..., :3], sepia_filter.T)
     sepia_img = np.clip(sepia_img, 0, 1) * 255
     sepia_imgs = [(sepia_img).astype(np.uint8) for _ in range(num_copies)]
-    return sepia_imgs
 
-def zip_sepia_images(sepia_imgs):
-    temp_dir = tempfile.mkdtemp()
-    zip_path = os.path.join(temp_dir, "sepia_images.zip")
-    with ZipFile(zip_path, 'w') as zipf:
-        for i, img in enumerate(sepia_imgs):
-            img_path = os.path.join(temp_dir, f"sepia_image_{i}.png")
-            gr.utils.save_image(img_path, img)
-            zipf.write(img_path, os.path.basename(img_path))
-    return zip_path
+    # Convert numpy images to PIL images for gallery
+    pil_images = [Image.fromarray(img) for img in sepia_imgs]
+
+    # Prepare ZIP file
+    zip_bytes_io = io.BytesIO()
+    with ZipFile(zip_bytes_io, 'w') as zip_file:
+        for i, img in enumerate(pil_images):
+            img_byte_arr = io.BytesIO()
+            img.save(img_byte_arr, format='PNG')
+            img_byte_arr = img_byte_arr.getvalue()
+            zip_file.writestr(f"sepia_image_{i}.png", img_byte_arr)
+    zip_bytes_io.seek(0)
+
+    return pil_images, zip_bytes_io
 
 with gr.Blocks() as demo:
     with gr.Row():
         input_img = gr.Image()
         num_copies = gr.Number(label="Number of Copies", value=1)
+    with gr.Row():
         gallery = gr.Gallery(label="Sepia Images")
-        download_btn = gr.File(label="Download ZIP", visible=False)
+        download_btn = gr.File(label="Download ZIP")
         
-    input_img.change(fn=lambda x: x, inputs=input_img, outputs=gallery)
-    num_copies.change(fn=lambda x: x, inputs=num_copies, outputs=gallery)
-    
-    generate_btn = gr.Button("Generate Sepia Images")
-    generate_btn.click(fn=sepia, inputs=[input_img, num_copies], outputs=gallery)
-    generate_btn.click(fn=zip_sepia_images, inputs=gallery, outputs=download_btn)
+    btn = gr.Button("Generate Sepia Images")
+    btn.click(fn=sepia_and_zip, 
+              inputs=[input_img, num_copies], 
+              outputs=[gallery, download_btn])
 
 if __name__ == "__main__":
     demo.launch()