utils/device.py

"""
utils/device.py
───────────────
Hardware detection helpers used by every model loader.
"""

import logging
import torch

logger = logging.getLogger(__name__)


def best_device() -> torch.device:
    """Return the best available compute device."""
    if torch.cuda.is_available():
        dev = torch.device("cuda")
        logger.info("Using CUDA: %s", torch.cuda.get_device_name(0))
    elif torch.backends.mps.is_available():
        dev = torch.device("mps")
        logger.info("Using Apple MPS")
    else:
        dev = torch.device("cpu")
        logger.info("Using CPU — inference will be slow")
    return dev


def torch_dtype(device: torch.device) -> torch.dtype:
    """Choose float16 on GPU, float32 on CPU (avoids CPU fp16 ops)."""
    return torch.float16 if device.type in ("cuda", "mps") else torch.float32


def vram_gb() -> float:
    """Return available VRAM in GB, 0.0 if no CUDA."""
    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(0)
        return props.total_memory / (1024 ** 3)
    return 0.0


def log_memory_stats() -> None:
    if torch.cuda.is_available():
        alloc = torch.cuda.memory_allocated() / (1024 ** 2)
        reserved = torch.cuda.memory_reserved() / (1024 ** 2)
        logger.debug("CUDA memory — allocated: %.1f MB  reserved: %.1f MB", alloc, reserved)


def clear_cache() -> None:
    """Free GPU cache between pipeline stages."""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    import gc
    gc.collect()