device_manager.py

from functools import wraps
import torch
from huggingface_hub import HfApi
import os
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DeviceManager:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(DeviceManager, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
            
        self._initialized = True
        self._current_device = None
        
    def check_zero_gpu_availability(self):
        try:
            # 檢查 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
    
    def get_optimal_device(self):
        if self._current_device is None:
            if self.check_zero_gpu_availability():
                try:
                    # 確保 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("Fallback to CPU due to GPU initialization failure")
            else:
                self._current_device = torch.device('cpu')
                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