DawnC commited on
Commit
4cac9e5
·
verified ·
1 Parent(s): 768bce1

Update scene_weaver_core.py

Browse files
Files changed (1) hide show
  1. scene_weaver_core.py +5 -1
scene_weaver_core.py CHANGED
@@ -137,13 +137,17 @@ class SceneWeaverCore:
137
 
138
  def _ultra_memory_cleanup(self):
139
  """Ultra aggressive memory cleanup for Colab stability"""
 
140
  logger.debug("🧹 Ultra memory cleanup...")
141
 
142
  # Multiple rounds of garbage collection
143
  for i in range(5):
144
  gc.collect()
145
 
146
- if torch.cuda.is_available():
 
 
 
147
  # Clear all cached memory
148
  torch.cuda.empty_cache()
149
  torch.cuda.ipc_collect()
 
137
 
138
  def _ultra_memory_cleanup(self):
139
  """Ultra aggressive memory cleanup for Colab stability"""
140
+ import os
141
  logger.debug("🧹 Ultra memory cleanup...")
142
 
143
  # Multiple rounds of garbage collection
144
  for i in range(5):
145
  gc.collect()
146
 
147
+ # On Hugging Face Spaces, skip CUDA operations in main process
148
+ is_spaces = os.getenv('SPACE_ID') is not None
149
+
150
+ if not is_spaces and torch.cuda.is_available():
151
  # Clear all cached memory
152
  torch.cuda.empty_cache()
153
  torch.cuda.ipc_collect()