Update app.py

app.py

@@ -279,28 +279,17 @@ with gr.Blocks(title="DiT Diffusion Model Generator", theme=gr.themes.Soft()) as
             progress_bar = gr.Progress(track_tqdm=True)
             
         with gr.Column(scale=2):
-            with gr.Box(elem_id="gallery_container"):
-                output_gallery = gr.Gallery(
-                    label="Generated Images",
-                    columns=4,
-                    rows=3,
-                    object_fit="contain",
-                    elem_id="scrollable_gallery"
-                )
+            output_gallery = gr.Gallery(
+                label="Generated Images",
+                columns=4,
+                rows=3,
+                height=550,
+                object_fit="contain",
+                allow_preview=True,
+                show_download_button=True
+            )
             error_message = gr.Textbox(label="Error", visible=False, max_lines=3, container=True, elem_id="error_box")
 
-    gr.HTML("""
-    <style>
-        #gallery_container {
-            height: 600px;
-            overflow-y: auto;
-            padding-right: 10px;
-        }
-        #scrollable_gallery {
-            min-height: 100%
-        }
-    </style>
-    """)
     
     dit_size.change(update_sample_slider, inputs=[dit_size],outputs=[num_samples])