Update device_manager.py

device_manager.py

@@ -1,8 +1,6 @@
-from functools import wraps
 import torch
 import os
 import logging
-import spaces
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -22,35 +20,37 @@ class DeviceManager:
             
         self._initialized = True
         self._current_device = None
+        self.initialize_device()
+    
+    def initialize_device(self):
         try:
             if os.environ.get('SPACE_ID'):
-                # 使用 spaces 的 GPU wrapper 進行初始化
-                @spaces.GPU
-                def init_gpu():
-                    return torch.device('cuda')
-                self._current_device = init_gpu()
-                logger.info("ZeroGPU initialized successfully")
+                # 嘗試初始化 CUDA 設備
+                if torch.cuda.is_available():
+                    self._current_device = torch.device('cuda')
+                    # 設置 CUDA 設備為可見
+                    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+                    logger.info("CUDA device initialized successfully")
+                else:
+                    raise RuntimeError("CUDA not available")
             else:
-                self._current_device = torch.device('cpu')
+                raise RuntimeError("Not in Spaces environment")
         except Exception as e:
-            logger.warning(f"Failed to initialize ZeroGPU: {e}")
+            logger.warning(f"Using CPU due to: {e}")
             self._current_device = torch.device('cpu')
     
     def get_optimal_device(self):
+        if self._current_device is None:
+            self.initialize_device()
         return self._current_device
 
-def device_handler(func):
-    """Decorator for handling device placement with ZeroGPU support"""
-    @spaces.GPU
-    @wraps(func)
-    async def wrapper(*args, **kwargs):
-        try:
-            return await func(*args, **kwargs)
-        except RuntimeError as e:
-            if "out of memory" in str(e) or "CUDA" in str(e):
-                logger.warning("ZeroGPU unavailable, falling back to CPU")
-                device_mgr = DeviceManager()
-                device_mgr._current_device = torch.device('cpu')
-                return await func(*args, **kwargs)
-            raise e
-    return wrapper
+def to_device(tensor_or_model, device=None):
+    """Helper function to move tensors or models to the appropriate device"""
+    if device is None:
+        device = DeviceManager().get_optimal_device()
+    
+    try:
+        return tensor_or_model.to(device)
+    except Exception as e:
+        logger.warning(f"Failed to move to {device}, using CPU: {e}")
+        return tensor_or_model.to('cpu')