Update diffusion_lens.py

diffusion_lens.py

@@ -10,8 +10,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Move the pipeline to the device
 pipeline.to(device)
 
-def get_images(prompt, skip_layers=0):
+def get_images(prompt, skip_layers):
     print('inside get images')
+    print(f'skipping {skip_layers}')
     pipeline_output = pipeline(prompt, skip_layers=skip_layers, num_images_per_prompt=1, return_tensors=False)
     print('after pipeline')
     images = pipeline_output.images