Update device_manager.py

device_manager.py

@@ -3,7 +3,6 @@ import torch
 from huggingface_hub import HfApi
 import os
 import logging
-import asyncio
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -23,7 +22,6 @@ class DeviceManager:
             
         self._initialized = True
         self._current_device = None
-        self._zero_gpu_available = None
         
     def check_zero_gpu_availability(self):
         try:
@@ -31,12 +29,9 @@ class DeviceManager:
                 api = HfApi()
                 space_info = api.get_space_runtime(os.environ['SPACE_ID'])
                 if hasattr(space_info, 'hardware') and space_info.hardware.get('zerogpu', False):
-                    self._zero_gpu_available = True
                     return True
         except Exception as e:
             logger.warning(f"Error checking ZeroGPU availability: {e}")
-        
-        self._zero_gpu_available = False
         return False
     
     def get_optimal_device(self):
@@ -45,26 +40,10 @@ class DeviceManager:
                 try:
                     self._current_device = torch.device('cuda')
                     logger.info("Using ZeroGPU")
-                except Exception as e:
-                    logger.warning(f"Failed to initialize ZeroGPU: {e}")
+                except Exception:
                     self._current_device = torch.device('cpu')
+                    logger.info("Failed to use ZeroGPU, falling back to CPU")
             else:
                 self._current_device = torch.device('cpu')
                 logger.info("Using CPU")
-        return self._current_device
-
-def device_handler(func):
-    """簡化版的 device handler"""
-    @wraps(func)
-    async def wrapper(*args, **kwargs):
-        device_mgr = DeviceManager()
-        try:
-            result = await func(*args, **kwargs)
-            return result
-        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._current_device = torch.device('cpu')
-                return await func(*args, **kwargs)
-            raise e
-    return wrapper
+        return self._current_device