Spaces:

mukaist
/

image

Running on Zero

multimodalart HF staff commited on Dec 5, 2023

Commit

69d5ab5

•

1 Parent(s): 608ae46

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -32,11 +32,11 @@ if torch.cuda.is_available():
         add_watermarker=False,
         variant="fp16"
     )
-    #if ENABLE_CPU_OFFLOAD:
-    #    pipe.enable_model_cpu_offload()
-    #else:
-    #    pipe.to(device)
-    #    print("Loaded on Device!")
     if USE_TORCH_COMPILE:
         pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

         add_watermarker=False,
         variant="fp16"
     )
+    if ENABLE_CPU_OFFLOAD:
+        pipe.enable_model_cpu_offload()
+    else:
+        pipe.to(device)
+        print("Loaded on Device!")
     if USE_TORCH_COMPILE:
         pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)