Update app.py

app.py

@@ -37,6 +37,11 @@ model_num_of_layers = {
     'Stable Diffusion 2.1': 22,
 }
 
+num_rows_dict = {
+    'Stable Diffusion 1.4': 2,
+    'Stable Diffusion 2.1': 4,
+}
+
 def generate_images(prompt, model, seed):
     seed = random.randint(0, MAX_SEED) if seed == -1 else seed
     print('calling diffusion lens with model:', model, 'and seed:', seed)
@@ -46,7 +51,7 @@ def generate_images(prompt, model, seed):
     for skip_layers in range(max_num_of_layers - 1, -1, -1):
         # Pass the model and seed to the get_images function
         images = get_images(prompt, skip_layers=skip_layers, model=model, seed=seed)
-        all_images.append((images[0], f'layer_{12 - skip_layers}'))
+        all_images.append((images[0], f'layer_{max_num_of_layers - skip_layers}'))
         yield all_images
 
 with gr.Blocks() as demo:
@@ -93,9 +98,9 @@ with gr.Blocks() as demo:
 
 
     generate_button = gr.Button("Generate Image")
-
+    num_rows = num_rows_dict[model]
     with gr.Column():
-        gallery = gr.Gallery(label="Generated Images", columns=4, rows=3, object_fit="contain", height="auto")
+        gallery = gr.Gallery(label="Generated Images", columns=6, rows=num_rows, object_fit="contain", height="auto")
 
     outputs = [gallery]