garmentiq.utils.device

  1"""Device resolution and accelerator memory management.
  2
  3GarmentIQ runs on CPU by default and never auto-detects an accelerator. These helpers
  4validate a requested device, fail fast with an actionable message when it is
  5unavailable, move processor output onto it, and release its cached memory afterwards.
  6"""
  7import torch
  8from typing import Union
  9
 10DeviceLike = Union[str, torch.device, None]
 11
 12
 13def resolve_device(device: DeviceLike = "cpu") -> torch.device:
 14    """
 15    Resolves and validates a user supplied device specification.
 16
 17    GarmentIQ runs on CPU by default. Hardware acceleration is opt-in: the caller
 18    explicitly requests it by passing `device="cuda"`, `device="cuda:1"`, or
 19    `device="mps"` (Apple Silicon). This function normalises the request into a
 20    `torch.device` and fails fast with an actionable message when the requested
 21    backend is unavailable, instead of silently degrading to CPU.
 22
 23    Args:
 24        device (Union[str, torch.device], optional): The requested device, e.g. `"cpu"`,
 25                                                     `"cuda"`, `"cuda:0"`, or `"mps"`.
 26                                                     `None` is treated as `"cpu"`.
 27                                                     Default is `"cpu"`.
 28
 29    Raises:
 30        ValueError: If `device` is not a valid device specification, or if the requested
 31                    accelerator (CUDA or MPS) is not available on the current machine.
 32
 33    Returns:
 34        torch.device: The validated device to run computation on.
 35    """
 36    if device is None:
 37        return torch.device("cpu")
 38
 39    if isinstance(device, torch.device):
 40        resolved = device
 41    else:
 42        try:
 43            resolved = torch.device(device)
 44        except (RuntimeError, TypeError, ValueError) as e:
 45            raise ValueError(
 46                f"Invalid device specification {device!r}. Expected values such as "
 47                f'"cpu", "cuda", "cuda:0", or "mps". Original error: {e}'
 48            )
 49
 50    if resolved.type == "cuda" and not torch.cuda.is_available():
 51        raise ValueError(
 52            f"Requested device {str(resolved)!r} but CUDA is not available in this "
 53            f"PyTorch installation. Use device='cpu', or device='mps' on Apple Silicon."
 54        )
 55
 56    if resolved.type == "mps" and not (
 57        hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
 58    ):
 59        raise ValueError(
 60            f"Requested device {str(resolved)!r} but the MPS backend is not available. "
 61            f"MPS requires macOS on Apple Silicon with a compatible PyTorch build. "
 62            f"Use device='cpu' instead."
 63        )
 64
 65    return resolved
 66
 67
 68def inputs_to_device(inputs, device: DeviceLike = "cpu"):
 69    """
 70    Moves a processor's batched output onto `device`, narrowing dtypes MPS cannot hold.
 71
 72    Hugging Face image processors emit some tensors as float64 — SAM's `input_points`
 73    and `input_boxes` are built from plain Python floats, for example. The MPS backend
 74    has no float64 support at all, so moving such a batch onto an Apple Silicon GPU
 75    raises `TypeError: Cannot convert a MPS Tensor to float64 dtype`. Narrowing those
 76    tensors to float32 is lossless in practice, because they carry pixel coordinates
 77    rather than values that need double precision.
 78
 79    The narrowing is applied **only** when the target is MPS. CPU and CUDA keep the
 80    processor's original dtypes, so their numerical results are unchanged.
 81
 82    Args:
 83        inputs (BatchFeature | BatchEncoding | dict): The processor output to move.
 84        device (Union[str, torch.device], optional): The device to move onto, e.g.
 85                                                     `"cpu"`, `"cuda"`, or `"mps"`.
 86                                                     Default is `"cpu"`.
 87
 88    Raises:
 89        ValueError: If the requested `device` is invalid or unavailable on this machine.
 90
 91    Returns:
 92        The same mapping, with its tensors placed on `device`.
 93    """
 94    resolved = resolve_device(device)
 95
 96    if resolved.type == "mps":
 97        for key, value in list(inputs.items()):
 98            if isinstance(value, torch.Tensor) and value.dtype == torch.float64:
 99                inputs[key] = value.to(torch.float32)
100
101    return inputs.to(resolved)
102
103
104def empty_cache(device: DeviceLike = "cpu") -> None:
105    """
106    Releases cached accelerator memory held by the allocator after inference or training.
107
108    PyTorch caching allocators keep freed blocks reserved for reuse, which inflates
109    reported memory usage and can starve other processes on the same accelerator.
110    This helper clears that cache for the backend actually in use, and is a no-op on
111    CPU (where no such cache exists) and for accelerators that are not available.
112
113    Args:
114        device (Union[str, torch.device], optional): The device whose cache should be
115                                                     released. Default is `"cpu"`.
116
117    Returns:
118        None
119    """
120    if device is None:
121        return
122
123    if isinstance(device, torch.device):
124        resolved = device
125    else:
126        try:
127            resolved = torch.device(device)
128        except (RuntimeError, TypeError, ValueError):
129            return
130
131    if resolved.type == "cuda" and torch.cuda.is_available():
132        torch.cuda.empty_cache()
133    elif resolved.type == "mps" and hasattr(torch, "mps"):
134        torch.mps.empty_cache()