Update app.py

app.py

@@ -292,10 +292,8 @@ def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_imag
     gc.collect()
     torch.cuda.empty_cache()
 
-    return gr.Gallery.update(
-        value=images,
-        columns=1 if num_images == 1 else 2,  # Show one column if one image, otherwise two
-        )
+    columns = 1 if num_images == 1 else 2
+    return gr.update(value=images, columns=columns)
     # return images
 
 ### Description
@@ -401,7 +399,7 @@ with gr.Blocks(css=css) as demo:
 
         with gr.Column():
             # gallery = gr.Gallery(label="Generated Images")
-            gallery = gr.Gallery(label="Generated Images").style(grid=[2], height="auto")  # Default to 2 columns
+            gallery = gr.Gallery(label="Generated Images", columns=2)  # Default to 2 columns
 
         submit.click(
             fn=randomize_seed_fn,