DawnC commited on
Commit
255703e
1 Parent(s): f90c43e

Delete device_manager.py

Browse files
Files changed (1) hide show
  1. 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