Changed Gallery Output

app.py

@@ -53,9 +53,11 @@ def infer(prompts, negative_prompts, image, num_inference_steps, seed, num_sampl
         jit=True,
     ).images
 
-    print(output[0].shape)
+    output = output.reshape((num_samples,) + output.shape[-3:])
+
+    print(output.shape)
     
-    final_image = [np.array(x[0]*255, dtype=np.uint8) for x in output]
+    final_image = [np.array(x*255, dtype=np.uint8) for x in output]
 
     del output
     
@@ -72,7 +74,7 @@ with gr.Blocks(css="h1 { text-align: center }") as demo:
         cond_img = gr.Image(label="Input")\
                      .style(height=400)
         output = gr.Gallery(label="Generated images")\
-                   .style(height=400, rows=[2], columns=[2])
+                   .style(height="auto", rows=[2], columns=[1, 2])
         
     # Submit & Clear
     with gr.Row():