Spaces:
Running
Running
Update diffusion_lens.py
Browse files- diffusion_lens.py +7 -1
diffusion_lens.py
CHANGED
@@ -3,9 +3,15 @@ from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
|
|
3 |
pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
|
4 |
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
|
5 |
|
|
|
|
|
|
|
|
|
|
|
|
|
6 |
def get_images(prompt, skip_layers=0):
|
7 |
print('inside get images')
|
8 |
-
pipeline_output = pipeline(prompt, skip_layers=skip_layers, num_images_per_prompt=1)
|
9 |
print('after pipeline')
|
10 |
images = pipeline_output.images
|
11 |
print('got images')
|
|
|
3 |
pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
|
4 |
pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
|
5 |
|
6 |
+
# Check if CUDA is available and set the device accordingly
|
7 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
8 |
+
|
9 |
+
# Move the pipeline to the device
|
10 |
+
pipeline.to(device)
|
11 |
+
|
12 |
def get_images(prompt, skip_layers=0):
|
13 |
print('inside get images')
|
14 |
+
pipeline_output = pipeline(prompt, skip_layers=skip_layers, num_images_per_prompt=1, return_tensors=False)
|
15 |
print('after pipeline')
|
16 |
images = pipeline_output.images
|
17 |
print('got images')
|