Spaces:
Running
on
Zero
Running
on
Zero
Update scene_weaver_core.py
Browse files- 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 |
-
|
|
|
|
|
|
|
|
|
|
| 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()
|