Josh Cox commited on
Commit
74b3e81
1 Parent(s): 158667b
Files changed (1) hide show
  1. 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("cuda")
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)