Josh Cox
commited on
Commit
·
74b3e81
1
Parent(s):
158667b
typo
Browse files- artist_lib.py +2 -1
artist_lib.py
CHANGED
@@ -87,6 +87,7 @@ def generate_tone(note, octave, duration):
|
|
87 |
return sr, audio
|
88 |
|
89 |
def draw(inp, this_model, force_new):
|
|
|
90 |
drawing = inp
|
91 |
if this_model == "stable-diffusion-2":
|
92 |
this_model_addr = "stabilityai/stable-diffusion-2"
|
@@ -111,7 +112,7 @@ def draw(inp, this_model, force_new):
|
|
111 |
print("generating drawing '", drawing, "'", drawing_filename)
|
112 |
pipe = StableDiffusionPipeline.from_pretrained(this_model_addr, torch_dtype=torch.float16)
|
113 |
pipe.enable_attention_slicing()
|
114 |
-
pipe = pipe.to(
|
115 |
image = pipe(drawing).images[0]
|
116 |
image.seek(0)
|
117 |
image.save(drawing_filename)
|
|
|
87 |
return sr, audio
|
88 |
|
89 |
def draw(inp, this_model, force_new):
|
90 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
91 |
drawing = inp
|
92 |
if this_model == "stable-diffusion-2":
|
93 |
this_model_addr = "stabilityai/stable-diffusion-2"
|
|
|
112 |
print("generating drawing '", drawing, "'", drawing_filename)
|
113 |
pipe = StableDiffusionPipeline.from_pretrained(this_model_addr, torch_dtype=torch.float16)
|
114 |
pipe.enable_attention_slicing()
|
115 |
+
pipe = pipe.to(device)
|
116 |
image = pipe(drawing).images[0]
|
117 |
image.seek(0)
|
118 |
image.save(drawing_filename)
|