+
+ If the link doesn't work, you can scroll down to view the UI generated directly in Colab. +
+{short_url}
diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index ab1825d94a..77cbda6b45 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -190,7 +190,11 @@ class UnslothTrainer: self._cuda_audio_used = False # --- Detect VLM --- - vision = is_vision_model(model_name) if not self.is_audio else False + vision = ( + is_vision_model(model_name, hf_token = hf_token) + if not self.is_audio + else False + ) self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image logger.info( @@ -558,7 +562,11 @@ class UnslothTrainer: self._cuda_audio_used = False # VLM: vision model with image dataset (mutually exclusive with audio paths) - vision = is_vision_model(model_name) if not self.is_audio else False + vision = ( + is_vision_model(model_name, hf_token = hf_token) + if not self.is_audio + else False + ) self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image self.model_name = model_name self.max_seq_length = max_seq_length diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1e31a91e26..3f361ca5eb 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -138,6 +138,47 @@ def _resolve_hf_cache_dir() -> Path: return Path.home() / ".cache" / "huggingface" / "hub" +def _is_model_directory(d: Path) -> bool: + """Return ``True`` when *d* looks like a model directory. + + A model directory must have **both** a config file (``config.json`` or + ``adapter_config.json``) **and** actual model weight files. Both + conditions are required: a bare directory with only loose ``.gguf`` + files (no config) might be a mixed collection, and a ``config.json`` + alone (no weights) is not a model directory. + + Excludes ``mmproj`` GGUF files (vision projectors) and non-weight + ``.bin`` files (``tokenizer.bin``, ``vocab.bin``, etc.) from the + weight check to avoid false positives. + """ + + def _is_weight_file(f: Path) -> bool: + suffix = f.suffix.lower() + if suffix == ".safetensors": + return True + if suffix == ".gguf": + return "mmproj" not in f.name.lower() + if suffix == ".bin": + name = f.name.lower() + return ( + name.startswith("pytorch_model") + or name.startswith("model") + or name.startswith("adapter_model") + or name.startswith("consolidated") + ) + return False + + try: + has_config = (d / "config.json").exists() or ( + d / "adapter_config.json" + ).exists() + if not has_config: + return False + return any(_is_weight_file(f) for f in d.iterdir() if f.is_file()) + except OSError: + return False + + def _scan_models_dir( models_dir: Path, *, @@ -146,6 +187,23 @@ def _scan_models_dir( if not models_dir.exists() or not models_dir.is_dir(): return [] + _is_self_model = _is_model_directory(models_dir) + + if _is_self_model: + try: + updated_at = models_dir.stat().st_mtime + except OSError: + updated_at = None + return [ + LocalModelInfo( + id = str(models_dir), + display_name = models_dir.name, + path = str(models_dir), + source = "models_dir", + updated_at = updated_at, + ), + ] + found: List[LocalModelInfo] = [] for child in models_dir.iterdir(): if limit is not None and len(found) >= limit: @@ -243,6 +301,25 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: if not lm_dir.exists() or not lm_dir.is_dir(): return [] + # If the directory itself is a model directory (has config AND weight + # files), it is not an LM Studio publisher structure -- return it as a + # single model entry. We cannot skip it silently because this function + # is the only scanner called for default LM Studio roots. + if _is_model_directory(lm_dir): + try: + updated_at = lm_dir.stat().st_mtime + except OSError: + updated_at = None + return [ + LocalModelInfo( + id = str(lm_dir), + display_name = lm_dir.name, + path = str(lm_dir), + source = "lmstudio", + updated_at = updated_at, + ), + ] + found: List[LocalModelInfo] = [] for child in lm_dir.iterdir(): try: @@ -263,6 +340,25 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: ) continue + # If the child directory itself looks like a model directory + # (has config AND weight files), surface it directly instead + # of descending into it as a publisher. + if _is_model_directory(child): + try: + updated_at = child.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id = str(child), + display_name = child.name, + path = str(child), + source = "lmstudio", + updated_at = updated_at, + ), + ) + continue + # child is a publisher directory -- scan its sub-directories for model_dir in child.iterdir(): try: @@ -618,7 +714,7 @@ async def get_model_config( config_dict = load_model_defaults(model_name) # Detect model capabilities (pass HF token for gated models) - is_vision = is_vision_model(model_name) + is_vision = is_vision_model(model_name, hf_token = hf_token) is_embedding = is_embedding_model(model_name, hf_token = hf_token) audio_type = detect_audio_type(model_name, hf_token = hf_token) diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py new file mode 100644 index 0000000000..fae1e95311 --- /dev/null +++ b/studio/backend/tests/test_vision_cache.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for is_vision_model() caching behaviour. + +The vision detection cache (``_vision_detection_cache``) mirrors the existing +``_audio_detection_cache`` pattern used by ``detect_audio_type()``. These +tests verify that: + +* Repeated calls for the same model hit the cache (no redundant work). +* Different models each trigger their own detection. +* Both True and False results are cached. +* The subprocess path (transformers 5.x models) is also cached. +* Exceptions that fall back to False are cached. +""" + +import sys +import types as _types +from pathlib import Path +from unittest.mock import patch, MagicMock + +import pytest + +# --------------------------------------------------------------------------- +# sys.path + logger stub — same pattern as the rest of the test suite +# --------------------------------------------------------------------------- +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from utils.models.model_config import ( + is_vision_model, + _is_vision_model_uncached, + _vision_detection_cache, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse = True) +def _clear_vision_cache(): + """Ensure every test starts with a fresh cache.""" + _vision_detection_cache.clear() + yield + _vision_detection_cache.clear() + + +# --------------------------------------------------------------------------- +# Cache hit / miss tests +# --------------------------------------------------------------------------- + + +class TestVisionCacheHitMiss: + """Verify the cache prevents redundant detection calls.""" + + @patch("utils.models.model_config._is_vision_model_uncached", return_value = True) + def test_second_call_uses_cache(self, mock_uncached): + """Calling is_vision_model() twice for the same model should invoke + the uncached function only once.""" + assert is_vision_model("org/my-vlm") is True + assert is_vision_model("org/my-vlm") is True + mock_uncached.assert_called_once_with("org/my-vlm", None) + + @patch("utils.models.model_config._is_vision_model_uncached", return_value = False) + def test_different_models_each_detected(self, mock_uncached): + """Different model names should each trigger detection.""" + is_vision_model("model-a") + is_vision_model("model-b") + assert mock_uncached.call_count == 2 + + @patch("utils.models.model_config._is_vision_model_uncached", return_value = True) + def test_cache_returns_correct_value(self, mock_uncached): + """The cached value must match what _is_vision_model_uncached returned.""" + first = is_vision_model("org/vlm") + second = is_vision_model("org/vlm") + assert first is True + assert second is True + + +class TestVisionCacheStoresFalse: + """Non-VLM results (False) must also be cached to avoid re-detection.""" + + @patch("utils.models.model_config._is_vision_model_uncached", return_value = False) + def test_false_result_cached(self, mock_uncached): + assert is_vision_model("org/text-only") is False + assert is_vision_model("org/text-only") is False + mock_uncached.assert_called_once() + assert _vision_detection_cache[("org/text-only", None)] is False + + +# --------------------------------------------------------------------------- +# Subprocess path (transformers 5.x) caching +# --------------------------------------------------------------------------- + + +class TestVisionCacheSubprocessPath: + """Models needing transformers 5.x go through _is_vision_model_subprocess. + The cache should prevent the subprocess from being spawned more than once + per model per process.""" + + @patch("utils.models.model_config._is_vision_model_subprocess", return_value = True) + @patch("utils.transformers_version.needs_transformers_5", return_value = True) + def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess): + """Subprocess should only fire on the first call; second is cached.""" + # First call: goes through uncached → subprocess + assert is_vision_model("unsloth/Qwen3.5-2B") is True + # Second call: cache hit, no subprocess + assert is_vision_model("unsloth/Qwen3.5-2B") is True + + mock_subprocess.assert_called_once() + assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None)] is True + + +# --------------------------------------------------------------------------- +# Exception handling — cache the False fallback +# --------------------------------------------------------------------------- + + +class TestVisionCacheOnException: + """When detection raises an exception, _is_vision_model_uncached catches + it and returns False. That False must be cached so subsequent calls don't + retry and fail again.""" + + @patch( + "utils.models.model_config.load_model_config", + side_effect = OSError("network down"), + ) + @patch("utils.transformers_version.needs_transformers_5", return_value = False) + def test_exception_result_cached(self, mock_needs_t5, mock_load_config): + """A real exception inside _is_vision_model_uncached should be caught, + return False, and that False should be cached for subsequent calls.""" + # First call: load_model_config raises → except branch → False + assert is_vision_model("broken/model") is False + # Second call: cache hit, load_model_config not called again + assert is_vision_model("broken/model") is False + mock_load_config.assert_called_once() + + +# --------------------------------------------------------------------------- +# Direct detection path (non-transformers-5 models) caching +# --------------------------------------------------------------------------- + + +class TestVisionCacheDirectPath: + """For models that do NOT need transformers 5.x, the detection goes through + load_model_config directly. The cache must work the same way.""" + + @patch("utils.transformers_version.needs_transformers_5", return_value = False) + @patch("utils.models.model_config.load_model_config") + def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5): + """A standard VLM detected via architecture suffix should be cached.""" + cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist + cfg.model_type = "gemma3" + cfg.architectures = ["Gemma3ForConditionalGeneration"] + mock_load_config.return_value = cfg + + assert is_vision_model("google/gemma-3-4b-it") is True + assert is_vision_model("google/gemma-3-4b-it") is True + # load_model_config should only be called once + mock_load_config.assert_called_once() + + @patch("utils.transformers_version.needs_transformers_5", return_value = False) + @patch("utils.models.model_config.load_model_config") + def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5): + """A standard text model (no VLM indicators) should cache False.""" + cfg = MagicMock(spec = []) # spec=[] means no attributes at all + cfg.model_type = "llama" + cfg.architectures = ["LlamaForCausalLM"] + mock_load_config.return_value = cfg + + # LlamaForCausalLM doesn't end with VLM suffixes, no vision_config, etc. + assert is_vision_model("meta-llama/Llama-3-8B") is False + assert is_vision_model("meta-llama/Llama-3-8B") is False + mock_load_config.assert_called_once() + + @patch("utils.transformers_version.needs_transformers_5", return_value = False) + @patch("utils.models.model_config.load_model_config") + def test_vision_config_attr_detected_and_cached( + self, mock_load_config, mock_needs_t5 + ): + """Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True.""" + cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist + cfg.model_type = "qwen2_vl" + cfg.architectures = ["Qwen2VLForCausalLM"] # Doesn't match VLM suffixes + cfg.vision_config = {"hidden_size": 1024} + mock_load_config.return_value = cfg + + assert is_vision_model("Qwen/Qwen2-VL-7B") is True + assert is_vision_model("Qwen/Qwen2-VL-7B") is True + mock_load_config.assert_called_once() + + @patch("utils.transformers_version.needs_transformers_5", return_value = False) + @patch("utils.models.model_config.load_model_config") + def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5): + """Audio-only models (csm, whisper) with ForConditionalGeneration + should be excluded from VLM detection and cached as False.""" + cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist + cfg.model_type = "whisper" + cfg.architectures = ["WhisperForConditionalGeneration"] + mock_load_config.return_value = cfg + + assert is_vision_model("openai/whisper-large-v3") is False + assert is_vision_model("openai/whisper-large-v3") is False + mock_load_config.assert_called_once() + + +# --------------------------------------------------------------------------- +# hf_token handling +# --------------------------------------------------------------------------- + + +class TestVisionCacheTokenHandling: + """The cache is keyed on (model_name, hf_token). + Different tokens for the same model should trigger separate detections + to handle gated models correctly.""" + + @patch("utils.models.model_config._is_vision_model_uncached", return_value = True) + def test_different_tokens_trigger_new_detection(self, mock_uncached): + """Calls with different tokens should trigger separate detections to + handle gated models correctly (e.g. unauthenticated probe → False, + then authenticated call should re-check).""" + assert is_vision_model("gated/model", hf_token = "token-a") is True + assert is_vision_model("gated/model", hf_token = "token-b") is True + assert mock_uncached.call_count == 2 + + @patch("utils.models.model_config._is_vision_model_uncached", return_value = True) + def test_same_token_uses_cache(self, mock_uncached): + """Repeated calls with identical model + token should hit cache.""" + assert is_vision_model("gated/model", hf_token = "token-a") is True + assert is_vision_model("gated/model", hf_token = "token-a") is True + mock_uncached.assert_called_once() diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index f97ea993eb..61226e52cb 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -26,7 +26,9 @@ import subprocess import sys from pathlib import Path from typing import List, Tuple +import hashlib import json +import threading import yaml @@ -548,12 +550,17 @@ except Exception as exc: def _is_vision_model_subprocess( model_name: str, hf_token: Optional[str] = None -) -> bool: +) -> Optional[bool]: """Run is_vision_model check in a subprocess with transformers 5.x. Same pattern as training/inference workers: spawn a clean subprocess with .venv_t5/ prepended to sys.path so AutoConfig recognizes newer architectures (glm4_moe_lite, etc.). + + Returns True/False for definitive results, or None for transient failures + (timeouts, subprocess errors) so callers can decide whether to cache + the result. Subprocess failures are treated as transient because they + can be caused by temporary HF/auth/network issues. """ token_arg = hf_token or "" @@ -580,7 +587,7 @@ def _is_vision_model_subprocess( model_name, stderr or result.stdout.strip(), ) - return False + return None data = json.loads(result.stdout.strip()) if "error" in data: @@ -589,7 +596,7 @@ def _is_vision_model_subprocess( model_name, data["error"], ) - return False + return None is_vlm = data["is_vision"] logger.info( @@ -604,10 +611,28 @@ def _is_vision_model_subprocess( except subprocess.TimeoutExpired: logger.warning("Vision check subprocess timed out for '%s'", model_name) - return False + return None except Exception as exc: logger.warning("Vision check subprocess failed for '%s': %s", model_name, exc) - return False + return None + + +def _token_fingerprint(token: Optional[str]) -> Optional[str]: + """Return a SHA256 digest of the token for use as a cache key. + + Avoids storing the raw bearer token in process memory as a dict key. + """ + if token is None: + return None + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +# Cache vision detection results per session to avoid repeated subprocess spawns. +# Keyed by (normalized_model_name, token_fingerprint) to handle gated models correctly. +# Only definitive results (True/False from successful detection) are cached; +# transient failures (network errors, timeouts) are NOT cached so they can be retried. +_vision_detection_cache: Dict[Tuple[str, Optional[str]], bool] = {} +_vision_cache_lock = threading.Lock() def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: @@ -616,13 +641,66 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: Works for fine-tuned models since they inherit the base architecture. For models that require transformers 5.x (e.g. GLM-4.7-Flash), the check - runs in a subprocess with .venv_t5/ activated — same pattern as the + runs in a subprocess with .venv_t5/ activated -- same pattern as the training and inference workers. + Results are cached per (model_name, token_fingerprint) for the lifetime of + the process to avoid repeated subprocess spawns and HuggingFace API calls. + Transient failures are not cached so they can be retried on the next call. + Args: model_name: Model identifier (HF repo or local path) hf_token: Optional HF token for accessing gated/private models """ + # Normalize model name for cache key to avoid duplicate entries for + # different casings of the same HF repo (e.g. "Org/Model" vs "org/model"). + try: + if is_local_path(model_name): + resolved_name = normalize_path(model_name) + else: + resolved_name = resolve_cached_repo_id_case(model_name) + except Exception as exc: + logger.debug( + "Could not normalize model name '%s' for cache key: %s", + model_name, + exc, + ) + resolved_name = model_name + cache_key = (resolved_name, _token_fingerprint(hf_token)) + + # Lock-free fast path for cache hits. Uses a sentinel to distinguish + # "key not found" from "value is False" in a single atomic dict.get() call. + _MISS = object() + cached = _vision_detection_cache.get(cache_key, _MISS) + if cached is not _MISS: + return cached + + # Compute outside the lock to avoid serializing long-running detection + # (subprocess spawns with 60s timeout, HF API calls) across all models. + # The tradeoff: two concurrent calls for the same uncached model may + # both run detection, but they produce the same result and the second + # write is a benign no-op. + result = _is_vision_model_uncached(resolved_name, hf_token) + # Only cache definitive results; None means a transient failure occurred + # and we should retry on the next call instead of locking in a wrong answer. + if result is not None: + with _vision_cache_lock: + _vision_detection_cache[cache_key] = result + return result + return False + + +def _is_vision_model_uncached( + model_name: str, hf_token: Optional[str] = None +) -> Optional[bool]: + """Uncached vision model detection -- called by is_vision_model(). + + Returns True/False for definitive results, or None when detection failed + due to a transient error (network, timeout, subprocess failure) so the + caller knows not to cache the result. + + Do not call directly; use is_vision_model() instead. + """ # Models that need transformers 5.x must be checked in a subprocess # because AutoConfig in the main process (transformers 4.57.x) doesn't # recognize their architectures. @@ -630,7 +708,7 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: if needs_transformers_5(model_name): logger.info( - "Model '%s' needs transformers 5.x — checking vision via subprocess", + "Model '%s' needs transformers 5.x -- checking vision via subprocess", model_name, ) return _is_vision_model_subprocess(model_name, hf_token = hf_token) @@ -681,7 +759,25 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: except Exception as e: logger.warning(f"Could not determine if {model_name} is vision model: {e}") - return False + # Permanent failures (model not found, gated, bad config) should be + # cached as False. Transient failures (network, timeout) should not. + try: + from huggingface_hub.errors import RepositoryNotFoundError, GatedRepoError + except ImportError: + try: + from huggingface_hub.utils import ( + RepositoryNotFoundError, + GatedRepoError, + ) + except ImportError: + RepositoryNotFoundError = GatedRepoError = None + if RepositoryNotFoundError is not None and isinstance( + e, (RepositoryNotFoundError, GatedRepoError) + ): + return False + if isinstance(e, (ValueError, json.JSONDecodeError)): + return False + return None VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm") @@ -1034,6 +1130,26 @@ def list_gguf_variants( return variants, has_vision +def _resolve_gguf_dir(p: Path) -> Optional[Path]: + """Resolve a path to the directory containing GGUF variants. + + If *p* is already a directory, returns it directly. If *p* is a ``.gguf`` + file whose parent directory has model metadata (``config.json`` or + ``adapter_config.json``), returns the parent -- all GGUFs in that + directory belong to the same model. Returns ``None`` for loose standalone + GGUFs (no config) to avoid cross-wiring unrelated models. + """ + if p.is_dir(): + return p + if p.is_file() and p.suffix.lower() == ".gguf": + parent = p.parent + if (parent / "config.json").exists() or ( + parent / "adapter_config.json" + ).exists(): + return parent + return None + + def list_local_gguf_variants( directory: str, ) -> tuple[list[GgufVariantInfo], bool]: @@ -1046,8 +1162,8 @@ def list_local_gguf_variants( Returns: (variants, has_vision): list of non-mmproj GGUF variants + vision flag. """ - p = Path(directory) - if not p.is_dir(): + p = _resolve_gguf_dir(Path(directory)) + if p is None: return [], False quant_totals: dict[str, int] = {} @@ -1087,8 +1203,8 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: Returns the resolved absolute path, or ``None`` if no match. """ - p = Path(directory) - if not p.is_dir(): + p = _resolve_gguf_dir(Path(directory)) + if p is None: return None matches = sorted( diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 1dbff145ee..cf1ba11d7b 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -225,6 +225,7 @@ const LoraCompareContent = memo(function LoraCompareContent({ modelType="base" pairId={pairId} initialThreadId={baseThreadId} + syncActiveThreadId={false} >