jingyangcarl commited on
Commit
a7740b0
·
1 Parent(s): a2072c0

debug zerogpu

Browse files
Files changed (1) hide show
  1. model.py +8 -1
model.py CHANGED
@@ -64,7 +64,14 @@ class Model:
64
  pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
65
  pipe.to(self.device)
66
  if self.device.type == "cuda":
67
- pipe.enable_xformers_memory_efficient_attention()
 
 
 
 
 
 
 
68
  torch.cuda.empty_cache()
69
  gc.collect()
70
  self.base_model_id = base_model_id
 
64
  pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
65
  pipe.to(self.device)
66
  if self.device.type == "cuda":
67
+ import os
68
+ if os.environ.get("SPACES_ZERO_GPU", "0") == "1"
69
+ # when running on ZeroGPU, enable CPU offload
70
+ # pipe.enable_xformers_memory_efficient_attention() doens't work
71
+ # pipe.enable_model_cpu_offload()
72
+ pass
73
+ else:
74
+ pipe.enable_xformers_memory_efficient_attention()
75
  torch.cuda.empty_cache()
76
  gc.collect()
77
  self.base_model_id = base_model_id