Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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.
|
45 |
-
|
|
|
|
|
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}")
|