PeterL1n commited on
Commit
83e9007
1 Parent(s): cbd20c7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -3
app.py CHANGED
@@ -40,9 +40,11 @@ def generate(prompt, option, progress=gr.Progress()):
40
  print(prompt, option)
41
  ckpt, step = opts[option]
42
  progress((0, step))
43
-
44
- pipe.to(device, dtype)
45
- safety_checker.to(device, dtype)
 
 
46
 
47
  if step != step_loaded:
48
  print(f"Switching checkpoint from {step_loaded} to {step}")
 
40
  print(prompt, option)
41
  ckpt, step = opts[option]
42
  progress((0, step))
43
+
44
+ if pipe.device != device:
45
+ pipe.to(device, dtype)
46
+ if safety_checker.device != device:
47
+ safety_checker.to(device, dtype)
48
 
49
  if step != step_loaded:
50
  print(f"Switching checkpoint from {step_loaded} to {step}")