fix: validate gpu counts and custom vram capacities

src/modelinfo/hardware.py

@@ -1,3 +1,4 @@
+import math
 import re
 import subprocess
 from typing import Optional, Tuple
@@ -209,9 +210,8 @@
             total_bytes = 0
             gpu_count = len(lines)
             for line in lines:
-                parts = line.split(":")
-                if len(parts) >= 2:
-                    total_bytes += int(parts[1].strip())
+                # Device prefixes contain a colon too (GPU[0] : VRAM ...).
+                total_bytes += int(line.rsplit(":", 1)[1].strip())
             display_name = (
                 f"AMD Multi-GPU ({gpu_count}x)" if gpu_count > 1 else "AMD GPU"
             )
@@ -222,18 +222,20 @@
 
 
 def _parse_intel_vram(size_str: str) -> Optional[float]:
-    match = re.search(r"([\d\.]+)\s*([a-zA-Z]*)", size_str)
+    match = re.fullmatch(r"\s*(\d+(?:\.\d*)?|\.\d+)\s*([a-zA-Z]*)\s*", size_str)
     if not match:
         return None
     val = float(match.group(1))
-    unit = match.group(2).lower()
-    if unit in ("gib", "gb"):
-        val *= 1024.0
-    elif unit in ("kib", "kb"):
-        val /= 1024.0
-    elif unit == "b":
-        val /= (1024.0 * 1024.0)
-    return val
+    factors = {"": 1.0, "mib": 1.0, "mb": 1.0,
+               "gib": 1024.0, "gb": 1024.0,
+               "tib": 1024.0**2, "tb": 1024.0**2,
+               "kib": 1 / 1024.0, "kb": 1 / 1024.0,
+               "b": 1 / 1024.0**2}
+    factor = factors.get(match.group(2).lower())
+    if factor is None or not math.isfinite(val) or val <= 0:
+        return None
+    result = val * factor
+    return result if math.isfinite(result) else None
 
 
 def _parse_xpu_smi_output(stdout: str) -> Tuple[list[str], float, int]:
@@ -283,6 +285,12 @@
 
 def _detect_apple_gpu() -> Optional[Tuple[str, float, int]]:
     try:
+        architecture = subprocess.run(
+            ["sysctl", "-n", "hw.optional.arm64"],
+            capture_output=True, text=True, check=True, timeout=2.0,
+        )
+        if architecture.stdout.strip() != "1":
+            return None
         result = subprocess.run(
             ["sysctl", "hw.memsize"],
             capture_output=True,
@@ -341,6 +349,8 @@
     match = re.match(r"^(\d+)x\s*(.+)$", lower_target)
     if match:
         gpu_count = int(match.group(1))
+        if gpu_count < 1:
+            raise ValueError("GPU count must be at least 1.")
         target_name = match.group(2)
     else:
         target_name = target
@@ -352,13 +362,16 @@
         display_name = f"{gpu_count}x {target_name}" if gpu_count > 1 else target_name
         return display_name, vram_gb, gpu_count
 
-    # If the user passed a pure number, assume GB
+    # Parse numeric capacity before normalization can remove a minus sign.
     try:
-        vram_gb = float(normalized) * gpu_count
-        display_name = f"Custom ({vram_gb} GB)"
-        return display_name, vram_gb, gpu_count
+        vram_gb = float(target_name.strip()) * gpu_count
     except ValueError:
         pass
+    else:
+        if not math.isfinite(vram_gb) or vram_gb <= 0:
+            raise ValueError("GPU VRAM must be a finite number greater than 0.")
+        display_name = f"Custom ({vram_gb} GB)"
+        return display_name, vram_gb, gpu_count
 
     import difflib
 

tests/test_apple_silicon_detection.py

@@ -0,0 +1,20 @@
+import subprocess
+import platform
+from modelinfo import hardware
+
+
+def test_intel_mac_memory_is_not_apple_silicon_vram(monkeypatch):
+    def run(command, **kwargs):
+        text = '0\n' if command == ['sysctl', '-n', 'hw.optional.arm64'] else 'hw.memsize: 17179869184\n'
+        return subprocess.CompletedProcess(command, 0, text)
+    monkeypatch.setattr(hardware.subprocess, 'run', run)
+    assert hardware._detect_apple_gpu() is None
+
+
+def test_apple_silicon_is_detected_under_translated_process(monkeypatch):
+    monkeypatch.setattr(platform, "machine", lambda: "x86_64")
+    def run(command, **kwargs):
+        text = '1\n' if command == ['sysctl', '-n', 'hw.optional.arm64'] else 'hw.memsize: 17179869184\n'
+        return subprocess.CompletedProcess(command, 0, text)
+    monkeypatch.setattr(hardware.subprocess, 'run', run)
+    assert hardware._detect_apple_gpu() == ('Apple Silicon (Unified Memory)', 12.0, 1)

tests/test_gpu_target_validation.py

@@ -0,0 +1,21 @@
+import pytest
+
+from modelinfo.hardware import resolve_gpu
+
+
+@pytest.mark.parametrize(
+    "target", ["0x RTX4090", "0x 16", "0", "-1", "nan", "inf", "2x nan", "2x -16"]
+)
+def test_gpu_targets_reject_nonpositive_or_nonfinite_capacity(target):
+    with pytest.raises(ValueError):
+        resolve_gpu(target)
+
+
+@pytest.mark.parametrize(
+    "target,capacity,count", [("0.5", 0.5, 1), ("2x 12", 24, 2), ("RTX-4090", 24, 1)]
+)
+def test_gpu_target_validation_preserves_valid_names_and_capacities(
+    target, capacity, count
+):
+    _, actual_capacity, actual_count = resolve_gpu(target)
+    assert (actual_capacity, actual_count) == (capacity, count)

tests/test_hardware.py

@@ -238,6 +238,8 @@ def test_detect_local_gpu_falls_back_to_apple_unified_memory(monkeypatch):
     def fake_run(command, **kwargs):
         if command[0] in {"nvidia-smi", "rocm-smi", "xpu-smi"}:
             raise FileNotFoundError(command[0])
+        if command == ["sysctl", "-n", "hw.optional.arm64"]:
+            return completed("1\n")
         assert command == ["sysctl", "hw.memsize"]
         return completed("hw.memsize: 17179869184\n")
 

tests/test_intel_memory_validation.py

@@ -0,0 +1,11 @@
+import pytest
+from modelinfo.hardware import _parse_intel_vram
+
+
+@pytest.mark.parametrize('text', ['-16 GiB', '16 XB', 'unavailable 16 GiB', '1.2.3 GiB', '0 MiB'])
+def test_invalid_intel_memory_is_not_reported_as_capacity(text):
+    assert _parse_intel_vram(text) is None
+
+
+def test_intel_terabyte_capacity_converts_to_mib():
+    assert _parse_intel_vram('1 TiB') == 1024 * 1024

tests/test_rocm_memory_prefix.py

@@ -0,0 +1,10 @@
+import subprocess
+from modelinfo import hardware
+
+
+def test_rocm_memory_rows_with_device_prefix(monkeypatch):
+    output = ('GPU[0] : VRAM Total Memory (B): 17179869184\n'
+              'GPU[0] : VRAM Total Used Memory (B): 1024\n'
+              'GPU[1] : VRAM Total Memory (B): 8589934592\n')
+    monkeypatch.setattr(hardware.subprocess, 'run', lambda *a, **kw: subprocess.CompletedProcess(a, 0, output))
+    assert hardware._detect_amd_gpu() == ('AMD Multi-GPU (2x)', 24.0, 2)