add gallery

app.py

@@ -36,7 +36,7 @@ def process_adain_api(files, style_strength=1.0, dataset_size=100):
                 if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                     zipf.write(os.path.join(output_dir, file), file)
         
-        return zip_path, f"Style transfer completed with strength {style_strength}"
+        return [os.path.join(output_dir, basename) for basename in os.listdir(output_dir)], zip_path, f"Style transfer completed with strength {style_strength}"
         
     except Exception as e:
         return None, f"Error: {str(e)}"
@@ -53,13 +53,14 @@ with gr.Blocks() as adain_demo:
             process_btn = gr.Button("Process", variant="primary")
             
         with gr.Column():
+            gallery_output = gr.Gallery(label="Gallery", columns=3, rows=2, object_fit="contain")
             download_output = gr.File(label="Download Results")
             status_output = gr.Textbox(label="Status", interactive=False)
     
     process_btn.click(
         fn=process_adain_api,
         inputs=[files_input, strength_input, size_input],
-        outputs=[download_output, status_output],
+        outputs=[gallery_output, download_output, status_output],
         api_name="adain_process"  # This creates the API endpoint
     )