Update device_manager.py

device_manager.py

@@ -25,25 +25,46 @@ class DeviceManager:
         
     def check_zero_gpu_availability(self):
         try:
-            if 'SPACE_ID' in os.environ:
-                api = HfApi()
-                space_info = api.get_space_runtime(os.environ['SPACE_ID'])
-                if hasattr(space_info, 'hardware') and space_info.hardware.get('zerogpu', False):
-                    return True
+            # 檢查 Hugging Face Space 環境變數
+            if not os.environ.get('SPACE_ID'):
+                return False
+                
+            # 檢查是否在 Spaces 環境中並且啟用了 ZeroGPU
+            if os.environ.get('ZERO_GPU_AVAILABLE') == '1':
+                return True
+                
+            return False
+            
         except Exception as e:
             logger.warning(f"Error checking ZeroGPU availability: {e}")
-        return False
+            return False
     
     def get_optimal_device(self):
         if self._current_device is None:
             if self.check_zero_gpu_availability():
                 try:
-                    self._current_device = torch.device('cuda')
-                    logger.info("Using ZeroGPU")
-                except Exception:
+                    # 確保 CUDA 可用
+                    if torch.cuda.is_available():
+                        self._current_device = torch.device('cuda')
+                        logger.info("Using ZeroGPU")
+                    else:
+                        raise RuntimeError("CUDA not available")
+                except Exception as e:
+                    logger.warning(f"Failed to initialize ZeroGPU: {e}")
                     self._current_device = torch.device('cpu')
-                    logger.info("Failed to use ZeroGPU, falling back to CPU")
+                    logger.info("Fallback to CPU due to GPU initialization failure")
             else:
                 self._current_device = torch.device('cpu')
-                logger.info("Using CPU")
-        return self._current_device
+                logger.info("Using CPU (ZeroGPU not available)")
+        return self._current_device
+
+    def move_to_device(self, tensor_or_model):
+        device = self.get_optimal_device()
+        try:
+            if hasattr(tensor_or_model, 'to'):
+                return tensor_or_model.to(device)
+        except Exception:
+            self._current_device = torch.device('cpu')
+            if hasattr(tensor_or_model, 'to'):
+                return tensor_or_model.to('cpu')
+        return tensor_or_model