From 59d5f24eb5b0eb4b1506864d95ade1840cf1fe0a Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 15:34:26 +0000 Subject: [PATCH] integrate global hardware detection at lifespan entrypoint --- studio/backend/core/__init__.py | 6 +- studio/backend/main.py | 32 ++-- studio/backend/tests/test_utils.py | 70 ++++---- studio/backend/utils/hardware/__init__.py | 24 +++ studio/backend/utils/hardware/hardware.py | 195 ++++++++++++++++++++++ studio/backend/utils/utils.py | 187 +-------------------- 6 files changed, 276 insertions(+), 238 deletions(-) create mode 100644 studio/backend/utils/hardware/__init__.py create mode 100644 studio/backend/utils/hardware/hardware.py diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 5fdb57ae7f..7b562c6cb0 100644 --- a/studio/backend/core/__init__.py +++ b/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, get_device, is_apple_silicon, clear_gpu_cache +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,7 +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', ] diff --git a/studio/backend/main.py b/studio/backend/main.py index b08ef3430e..de468c183a 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -15,13 +15,17 @@ 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 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) @@ -78,23 +82,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 +103,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), diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 32c82cacb0..89b16dcd0e 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -1,5 +1,5 @@ """ -Tests for utils/utils.py — device detection, GPU memory, error formatting. +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) @@ -16,14 +16,15 @@ from unittest.mock import patch, MagicMock import pytest import torch -from utils.utils import ( +from utils.hardware import ( get_device, is_apple_silicon, clear_gpu_cache, get_gpu_memory_info, log_gpu_memory, - format_error_message, + DeviceType, ) +from utils.utils import format_error_message # ========== Helpers ========== @@ -42,31 +43,30 @@ def _actual_device() -> str: class TestGetDevice: """Tests for get_device() — should agree with the real hardware.""" - def test_returns_valid_string(self): + def test_returns_valid_device_type(self): result = get_device() - assert result in ("cuda", "mps", "cpu") + assert result in (DeviceType.CUDA, DeviceType.MPS, DeviceType.CPU) def test_matches_actual_hardware(self): - assert get_device() == _actual_device() + assert get_device().value == _actual_device() # --- Mocked paths to cover all branches regardless of hardware --- def test_returns_cuda_when_cuda_available(self): with patch("torch.cuda.is_available", return_value=True): - assert get_device() == "cuda" + assert get_device() == DeviceType.CUDA def test_returns_mps_when_only_mps_available(self): mock_mps = MagicMock() mock_mps.is_available.return_value = True with patch("torch.cuda.is_available", return_value=False), \ patch.object(torch.backends, "mps", mock_mps, create=True): - assert get_device() == "mps" + assert get_device() == DeviceType.MPS def test_returns_cpu_when_nothing_available(self): - # Patch out mps entirely so hasattr(..., "mps") returns False with patch("torch.cuda.is_available", return_value=False), \ patch("builtins.hasattr", side_effect=lambda obj, name: False if name == "mps" else hasattr(obj, name)): - assert get_device() == "cpu" + assert get_device() == DeviceType.CPU # ========== is_apple_silicon() ========== @@ -77,19 +77,22 @@ class TestIsAppleSilicon: assert isinstance(is_apple_silicon(), bool) def test_true_on_darwin_arm64(self): - with patch("platform.system", return_value="Darwin"), \ - patch("platform.machine", return_value="arm64"): + 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("platform.system", return_value="Linux"), \ - patch("platform.machine", return_value="x86_64"): + 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("platform.system", return_value="Darwin"), \ - patch("platform.machine", return_value="x86_64"): + 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 @@ -103,7 +106,7 @@ class TestClearGpuCache: clear_gpu_cache() def test_calls_cuda_cache_when_cuda_available(self): - with patch("torch.cuda.is_available", return_value=True), \ + 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() @@ -111,21 +114,16 @@ class TestClearGpuCache: mock_ipc.assert_called_once() def test_calls_mps_cache_when_mps_available(self): - mock_mps_backend = MagicMock() - mock_mps_backend.is_available.return_value = True - mock_mps_module = MagicMock() mock_mps_module.empty_cache = MagicMock() - with patch("torch.cuda.is_available", return_value=False), \ - patch.object(torch.backends, "mps", mock_mps_backend, create=True), \ + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS), \ patch.object(torch, "mps", mock_mps_module, create=True): clear_gpu_cache() mock_mps_module.empty_cache.assert_called_once() def test_noop_on_cpu(self): - with patch("torch.cuda.is_available", return_value=False), \ - patch("builtins.hasattr", side_effect=lambda obj, name: False if name == "mps" else hasattr(obj, name)): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): # Must not raise clear_gpu_cache() @@ -150,7 +148,7 @@ class TestGetGpuMemoryInfo: """The reported backend must agree with get_device().""" result = get_gpu_memory_info() device = get_device() - assert result["backend"] == device + assert result["backend"] == device.value # --- When a GPU IS available --- @@ -174,7 +172,7 @@ class TestGetGpuMemoryInfo: mock_props.total_memory = 16 * (1024 ** 3) # 16 GB mock_props.name = "NVIDIA Test GPU" - with patch("utils.utils.get_device", return_value="cuda"), \ + 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)), \ @@ -201,7 +199,7 @@ class TestGetGpuMemoryInfo: mock_mps = MagicMock() mock_mps.current_allocated_memory.return_value = 2 * (1024 ** 3) - with patch("utils.utils.get_device", return_value="mps"), \ + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS), \ patch.dict("sys.modules", {"psutil": mock_psutil}), \ patch.object(torch, "mps", mock_mps, create=True): result = get_gpu_memory_info() @@ -215,7 +213,7 @@ class TestGetGpuMemoryInfo: # --- CPU-only path --- def test_cpu_path_returns_unavailable(self): - with patch("utils.utils.get_device", return_value="cpu"): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): result = get_gpu_memory_info() assert result["available"] is False @@ -224,7 +222,7 @@ class TestGetGpuMemoryInfo: # --- Error resilience --- def test_cuda_error_returns_unavailable(self): - with patch("utils.utils.get_device", return_value="cuda"), \ + 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() @@ -249,9 +247,9 @@ class TestLogGpuMemory: "utilization_pct": 12.5, "free_gb": 14.0, } - with patch("utils.utils.get_gpu_memory_info", return_value=fake_info): + with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info): import logging - with caplog.at_level(logging.INFO, logger="utils.utils"): + with caplog.at_level(logging.INFO, logger="utils.hardware.hardware"): log_gpu_memory("unit-test") assert "unit-test" in caplog.text @@ -260,9 +258,9 @@ class TestLogGpuMemory: def test_logs_cpu_fallback_when_no_gpu(self, caplog): fake_info = {"available": False, "backend": "cpu"} - with patch("utils.utils.get_gpu_memory_info", return_value=fake_info): + with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info): import logging - with caplog.at_level(logging.INFO, logger="utils.utils"): + with caplog.at_level(logging.INFO, logger="utils.hardware.hardware"): log_gpu_memory("cpu-test") assert "No GPU available" in caplog.text @@ -297,7 +295,7 @@ class TestFormatErrorMessage: def test_cuda_oom(self): err = Exception("CUDA out of memory") - with patch("utils.utils.get_device", return_value="cuda"): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA): msg = format_error_message(err, "big/model") assert "GPU" in msg assert "big/model" not in msg # should use short name @@ -307,7 +305,7 @@ class TestFormatErrorMessage: def test_mps_oom(self): err = Exception("MPS backend out of memory") - with patch("utils.utils.get_device", return_value="mps"): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS): msg = format_error_message(err, "unsloth/huge-model") assert "Apple Silicon" in msg @@ -315,7 +313,7 @@ class TestFormatErrorMessage: def test_cpu_oom(self): err = Exception("not enough memory to allocate") - with patch("utils.utils.get_device", return_value="cpu"): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): msg = format_error_message(err, "any/model") assert "system" in msg.lower() diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py new file mode 100644 index 0000000000..b992fc191d --- /dev/null +++ b/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', +] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py new file mode 100644 index 0000000000..ea223096b1 --- /dev/null +++ b/studio/backend/utils/hardware/hardware.py @@ -0,0 +1,195 @@ +""" +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" + MPS = "mps" + 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" +pass + + +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. MPS (Apple Silicon via PyTorch MPS backend) + 3. CPU (fallback) + """ + global DEVICE + + # --- Try PyTorch first (covers CUDA and MPS) --- + try: + 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 + + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + DEVICE = DeviceType.MPS + chip = platform.processor() or platform.machine() + logger.info(f"Hardware detected: MPS — Apple Silicon ({chip})") + return DEVICE + + except ImportError: + logger.warning("PyTorch not installed — falling back to CPU") + + DEVICE = DeviceType.CPU + logger.info("Hardware detected: CPU (no GPU backend available)") + return DEVICE +pass + + +# ========== 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 +pass + + +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.MPS: + import torch + if hasattr(torch.mps, "empty_cache"): + torch.mps.empty_cache() +pass + + +def get_gpu_memory_info() -> Dict[str, Any]: + """ + Get GPU memory information. + Supports CUDA (NVIDIA), MPS (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)} + + # ---- MPS path (Apple Silicon) ---- + if device == DeviceType.MPS: + try: + import torch + import psutil + allocated = torch.mps.current_allocated_memory() if hasattr(torch.mps, "current_allocated_memory") else 0 + total = psutil.virtual_memory().total + + 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, # MPS doesn't have a separate reserved pool + "free_gb": (total - allocated) / (1024**3), + "utilization_pct": (allocated / total) * 100 if total else 0, + } + except Exception as e: + logger.error(f"Error getting MPS GPU info: {e}") + return {"available": False, "backend": device.value, "error": str(e)} + + # ---- CPU-only ---- + return {"available": False, "backend": "cpu"} +pass + + +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)") +pass diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 7bae2dd086..ee88406965 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -1,61 +1,16 @@ """ Shared backend utilities """ -import gradio as gr import os -import platform import logging from contextlib import contextmanager from pathlib import Path -from typing import Optional, Dict, Any import shutil import tempfile logger = logging.getLogger(__name__) -# ========== Device Detection & Management ========== - -def get_device() -> str: - """ - Detect the best available compute device. - - Returns: - "cuda" on NVIDIA GPUs, "mps" on Apple Silicon, "cpu" otherwise. - """ - import torch - - if torch.cuda.is_available(): - return "cuda" - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return "mps" - return "cpu" -pass - - -def is_apple_silicon() -> bool: - """Check if running on Apple Silicon hardware.""" - return platform.system() == "Darwin" and platform.machine() == "arm64" -pass - - -def clear_gpu_cache(): - """ - Clear GPU memory cache for the current device. - Safe to call on any platform — no-ops gracefully when the backend is unavailable. - """ - import torch - import gc - - gc.collect() - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - if hasattr(torch.mps, "empty_cache"): - torch.mps.empty_cache() -pass @contextmanager def without_hf_auth(): @@ -141,149 +96,11 @@ def format_error_message(error: Exception, model_name: str) -> str: return "Invalid HF token. Please check your token and try again." if "memory" in error_str or "cuda" in error_str or "mps" in error_str or "out of memory" in error_str: + from utils.hardware import get_device device = get_device() - device_label = {"cuda": "GPU", "mps": "Apple Silicon GPU", "cpu": "system"}.get(device, "GPU") + device_label = {"cuda": "GPU", "mps": "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. - Supports CUDA (NVIDIA), MPS (Apple Silicon), and CPU-only environments. - """ - import torch - - device_type = get_device() - - # ---- CUDA path ---- - if device_type == "cuda": - 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, - "backend": "cuda", - "device": device, - "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": "cuda", "error": str(e)} - - # ---- MPS path (Apple Silicon) ---- - if device_type == "mps": - try: - allocated = torch.mps.current_allocated_memory() if hasattr(torch.mps, "current_allocated_memory") else 0 - # MPS doesn't expose total VRAM directly — use unified memory from psutil as a proxy - import psutil - total = psutil.virtual_memory().total - - return { - "available": True, - "backend": "mps", - "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, # MPS doesn't have a separate reserved pool - "free_gb": (total - allocated) / (1024**3), - "utilization_pct": (allocated / total) * 100 if total else 0, - } - except Exception as e: - logger.error(f"Error getting MPS GPU info: {e}") - return {"available": False, "backend": "mps", "error": str(e)} - - # ---- CPU-only ---- - return {"available": False, "backend": "cpu"} -pass - -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)") -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=[])