Spaces:
Running on Zero
Running on Zero
File size: 3,354 Bytes
ffe9fdb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """GPU configuration for HuggingFace ZeroGPU compatibility.
This module provides configuration constants and utilities for managing GPU resources
in both local and HuggingFace Spaces ZeroGPU environments.
"""
import os
import torch
class GPUConfig:
"""GPU configuration constants and environment detection."""
# Environment detection
IS_ZEROGPU: bool = os.environ.get("SPACES_ZERO_GPU") is not None
IS_SPACES: bool = os.environ.get("SPACE_ID") is not None
# Device configuration
GPU_AVAILABLE: bool = torch.cuda.is_available()
DEFAULT_DEVICE: torch.device = torch.device(
"cuda" if GPU_AVAILABLE and not IS_ZEROGPU else "cpu"
)
# Duration limits for @spaces.GPU decorator (seconds)
# These values are based on typical processing times per workflow
SEPARATION_DURATION: int = 90 # Speaker separation (longest operation)
EXTRACTION_DURATION: int = 60 # Speaker extraction
DENOISING_DURATION: int = 45 # Voice denoising (fastest operation)
MAX_DURATION: int = 120 # Maximum allowed by ZeroGPU
# Resource management
CLEANUP_TIMEOUT: float = 2.0 # Maximum time for GPU cleanup (SC-004)
ENABLE_CACHE_CLEARING: bool = True # Clear CUDA cache after operations
@classmethod
def get_device(cls) -> torch.device:
"""Get the appropriate device for model operations.
Returns:
torch.device: CUDA device if available and not in ZeroGPU mode, else CPU
"""
return cls.DEFAULT_DEVICE
@classmethod
def get_environment_type(cls) -> str:
"""Get a string describing the current execution environment.
Returns:
str: One of "zerogpu", "local_gpu", "spaces_cpu", or "local_cpu"
"""
if cls.IS_ZEROGPU:
return "zerogpu"
elif cls.IS_SPACES:
return "spaces_cpu"
elif cls.GPU_AVAILABLE:
return "local_gpu"
else:
return "local_cpu"
@classmethod
def validate_duration(cls, duration: int, max_duration: int = None) -> int:
"""Validate and clamp duration to acceptable limits.
Args:
duration: Requested duration in seconds
max_duration: Maximum allowed duration (defaults to MAX_DURATION)
Returns:
int: Clamped duration value
Raises:
ValueError: If duration is less than 1 second
"""
if duration < 1:
raise ValueError(f"Duration must be at least 1 second, got {duration}")
max_limit = max_duration if max_duration is not None else cls.MAX_DURATION
if duration > max_limit:
return max_limit
return duration
@classmethod
def info(cls) -> dict:
"""Get a dictionary of current GPU configuration.
Returns:
dict: Configuration information for debugging and logging
"""
return {
"environment_type": cls.get_environment_type(),
"is_zerogpu": cls.IS_ZEROGPU,
"is_spaces": cls.IS_SPACES,
"gpu_available": cls.GPU_AVAILABLE,
"default_device": str(cls.DEFAULT_DEVICE),
"separation_duration": cls.SEPARATION_DURATION,
"extraction_duration": cls.EXTRACTION_DURATION,
"denoising_duration": cls.DENOISING_DURATION,
}
|