Update diffusion_lens.py

diffusion_lens.py

@@ -3,9 +3,15 @@ from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
 pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
 pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
 
+# Check if CUDA is available and set the device accordingly
+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):
     print('inside get images')
-    pipeline_output = pipeline(prompt, skip_layers=skip_layers, num_images_per_prompt=1)
+    pipeline_output = pipeline(prompt, skip_layers=skip_layers, num_images_per_prompt=1, return_tensors=False)
     print('after pipeline')
     images = pipeline_output.images
     print('got images')