Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from PIL import Image
+import io
 
 # Function to process and display the images in grid format
 def create_collage(uploaded_files, grid_layout):
@@ -35,7 +36,11 @@ def create_collage(uploaded_files, grid_layout):
         col = i % cols
         collage.paste(img, (col * img.width, row * img.height))
 
-    return collage
+    # Save the collage as a PNG file in memory
+    output_io = io.BytesIO()
+    collage.save(output_io, format="PNG")
+    output_io.seek(0)  # Reset the pointer to the start of the file
+    return output_io
 
 # Gradio Interface
 with gr.Blocks() as app:
@@ -61,8 +66,8 @@ with gr.Blocks() as app:
         uploaded_files.change(create_collage, inputs=[uploaded_files, grid_layout], outputs=output)
         grid_layout.change(create_collage, inputs=[uploaded_files, grid_layout], outputs=output)
 
-        # Step 6: Download button for the final collage
-        download_button = gr.Button("Download Collage")
+        # Step 6: Download button for the final collage as PNG
+        download_button = gr.Button("Download Collage (PNG)")
         download_button.click(create_collage, inputs=[uploaded_files, grid_layout], outputs=output)
 
 app.launch(share=True)  # Public link to share the app