jbilcke-hf HF Staff commited on
Commit
bc207eb
·
1 Parent(s): 37e49b8
Files changed (1) hide show
  1. app.py +5 -2
app.py CHANGED
@@ -6,7 +6,10 @@ WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "/data/weights")
6
  # IMPORTANT: Set environment variables BEFORE importing any modules that use them
7
  os.environ["MODEL_BASE"] = os.path.join(WEIGHTS_PATH, "stdmodels")
8
  os.environ["DISABLE_SP"] = "1"
9
- os.environ["CPU_OFFLOAD"] = "1"
 
 
 
10
 
11
  import torch
12
  import gradio as gr
@@ -60,7 +63,7 @@ def create_args():
60
  args.use_linear_quadratic_schedule = False
61
  args.linear_schedule_end = 0.25
62
  args.use_deepcache = False
63
- args.cpu_offload = True
64
  args.use_sage = True
65
  args.save_path = './results/'
66
  args.save_path_suffix = ''
 
6
  # IMPORTANT: Set environment variables BEFORE importing any modules that use them
7
  os.environ["MODEL_BASE"] = os.path.join(WEIGHTS_PATH, "stdmodels")
8
  os.environ["DISABLE_SP"] = "1"
9
+ # Configure CPU_OFFLOAD in system environment variables:
10
+ # Set CPU_OFFLOAD=1 to enable CPU offloading (for low VRAM, but slower)
11
+ # Set CPU_OFFLOAD=0 to disable CPU offloading (requires more VRAM, but faster)
12
+ # os.environ["CPU_OFFLOAD"] = "1"
13
 
14
  import torch
15
  import gradio as gr
 
63
  args.use_linear_quadratic_schedule = False
64
  args.linear_schedule_end = 0.25
65
  args.use_deepcache = False
66
+ args.cpu_offload = os.environ.get("CPU_OFFLOAD", "0") == "1"
67
  args.use_sage = True
68
  args.save_path = './results/'
69
  args.save_path_suffix = ''