Merge pull request #27 from unslothai/feature/implement-silicon-utils-compatibility

[MLX] Add Hardware Detection Module & Apple Silicon (MLX) Compatibility

studio/backend/core/__init__.py

@@ -13,7 +13,8 @@ from utils.models import is_vision_model, ModelConfig, scan_trained_loras, load_
 
 # Utilities (from utils)
 from utils.paths import normalize_path, is_local_path, is_model_cached
-from utils.utils import without_hf_auth, format_error_message, get_gpu_memory_info, search_hf_models
+from utils.utils import without_hf_auth, format_error_message
+from utils.hardware import get_device, is_apple_silicon, clear_gpu_cache, get_gpu_memory_info, log_gpu_memory, DeviceType
 from utils.datasets import format_and_template_dataset
 
 __all__ = [
@@ -37,7 +38,6 @@ __all__ = [
     'get_base_model_from_lora',
 
     # Utils
-    'search_hf_models',
     'format_and_template_dataset',
     'normalize_path',
     'is_local_path',
@@ -45,4 +45,9 @@ __all__ = [
     'without_hf_auth',
     'format_error_message',
     'get_gpu_memory_info',
+    'log_gpu_memory',
+    'get_device',
+    'is_apple_silicon',
+    'clear_gpu_cache',
+    'DeviceType',
 ]

studio/backend/main.py

@@ -15,13 +15,18 @@ from datetime import datetime
 # Import routers
 from routes import training_router, models_router, inference_router, datasets_router, auth_router
 from auth import storage
+from utils.hardware import detect_hardware
+import utils.hardware.hardware as _hw_module
 
 UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache"
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Startup: print setup token if needed. Shutdown: clean up compiled cache."""
+    """Startup: detect hardware, print setup token if needed. Shutdown: clean up compiled cache."""
+    # Detect hardware first — sets DEVICE global used everywhere
+    detect_hardware()
+
     if not storage.is_initialized():
         setup_token = secrets.token_urlsafe(32)
         storage.save_setup_token(setup_token)
@@ -32,7 +37,8 @@ async def lifespan(app: FastAPI):
         print("This token can only be used once.")
         print("=" * 60 + "\n")
     yield
-    # Cleanup: remove Unsloth compiled cache on shutdown
+    # Cleanup
+    _hw_module.DEVICE = None
     shutil.rmtree(UNSLOTH_CACHE_DIR, ignore_errors=True)
 
 
@@ -78,23 +84,20 @@ async def health_check():
 @app.get("/api/system")
 async def get_system_info():
     """Get system information"""
-    import torch
     import platform
     import psutil
+    from utils.hardware import get_device, get_gpu_memory_info, DeviceType
 
-    # GPU Info
-    gpu_info = {"available": False, "devices": []}
-    if torch.cuda.is_available():
-        gpu_info["available"] = True
-        for i in range(torch.cuda.device_count()):
-            props = torch.cuda.get_device_properties(i)
-            gpu_info["devices"].append(
-                {
-                    "index": i,
-                    "name": props.name,
-                    "memory_total_gb": round(props.total_memory / 1e9, 2),
-                }
-            )
+    # GPU Info — uses the hardware module (works on CUDA, MPS, CPU)
+    mem_info = get_gpu_memory_info()
+    gpu_info = {"available": mem_info.get("available", False), "devices": []}
+
+    if mem_info.get("available"):
+        gpu_info["devices"].append({
+            "index": mem_info.get("device", 0),
+            "name": mem_info.get("device_name", "Unknown"),
+            "memory_total_gb": round(mem_info.get("total_gb", 0), 2),
+        })
 
     # CPU & Memory
     memory = psutil.virtual_memory()
@@ -102,6 +105,7 @@ async def get_system_info():
     return {
         "platform": platform.platform(),
         "python_version": platform.python_version(),
+        "device_backend": get_device().value,
         "cpu_count": psutil.cpu_count(),
         "memory": {
             "total_gb": round(memory.total / 1e9, 2),

studio/backend/tests/conftest.py

@@ -0,0 +1,12 @@
+"""
+Shared pytest configuration for the backend test suite.
+Ensures that the backend root is on sys.path so that
+`import utils.utils` (and similar flat imports) resolve correctly.
+"""
+import sys
+from pathlib import Path
+
+# Add backend root to sys.path (mirrors how the app itself is launched)
+_backend_root = Path(__file__).resolve().parent.parent
+if str(_backend_root) not in sys.path:
+    sys.path.insert(0, str(_backend_root))

studio/backend/tests/test_utils.py

@@ -0,0 +1,350 @@
+"""
+Tests for utils/hardware and utils/utils — device detection, GPU memory, error formatting.
+
+These tests are designed to pass on ANY platform:
+  • NVIDIA GPU  (CUDA backend, requires torch)
+  • Apple Silicon (MLX backend, requires mlx)
+  • CPU-only     (no GPU at all)
+
+No ML framework is imported at the top level.
+Tests that need torch/mlx internals for mocking are skipped when unavailable.
+
+Run with:
+    cd studio/backend
+    python -m pytest tests/test_utils.py -v
+"""
+import platform
+from unittest.mock import patch, MagicMock
+
+import pytest
+
+# --- Conditional framework imports ---
+try:
+    import torch
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+try:
+    import mlx.core as mx
+    HAS_MLX = True
+except ImportError:
+    HAS_MLX = False
+
+needs_torch = pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not installed")
+needs_mlx = pytest.mark.skipif(not HAS_MLX, reason="MLX not installed")
+
+from utils.hardware import (
+    get_device,
+    detect_hardware,
+    is_apple_silicon,
+    clear_gpu_cache,
+    get_gpu_memory_info,
+    log_gpu_memory,
+    DeviceType,
+)
+import utils.hardware.hardware as _hw_module
+from utils.utils import format_error_message
+
+
+# ========== Helpers ==========
+
+def _actual_device() -> str:
+    """Return the real device string for the current machine."""
+    if HAS_TORCH and torch.cuda.is_available():
+        return "cuda"
+    if is_apple_silicon() and HAS_MLX:
+        return "mlx"
+    return "cpu"
+
+
+def _reset_and_detect():
+    """Reset the cached DEVICE global and re-run detection."""
+    _hw_module.DEVICE = None
+    return detect_hardware()
+
+
+# ========== get_device() ==========
+
+class TestGetDevice:
+    """Tests for get_device() — should agree with the real hardware."""
+
+    def setup_method(self):
+        self._saved_device = _hw_module.DEVICE
+
+    def teardown_method(self):
+        _hw_module.DEVICE = self._saved_device
+
+    def test_returns_valid_device_type(self):
+        result = get_device()
+        assert result in (DeviceType.CUDA, DeviceType.MLX, DeviceType.CPU)
+
+    def test_matches_actual_hardware(self):
+        assert get_device().value == _actual_device()
+
+    # --- Mocked paths ---
+
+    @needs_torch
+    def test_returns_cuda_when_cuda_available(self):
+        with patch("utils.hardware.hardware._has_torch", return_value=True), \
+             patch("torch.cuda.is_available", return_value=True):
+            assert _reset_and_detect() == DeviceType.CUDA
+
+    @needs_mlx
+    def test_returns_mlx_when_on_apple_silicon_with_mlx(self):
+        with patch("utils.hardware.hardware._has_torch", return_value=False), \
+             patch("utils.hardware.hardware.is_apple_silicon", return_value=True), \
+             patch("utils.hardware.hardware._has_mlx", return_value=True):
+            assert _reset_and_detect() == DeviceType.MLX
+
+    def test_returns_cpu_when_nothing_available(self):
+        with patch("utils.hardware.hardware._has_torch", return_value=False), \
+             patch("utils.hardware.hardware.is_apple_silicon", return_value=False), \
+             patch("utils.hardware.hardware._has_mlx", return_value=False):
+            assert _reset_and_detect() == DeviceType.CPU
+
+
+# ========== is_apple_silicon() ==========
+
+class TestIsAppleSilicon:
+
+    def test_returns_bool(self):
+        assert isinstance(is_apple_silicon(), bool)
+
+    def test_true_on_darwin_arm64(self):
+        with patch("utils.hardware.hardware.platform") as mock_plat:
+            mock_plat.system.return_value = "Darwin"
+            mock_plat.machine.return_value = "arm64"
+            assert is_apple_silicon() is True
+
+    def test_false_on_linux_x86(self):
+        with patch("utils.hardware.hardware.platform") as mock_plat:
+            mock_plat.system.return_value = "Linux"
+            mock_plat.machine.return_value = "x86_64"
+            assert is_apple_silicon() is False
+
+    def test_false_on_darwin_x86(self):
+        """Intel Mac should return False."""
+        with patch("utils.hardware.hardware.platform") as mock_plat:
+            mock_plat.system.return_value = "Darwin"
+            mock_plat.machine.return_value = "x86_64"
+            assert is_apple_silicon() is False
+
+
+# ========== clear_gpu_cache() ==========
+
+class TestClearGpuCache:
+    """clear_gpu_cache() must never raise, regardless of platform."""
+
+    def test_does_not_raise(self):
+        clear_gpu_cache()
+
+    @needs_torch
+    def test_calls_cuda_cache_when_cuda(self):
+        with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
+             patch("torch.cuda.empty_cache") as mock_empty, \
+             patch("torch.cuda.ipc_collect") as mock_ipc:
+            clear_gpu_cache()
+            mock_empty.assert_called_once()
+            mock_ipc.assert_called_once()
+
+    @needs_mlx
+    def test_mlx_does_not_raise(self):
+        """MLX cache clear is a no-op — should just succeed."""
+        with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX):
+            clear_gpu_cache()
+
+    def test_noop_on_cpu(self):
+        with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU):
+            clear_gpu_cache()
+
+
+# ========== get_gpu_memory_info() ==========
+
+class TestGetGpuMemoryInfo:
+
+    def test_returns_dict(self):
+        result = get_gpu_memory_info()
+        assert isinstance(result, dict)
+
+    def test_has_available_key(self):
+        assert "available" in get_gpu_memory_info()
+
+    def test_has_backend_key(self):
+        assert "backend" in get_gpu_memory_info()
+
+    def test_backend_matches_device(self):
+        result = get_gpu_memory_info()
+        assert result["backend"] == get_device().value
+
+    # --- When a GPU IS available ---
+
+    @pytest.mark.skipif(
+        _actual_device() == "cpu",
+        reason="No GPU available on this machine"
+    )
+    def test_gpu_available_fields(self):
+        result = get_gpu_memory_info()
+        assert result["available"] is True
+        assert result["total_gb"] > 0
+        assert result["allocated_gb"] >= 0
+        assert result["free_gb"] >= 0
+        assert 0 <= result["utilization_pct"] <= 100
+        assert "device_name" in result
+
+    # --- CUDA-specific mocked test ---
+
+    @needs_torch
+    def test_cuda_path_returns_correct_fields(self):
+        mock_props = MagicMock()
+        mock_props.total_memory = 16 * (1024 ** 3)
+        mock_props.name = "NVIDIA Test GPU"
+
+        with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
+             patch("torch.cuda.current_device", return_value=0), \
+             patch("torch.cuda.get_device_properties", return_value=mock_props), \
+             patch("torch.cuda.memory_allocated", return_value=4 * (1024 ** 3)), \
+             patch("torch.cuda.memory_reserved", return_value=6 * (1024 ** 3)):
+            result = get_gpu_memory_info()
+
+        assert result["available"] is True
+        assert result["backend"] == "cuda"
+        assert result["device_name"] == "NVIDIA Test GPU"
+        assert abs(result["total_gb"] - 16.0) < 0.01
+        assert abs(result["allocated_gb"] - 4.0) < 0.01
+        assert abs(result["free_gb"] - 12.0) < 0.01
+        assert abs(result["utilization_pct"] - 25.0) < 0.1
+
+    # --- MLX-specific mocked test ---
+
+    @needs_mlx
+    def test_mlx_path_returns_correct_fields(self):
+        mock_psutil_mem = MagicMock()
+        mock_psutil_mem.total = 32 * (1024 ** 3)  # 32 GB unified
+
+        mock_psutil = MagicMock()
+        mock_psutil.virtual_memory.return_value = mock_psutil_mem
+
+        with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX), \
+             patch.dict("sys.modules", {"psutil": mock_psutil}):
+            result = get_gpu_memory_info()
+
+        assert result["available"] is True
+        assert result["backend"] == "mlx"
+        assert "Apple Silicon" in result["device_name"]
+        assert abs(result["total_gb"] - 32.0) < 0.01
+
+    # --- CPU-only path ---
+
+    def test_cpu_path_returns_unavailable(self):
+        with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU):
+            result = get_gpu_memory_info()
+        assert result["available"] is False
+        assert result["backend"] == "cpu"
+
+    # --- Error resilience ---
+
+    @needs_torch
+    def test_cuda_error_returns_unavailable(self):
+        with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \
+             patch("torch.cuda.current_device", side_effect=RuntimeError("CUDA init failed")):
+            result = get_gpu_memory_info()
+        assert result["available"] is False
+        assert "error" in result
+
+
+# ========== log_gpu_memory() ==========
+
+class TestLogGpuMemory:
+
+    def test_does_not_raise(self):
+        log_gpu_memory("test")
+
+    def test_logs_gpu_info_when_available(self, caplog):
+        fake_info = {
+            "available": True,
+            "backend": "cuda",
+            "device_name": "FakeGPU",
+            "allocated_gb": 2.0,
+            "total_gb": 16.0,
+            "utilization_pct": 12.5,
+            "free_gb": 14.0,
+        }
+        import logging
+        with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \
+             caplog.at_level(logging.INFO, logger="utils.hardware.hardware"):
+            log_gpu_memory("unit-test")
+
+        assert "unit-test" in caplog.text
+        assert "CUDA" in caplog.text
+        assert "FakeGPU" in caplog.text
+
+    def test_logs_cpu_fallback_when_no_gpu(self, caplog):
+        fake_info = {"available": False, "backend": "cpu"}
+        import logging
+        with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \
+             caplog.at_level(logging.INFO, logger="utils.hardware.hardware"):
+            log_gpu_memory("cpu-test")
+
+        assert "No GPU available" in caplog.text
+
+
+# ========== format_error_message() ==========
+
+class TestFormatErrorMessage:
+
+    def test_not_found(self):
+        err = Exception("Repository not found for unsloth/test")
+        msg = format_error_message(err, "unsloth/test")
+        assert "not found" in msg.lower()
+        assert "test" in msg
+
+    def test_unauthorized(self):
+        err = Exception("401 Unauthorized")
+        msg = format_error_message(err, "some/model")
+        assert "authentication" in msg.lower() or "unauthorized" in msg.lower()
+
+    def test_gated_model(self):
+        err = Exception("Access to model requires authentication")
+        msg = format_error_message(err, "meta/llama")
+        assert "authentication" in msg.lower()
+
+    def test_invalid_token(self):
+        err = Exception("Invalid user token")
+        msg = format_error_message(err, "any/model")
+        assert "invalid" in msg.lower()
+
+    # --- OOM on CUDA ---
+
+    @needs_torch
+    def test_cuda_oom(self):
+        err = Exception("CUDA out of memory")
+        with patch("utils.hardware.get_device", return_value=DeviceType.CUDA):
+            msg = format_error_message(err, "big/model")
+        assert "GPU" in msg
+        assert "big/model" not in msg
+        assert "model" in msg
+
+    # --- OOM on MLX ---
+
+    @needs_mlx
+    def test_mlx_oom(self):
+        err = Exception("MLX backend out of memory")
+        with patch("utils.hardware.get_device", return_value=DeviceType.MLX):
+            msg = format_error_message(err, "unsloth/huge-model")
+        assert "Apple Silicon" in msg
+
+    # --- OOM on CPU ---
+
+    def test_cpu_oom(self):
+        err = Exception("not enough memory to allocate")
+        with patch("utils.hardware.get_device", return_value=DeviceType.CPU):
+            msg = format_error_message(err, "any/model")
+        assert "system" in msg.lower()
+
+    # --- Generic fallback ---
+
+    def test_generic_error(self):
+        err = Exception("Something completely unexpected")
+        msg = format_error_message(err, "any/model")
+        assert msg == "Something completely unexpected"

studio/backend/utils/hardware/__init__.py

@@ -0,0 +1,24 @@
+"""
+Hardware detection and GPU utilities
+"""
+from .hardware import (
+    DeviceType,
+    DEVICE,
+    detect_hardware,
+    get_device,
+    is_apple_silicon,
+    clear_gpu_cache,
+    get_gpu_memory_info,
+    log_gpu_memory,
+)
+
+__all__ = [
+    'DeviceType',
+    'DEVICE',
+    'detect_hardware',
+    'get_device',
+    'is_apple_silicon',
+    'clear_gpu_cache',
+    'get_gpu_memory_info',
+    'log_gpu_memory',
+]

studio/backend/utils/hardware/hardware.py

@@ -0,0 +1,208 @@
+"""
+Hardware detection — run once at startup, read everywhere.
+
+Usage:
+    # At FastAPI lifespan startup:
+    from utils.hardware import detect_hardware
+    detect_hardware()
+
+    # Anywhere else:
+    from utils.hardware import DEVICE, DeviceType, is_apple_silicon
+    if DEVICE == DeviceType.CUDA:
+        import torch
+        ...
+"""
+import platform
+import logging
+from enum import Enum
+from typing import Optional, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+
+# ========== Device Enum ==========
+
+class DeviceType(str, Enum):
+    """Supported compute backends. Inherits from str so it serializes cleanly in JSON."""
+    CUDA = "cuda"
+    MLX  = "mlx"
+    CPU  = "cpu"
+
+
+# ========== Global State (set once by detect_hardware) ==========
+
+DEVICE: Optional[DeviceType] = None
+
+
+# ========== Detection ==========
+
+def is_apple_silicon() -> bool:
+    """Check if running on Apple Silicon hardware (pure platform check, no ML imports)."""
+    return platform.system() == "Darwin" and platform.machine() == "arm64"
+
+
+def _has_torch() -> bool:
+    """Check if PyTorch is importable."""
+    try:
+        import torch
+        return True
+    except ImportError:
+        return False
+
+
+def _has_mlx() -> bool:
+    """Check if MLX is importable."""
+    try:
+        import mlx.core
+        return True
+    except ImportError:
+        return False
+
+
+def detect_hardware() -> DeviceType:
+    """
+    Detect the best available compute device and set the module-level DEVICE global.
+
+    Should be called exactly once during FastAPI lifespan startup.
+    Safe to call multiple times (idempotent).
+
+    Detection order:
+      1. CUDA  (NVIDIA GPU, requires torch)
+      2. MLX   (Apple Silicon via MLX framework)
+      3. CPU   (fallback)
+    """
+    global DEVICE
+
+    # --- CUDA: try PyTorch ---
+    if _has_torch():
+        import torch
+        if torch.cuda.is_available():
+            DEVICE = DeviceType.CUDA
+            device_name = torch.cuda.get_device_properties(0).name
+            logger.info(f"Hardware detected: CUDA — {device_name}")
+            return DEVICE
+
+    # --- MLX: Apple Silicon ---
+    if is_apple_silicon() and _has_mlx():
+        DEVICE = DeviceType.MLX
+        chip = platform.processor() or platform.machine()
+        logger.info(f"Hardware detected: MLX — Apple Silicon ({chip})")
+        return DEVICE
+
+    # --- Fallback ---
+    DEVICE = DeviceType.CPU
+    logger.info("Hardware detected: CPU (no GPU backend available)")
+    return DEVICE
+
+
+# ========== Convenience helpers ==========
+
+def get_device() -> DeviceType:
+    """
+    Return the detected device. Auto-detects if detect_hardware() hasn't been called yet.
+    Prefer calling detect_hardware() explicitly at startup instead.
+    """
+    global DEVICE
+    if DEVICE is None:
+        detect_hardware()
+    return DEVICE
+
+
+def clear_gpu_cache():
+    """
+    Clear GPU memory cache for the current device.
+    Safe to call on any platform — no-ops gracefully.
+    """
+    import gc
+    gc.collect()
+
+    device = get_device()
+
+    if device == DeviceType.CUDA:
+        import torch
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+    elif device == DeviceType.MLX:
+        # MLX manages memory automatically; no explicit cache clear needed.
+        # mlx.core has no empty_cache equivalent — gc.collect() above is enough.
+        pass
+
+
+def get_gpu_memory_info() -> Dict[str, Any]:
+    """
+    Get GPU memory information.
+    Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only environments.
+    """
+    device = get_device()
+
+    # ---- CUDA path ----
+    if device == DeviceType.CUDA:
+        try:
+            import torch
+            idx = torch.cuda.current_device()
+            props = torch.cuda.get_device_properties(idx)
+
+            total = props.total_memory
+            allocated = torch.cuda.memory_allocated(idx)
+            reserved = torch.cuda.memory_reserved(idx)
+
+            return {
+                "available": True,
+                "backend": device.value,
+                "device": idx,
+                "device_name": props.name,
+                "total_gb": total / (1024**3),
+                "allocated_gb": allocated / (1024**3),
+                "reserved_gb": reserved / (1024**3),
+                "free_gb": (total - allocated) / (1024**3),
+                "utilization_pct": (allocated / total) * 100,
+            }
+        except Exception as e:
+            logger.error(f"Error getting CUDA GPU info: {e}")
+            return {"available": False, "backend": device.value, "error": str(e)}
+
+    # ---- MLX path (Apple Silicon) ----
+    if device == DeviceType.MLX:
+        try:
+            import mlx.core as mx
+            import psutil
+
+            # MLX uses unified memory — report system memory as the pool
+            total = psutil.virtual_memory().total
+            # MLX doesn't expose per-process GPU allocation; report 0 as allocated
+            allocated = 0
+
+            return {
+                "available": True,
+                "backend": device.value,
+                "device": 0,
+                "device_name": f"Apple Silicon ({platform.processor() or platform.machine()})",
+                "total_gb": total / (1024**3),
+                "allocated_gb": allocated / (1024**3),
+                "reserved_gb": 0,
+                "free_gb": (total - allocated) / (1024**3),
+                "utilization_pct": (allocated / total) * 100 if total else 0,
+            }
+        except Exception as e:
+            logger.error(f"Error getting MLX GPU info: {e}")
+            return {"available": False, "backend": device.value, "error": str(e)}
+
+    # ---- CPU-only ----
+    return {"available": False, "backend": "cpu"}
+
+
+def log_gpu_memory(context: str):
+    """Log GPU memory usage with context."""
+    memory_info = get_gpu_memory_info()
+    if memory_info.get("available"):
+        backend = memory_info.get("backend", "unknown").upper()
+        device_name = memory_info.get("device_name", "")
+        label = f"{backend}" + (f" ({device_name})" if device_name else "")
+        logger.info(
+            f"GPU Memory [{context}] {label}: "
+            f"{memory_info['allocated_gb']:.2f}GB/{memory_info['total_gb']:.2f}GB "
+            f"({memory_info['utilization_pct']:.1f}% used, "
+            f"{memory_info['free_gb']:.2f}GB free)"
+        )
+    else:
+        logger.info(f"GPU Memory [{context}]: No GPU available (CPU-only)")

studio/backend/utils/utils.py

@@ -1,18 +1,17 @@
 """
 Shared backend utilities
 """
-import gradio as gr
 import os
 import logging
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Optional, Dict, Any
 import shutil
 import tempfile
 
 
 logger = logging.getLogger(__name__)
 
+
 @contextmanager
 def without_hf_auth():
     """
@@ -96,113 +95,12 @@ def format_error_message(error: Exception, model_name: str) -> str:
     if "invalid user token" in error_str:
         return "Invalid HF token. Please check your token and try again."
 
-    if "memory" in error_str or "cuda" in error_str or "out of memory" in error_str:
-        return f"Not enough GPU memory to load '{model_short}'. Try a smaller model or free GPU memory."
+    if "memory" in error_str or "cuda" in error_str or "mlx" in error_str or "out of memory" in error_str:
+        from utils.hardware import get_device
+        device = get_device()
+        device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get(device.value, "GPU")
+        return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory."
 
     # Generic fallback
     return str(error)
 pass
-
-def get_gpu_memory_info() -> Dict[str, Any]:
-    """Get GPU memory information."""
-    import torch
-
-    if not torch.cuda.is_available():
-        return {"available": False}
-
-    try:
-        device = torch.cuda.current_device()
-        props = torch.cuda.get_device_properties(device)
-
-        total = props.total_memory
-        allocated = torch.cuda.memory_allocated(device)
-        reserved = torch.cuda.memory_reserved(device)
-
-        return {
-            "available": True,
-            "device": device,
-            "total_gb": total / (1024**3),
-            "allocated_gb": allocated / (1024**3),
-            "reserved_gb": reserved / (1024**3),
-            "free_gb": (total - allocated) / (1024**3),
-            "utilization_pct": (allocated / total) * 100
-        }
-    except Exception as e:
-        logger.error(f"Error getting GPU info: {e}")
-        return {"available": False, "error": str(e)}
-pass
-
-def log_gpu_memory(context: str):
-    """Log GPU memory usage with context."""
-    memory_info = get_gpu_memory_info()
-    if memory_info.get("available"):
-        logger.info(
-            f"GPU Memory [{context}]: "
-            f"{memory_info['allocated_gb']:.2f}GB/{memory_info['total_gb']:.2f}GB "
-            f"({memory_info['utilization_pct']:.1f}% used, "
-            f"{memory_info['free_gb']:.2f}GB free)"
-        )
-    else:
-        logger.info(f"GPU Memory [{context}]: No CUDA GPU available")
-pass
-
-"""
-Model utility functions - search, discovery, etc.
-"""
-
-
-def search_hf_models(search_query: str, hf_token: Optional[str] = None):
-    """
-    Search HuggingFace model hub.
-    """
-    import requests
-
-    if not search_query or not search_query.strip():
-        return gr.update(choices=[])
-
-    # Simple debouncing: only search if query is at least 2 characters
-    if len(search_query.strip()) < 2:
-        return gr.update(choices=[])
-
-    try:
-        headers = {}
-        if hf_token and hf_token.strip():
-            headers["Authorization"] = f"Bearer {hf_token.strip()}"
-
-        url = "https://huggingface.co/api/models"
-        params = {
-            "search": search_query,
-            "pipeline_tag": "text-generation",
-            "library": "transformers",
-            "limit": 15,
-            "sort": "downloads",
-            "direction": -1
-        }
-
-        response = requests.get(url, headers=headers, params=params, timeout=10)
-
-        if response.status_code == 200:
-            models = response.json()
-            unsloth_results = []
-            other_results = []
-
-            for model in models:
-                model_id = model.get("modelId", "")
-                if model_id and "gguf" not in model_id.lower():
-                    result = (f"{model_id}", model_id)
-
-                    if model_id.startswith("unsloth/"):
-                        unsloth_results.append(result)
-                    else:
-                        other_results.append(result)
-
-            # Combine with unsloth models first
-            search_results = unsloth_results + other_results
-            return gr.update(choices=search_results)
-        else:
-            logger.warning(f"HF API returned status {response.status_code}")
-            return gr.update(choices=[])
-
-    except Exception as e:
-        logger.warning(f"Model search failed: {e}")
-        return gr.update(choices=[])