Update device_manager.py

device_manager.py

@@ -1,6 +1,8 @@
 import torch
 import os
 import logging
+import spaces
+from functools import wraps
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -17,40 +19,49 @@ class DeviceManager:
     def __init__(self):
         if self._initialized:
             return
-            
         self._initialized = True
-        self._current_device = None
-        self.initialize_device()
-    
-    def initialize_device(self):
+        self.device = self._initialize_device()
+        
+    def _initialize_device(self):
+        """初始化並確定使用的設備"""
         try:
-            if os.environ.get('SPACE_ID'):
-                # 嘗試初始化 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:
-                raise RuntimeError("Not in Spaces environment")
+            # 檢查是否在 Spaces 環境且有 ZeroGPU
+            if os.environ.get('SPACE_ID') and torch.cuda.is_available():
+                logger.info("ZeroGPU environment detected")
+                return 'cuda'
         except Exception as e:
-            logger.warning(f"Using CPU due to: {e}")
-            self._current_device = torch.device('cpu')
+            logger.warning(f"Unable to initialize ZeroGPU: {e}")
+        
+        logger.info("Using CPU")
+        return 'cpu'
     
-    def get_optimal_device(self):
-        if self._current_device is None:
-            self.initialize_device()
-        return self._current_device
-
-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()
+    def get_device(self):
+        """獲取當前設備"""
+        return self.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')
+    def to_device(self, model_or_tensor):
+        """將模型或張量移到正確的設備上"""
+        try:
+            if hasattr(model_or_tensor, 'to'):
+                return model_or_tensor.to(self.device)
+        except Exception as e:
+            logger.warning(f"Failed to move to {self.device}, using CPU: {e}")
+            self.device = 'cpu'
+            return model_or_tensor.to('cpu')
+        return model_or_tensor
+
+def adaptive_gpu(duration=60):
+    """結合 spaces.GPU 和 CPU 降級的裝飾器"""
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            device_mgr = DeviceManager()
+            if device_mgr.get_device() == 'cuda':
+                # 在 ZeroGPU 環境中使用 spaces.GPU
+                decorated = spaces.GPU(duration=duration)(func)
+                return await decorated(*args, **kwargs)
+            else:
+                # 在 CPU 環境中直接執行
+                return await func(*args, **kwargs)
+        return wrapper
+    return decorator