Update app.py

app.py

@@ -58,13 +58,13 @@ model_num_of_layers = {
 }
 
 
-def generate_images(prompt, model, seed):
+def generate_images(prompt, model, seed, skip):
     seed = random.randint(0, MAX_SEED) if seed == -1 else seed
     print('calling diffusion lens with model:', model, 'and seed:', seed)
     gr.Info('Generating images from intermediate layers..')
     all_images = []  # Initialize a list to store all images
     max_num_of_layers = model_num_of_layers[model]
-    for skip_layers in range(max_num_of_layers - 1, -1, -1):
+    for skip_layers in range(max_num_of_layers - 1, -1, -1 * skip):
         # 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_{max_num_of_layers - skip_layers}'))
@@ -106,10 +106,19 @@ with gr.Blocks() as demo:
         label="Seed Value",
     )
 
+    skip = gr.Slider(
+        minimum=0,
+        maximum=6,
+        value=4,
+        step=1,
+        label="Number of layers to skip between each generaion",
+    )
+
     inputs = [
         prompt,
         model,
         seed,
+        skip,
     ]