VincentGOURBIN/MeetingNotes-Voxtral-Analysis

1"""2Zero GPU management for Hugging Face Spaces.3 4This module provides decorators and utilities for efficient GPU usage5in HF Spaces environment with automatic resource management.6"""7 8import functools9import gc10import os11import torch12from typing import Callable, Any13 14# Import spaces if available (HF Spaces environment)15try:16    import spaces17except ImportError:18    spaces = None19 20 21class ZeroGPUManager:22    """Manager for Zero GPU operations in HF Spaces."""23    24    def __init__(self):25        # Device selection with MPS support for local Mac testing26        if torch.backends.mps.is_available():27            self.device = "mps"28            self.dtype = torch.float16  # MPS works better with float1629            print("🚀 Using MPS (Apple Silicon) for local testing")30        elif torch.cuda.is_available():31            self.device = "cuda"32            self.dtype = torch.bfloat16  # CUDA supports bfloat1633            print("🚀 Using CUDA GPU")34        else:35            self.device = "cpu"36            self.dtype = torch.float16  # CPU with float16 to save memory37            print("⚠️ Using CPU")38        39        self.is_spaces = os.getenv("SPACE_ID") is not None40    41    @staticmethod42    def gpu_task(duration: int = 60):43        """44        Decorator for GPU-intensive tasks.45        46        Args:47            duration: Expected duration in seconds for GPU allocation48        """49        def decorator(func: Callable) -> Callable:50            if spaces is not None and hasattr(spaces, 'GPU'):51                # Use HF Spaces GPU decorator52                return spaces.GPU(duration=duration)(func)53            else:54                # Fallback for local development55                return func56        return decorator57    58    @staticmethod59    def cleanup_gpu():60        """Clean up GPU memory after processing (CUDA/MPS/CPU)."""61        if torch.backends.mps.is_available():62            torch.mps.empty_cache()63        elif torch.cuda.is_available():64            torch.cuda.empty_cache()65        gc.collect()66    67    def get_device(self) -> str:68        """Get the appropriate device for processing."""69        return self.device70    71    def is_gpu_available(self) -> bool:72        """Check if GPU (CUDA or MPS) is available."""73        return torch.cuda.is_available() or torch.backends.mps.is_available()74    75    def is_spaces_environment(self) -> bool:76        """Check if running in HF Spaces environment."""77        return self.is_spaces78    79    def get_memory_info(self) -> dict:80        """Get current GPU memory information (CUDA or MPS)."""81        if torch.cuda.is_available():82            return {83                "available": True,84                "device": "cuda",85                "allocated": torch.cuda.memory_allocated(),86                "cached": torch.cuda.memory_reserved(),87                "total": torch.cuda.get_device_properties(0).total_memory88            }89        elif torch.backends.mps.is_available():90            return {91                "available": True,92                "device": "mps",93                "allocated": torch.mps.current_allocated_memory(),94                "driver_allocated": torch.mps.driver_allocated_memory(),95                # MPS doesn't have total memory info readily available96                "total": "N/A (MPS)"97            }98        else:99            return {"available": False, "device": "cpu"}100 101 102# Convenience decorators103def gpu_inference(duration: int = 60):104    """Decorator for GPU inference tasks."""105    return ZeroGPUManager.gpu_task(duration=duration)106 107 108def gpu_model_loading(duration: int = 120):109    """Decorator for GPU model loading tasks."""110    return ZeroGPUManager.gpu_task(duration=duration)111 112 113def gpu_long_task(duration: int = 300):114    """Decorator for long GPU processing tasks."""115    return ZeroGPUManager.gpu_task(duration=duration)