From b1bc98723969591bac8f5478953e30f6f314965c Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Wed, 9 Sep 2026 14:10:22 -0400
Subject: [PATCH 1/4] fix: validate gpu counts and custom vram capacities

How to test: python -m pytest -q (75 passed); eight invalid target cases failed before the fix.
---
 src/modelinfo/hardware.py           | 14 ++++++++++----
 tests/test_gpu_target_validation.py | 21 +++++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_gpu_target_validation.py

diff --git a/src/modelinfo/hardware.py b/src/modelinfo/hardware.py
index a785944..4aeffd1 100644
--- a/src/modelinfo/hardware.py
+++ b/src/modelinfo/hardware.py
@@ -1,3 +1,4 @@
+import math
 import re
 import subprocess
 from typing import Optional, Tuple
@@ -341,6 +342,8 @@ def resolve_gpu(target: str) -> Tuple[str, float, int]:
     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 +355,16 @@ def resolve_gpu(target: str) -> Tuple[str, float, int]:
         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
 
diff --git a/tests/test_gpu_target_validation.py b/tests/test_gpu_target_validation.py
new file mode 100644
index 0000000..6b6e597
--- /dev/null
+++ b/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)

From bd5703e152851976e6dd012cf2d53a5237058d51 Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Tue, 15 Sep 2026 00:20:11 -0400
Subject: [PATCH 2/4] Parse ROCm memory totals after device prefixes

---
 src/modelinfo/hardware.py        |  5 ++---
 tests/test_rocm_memory_prefix.py | 10 ++++++++++
 2 files changed, 12 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_rocm_memory_prefix.py

diff --git a/src/modelinfo/hardware.py b/src/modelinfo/hardware.py
index 4aeffd1..8ce49f8 100644
--- a/src/modelinfo/hardware.py
+++ b/src/modelinfo/hardware.py
@@ -210,9 +210,8 @@ def _detect_amd_gpu() -> Optional[Tuple[str, float, int]]:
             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"
             )
diff --git a/tests/test_rocm_memory_prefix.py b/tests/test_rocm_memory_prefix.py
new file mode 100644
index 0000000..d3f51eb
--- /dev/null
+++ b/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)

From e40cc0333e29333ffe886f4cad889ffa165577e6 Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Tue, 15 Sep 2026 00:21:23 -0400
Subject: [PATCH 3/4] Check Apple Silicon hardware before applying
 unified-memory estimates

---
 src/modelinfo/hardware.py             |  6 ++++++
 tests/test_apple_silicon_detection.py | 20 ++++++++++++++++++++
 tests/test_hardware.py                |  2 ++
 3 files changed, 28 insertions(+)
 create mode 100644 tests/test_apple_silicon_detection.py

diff --git a/src/modelinfo/hardware.py b/src/modelinfo/hardware.py
index 8ce49f8..5168f31 100644
--- a/src/modelinfo/hardware.py
+++ b/src/modelinfo/hardware.py
@@ -283,6 +283,12 @@ def _detect_intel_gpu() -> Optional[Tuple[str, float, int]]:
 
 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,
diff --git a/tests/test_apple_silicon_detection.py b/tests/test_apple_silicon_detection.py
new file mode 100644
index 0000000..cf12ae7
--- /dev/null
+++ b/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)
diff --git a/tests/test_hardware.py b/tests/test_hardware.py
index 60d84c4..3d0b47e 100644
--- a/tests/test_hardware.py
+++ b/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")
 

From f7930e7aba632884367905552effb275ffd7d1c7 Mon Sep 17 00:00:00 2001
From: Rupayon Haldar <80724680+rupayon123@users.noreply.github.com>
Date: Tue, 15 Sep 2026 00:22:46 -0400
Subject: [PATCH 4/4] Validate complete Intel GPU memory values and units

---
 src/modelinfo/hardware.py             | 20 +++++++++++---------
 tests/test_intel_memory_validation.py | 11 +++++++++++
 2 files changed, 22 insertions(+), 9 deletions(-)
 create mode 100644 tests/test_intel_memory_validation.py

diff --git a/src/modelinfo/hardware.py b/src/modelinfo/hardware.py
index 5168f31..1ef5c10 100644
--- a/src/modelinfo/hardware.py
+++ b/src/modelinfo/hardware.py
@@ -222,18 +222,20 @@ def _detect_amd_gpu() -> Optional[Tuple[str, float, int]]:
 
 
 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]:
diff --git a/tests/test_intel_memory_validation.py b/tests/test_intel_memory_validation.py
new file mode 100644
index 0000000..8aec260
--- /dev/null
+++ b/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