Update inpainting_module.py

inpainting_module.py

@@ -186,13 +186,17 @@ class InpaintingModule:
         for _ in range(rounds):
             gc.collect()
 
-        if torch.cuda.is_available():
+        # On Hugging Face Spaces, avoid CUDA operations in main process
+        # CUDA operations must only happen within @spaces.GPU decorated functions
+        is_spaces = os.getenv('SPACE_ID') is not None
+
+        if not is_spaces and torch.cuda.is_available():
             torch.cuda.empty_cache()
             if aggressive:
                 torch.cuda.ipc_collect()
                 torch.cuda.synchronize()
 
-        logger.debug(f"Memory cleanup completed (aggressive={aggressive})")
+        logger.debug(f"Memory cleanup completed (aggressive={aggressive}, spaces={is_spaces})")
 
     def _check_memory_status(self) -> Dict[str, float]:
         """
@@ -203,7 +207,10 @@ class InpaintingModule:
         dict
             Memory statistics including allocated, total, and usage ratio
         """
-        if not torch.cuda.is_available():
+        # On Spaces, skip CUDA checks in main process
+        is_spaces = os.getenv('SPACE_ID') is not None
+
+        if is_spaces or not torch.cuda.is_available():
             return {"available": True, "usage_ratio": 0.0}
 
         allocated = torch.cuda.memory_allocated() / 1024**3