Commit
•
69d5ab5
1
Parent(s):
608ae46
Update app.py
Browse files
app.py
CHANGED
@@ -32,11 +32,11 @@ if torch.cuda.is_available():
|
|
32 |
add_watermarker=False,
|
33 |
variant="fp16"
|
34 |
)
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
|
41 |
if USE_TORCH_COMPILE:
|
42 |
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
|
|
|
32 |
add_watermarker=False,
|
33 |
variant="fp16"
|
34 |
)
|
35 |
+
if ENABLE_CPU_OFFLOAD:
|
36 |
+
pipe.enable_model_cpu_offload()
|
37 |
+
else:
|
38 |
+
pipe.to(device)
|
39 |
+
print("Loaded on Device!")
|
40 |
|
41 |
if USE_TORCH_COMPILE:
|
42 |
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
|