Spaces:
Running
on
Zero
Running
on
Zero
Delete device_manager.py
Browse files- device_manager.py +0 -67
device_manager.py
DELETED
@@ -1,67 +0,0 @@
|
|
1 |
-
import torch
|
2 |
-
import os
|
3 |
-
import logging
|
4 |
-
import spaces
|
5 |
-
from functools import wraps
|
6 |
-
|
7 |
-
logging.basicConfig(level=logging.INFO)
|
8 |
-
logger = logging.getLogger(__name__)
|
9 |
-
|
10 |
-
class DeviceManager:
|
11 |
-
_instance = None
|
12 |
-
|
13 |
-
def __new__(cls):
|
14 |
-
if cls._instance is None:
|
15 |
-
cls._instance = super(DeviceManager, cls).__new__(cls)
|
16 |
-
cls._instance._initialized = False
|
17 |
-
return cls._instance
|
18 |
-
|
19 |
-
def __init__(self):
|
20 |
-
if self._initialized:
|
21 |
-
return
|
22 |
-
self._initialized = True
|
23 |
-
self.device = self._initialize_device()
|
24 |
-
|
25 |
-
def _initialize_device(self):
|
26 |
-
"""初始化並確定使用的設備"""
|
27 |
-
try:
|
28 |
-
# 檢查是否在 Spaces 環境且有 ZeroGPU
|
29 |
-
if os.environ.get('SPACE_ID') and torch.cuda.is_available():
|
30 |
-
logger.info("ZeroGPU environment detected")
|
31 |
-
return 'cuda'
|
32 |
-
except Exception as e:
|
33 |
-
logger.warning(f"Unable to initialize ZeroGPU: {e}")
|
34 |
-
|
35 |
-
logger.info("Using CPU")
|
36 |
-
return 'cpu'
|
37 |
-
|
38 |
-
def get_device(self):
|
39 |
-
"""獲取當前設備"""
|
40 |
-
return self.device
|
41 |
-
|
42 |
-
def to_device(self, model_or_tensor):
|
43 |
-
"""將模型或張量移到正確的設備上"""
|
44 |
-
try:
|
45 |
-
if hasattr(model_or_tensor, 'to'):
|
46 |
-
return model_or_tensor.to(self.device)
|
47 |
-
except Exception as e:
|
48 |
-
logger.warning(f"Failed to move to {self.device}, using CPU: {e}")
|
49 |
-
self.device = 'cpu'
|
50 |
-
return model_or_tensor.to('cpu')
|
51 |
-
return model_or_tensor
|
52 |
-
|
53 |
-
def adaptive_gpu(duration=60):
|
54 |
-
"""結合 spaces.GPU 和 CPU 降級的裝飾器"""
|
55 |
-
def decorator(func):
|
56 |
-
@wraps(func)
|
57 |
-
async def wrapper(*args, **kwargs):
|
58 |
-
device_mgr = DeviceManager()
|
59 |
-
if device_mgr.get_device() == 'cuda':
|
60 |
-
# 在 ZeroGPU 環境中使用 spaces.GPU
|
61 |
-
decorated = spaces.GPU(duration=duration)(func)
|
62 |
-
return await decorated(*args, **kwargs)
|
63 |
-
else:
|
64 |
-
# 在 CPU 環境中直接執行
|
65 |
-
return await func(*args, **kwargs)
|
66 |
-
return wrapper
|
67 |
-
return decorator
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|