diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 66aa9a2db6..3d911f62f1 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -280,7 +280,9 @@ class ExportBackend: self._audio_type = detect_audio_type( model_id, hf_token = token, local_files_only = local_files_only ) - self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token) + self.is_vision = not self._audio_type and is_vision_model( + model_id, hf_token = token, local_files_only = local_files_only + ) if self._audio_type == "csm": from unsloth import FastModel diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py index 97d07e6874..cf754ec06e 100644 --- a/studio/backend/tests/test_vision_cache.py +++ b/studio/backend/tests/test_vision_cache.py @@ -71,7 +71,7 @@ class TestVisionCacheHitMiss: """Two calls for the same model invoke the uncached fn 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) + mock_uncached.assert_called_once_with("org/my-vlm", None, local_files_only = False) @patch("utils.models.model_config._is_vision_model_uncached", return_value = False) def test_different_models_each_detected(self, mock_uncached): @@ -97,7 +97,7 @@ class TestVisionCacheStoresFalse: 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 + assert _vision_detection_cache[("org/text-only", None, False)] is False # Subprocess path (transformers 5.x) caching @@ -120,7 +120,7 @@ class TestVisionCacheSubprocessPath: 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 + assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None, False)] is True @patch("utils.models.model_config._raw_config_has_vision_config", return_value = True) @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None) @@ -133,7 +133,9 @@ class TestVisionCacheSubprocessPath: assert is_vision_model("unsloth/gemma-4-E4B-it") is True assert is_vision_model("unsloth/gemma-4-E4B-it") is True - mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None) + mock_raw_config.assert_called_once_with( + "unsloth/gemma-4-E4B-it", hf_token = None, local_files_only = False + ) mock_subprocess.assert_not_called() @@ -405,6 +407,40 @@ class TestVisionCacheTokenHandling: mock_uncached.assert_called_once() +class TestVisionCacheLocalOnly: + """local_files_only is part of the cache key (mirrors the audio cache). An + offline probe only sees the on-disk cache, so its negative result must never + be reused by a later online probe that can reach the Hub -- otherwise a VLM + is routed through the text loader until restart.""" + + def test_local_only_negative_does_not_poison_online(self, monkeypatch): + import utils.models.model_config as mc + + mc._vision_detection_cache.clear() + monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n) + + seen = [] + + def _probe(name, hf_token = None, local_files_only = False): + seen.append(local_files_only) + # Offline: cannot run the transformers-5 subprocess / fetch -> not a VLM. + # Online: the remote config reveals a VLM. + return False if local_files_only else True + + monkeypatch.setattr(mc, "_is_vision_model_uncached", _probe) + + # Offline probe caches False under a local-only key. + assert mc.is_vision_model("some/vlm", local_files_only = True) is False + # A later online probe must re-run (different key) and detect the VLM. + assert mc.is_vision_model("some/vlm", local_files_only = False) is True + assert seen == [True, False] + # The online positive is then cached for subsequent online callers. + assert mc.is_vision_model("some/vlm", local_files_only = False) is True + assert seen == [True, False] + mc._vision_detection_cache.clear() + + # --------------------------------------------------------------------------- # Direct unit tests for _raw_config_has_vision_config # --------------------------------------------------------------------------- diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index f359b74dc5..7530a4b16d 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -471,6 +471,7 @@ def load_model_config( use_auth: bool = False, token: Optional[str] = None, trust_remote_code: bool = False, + local_files_only: bool = False, ): """Load model config with optional authentication control. @@ -478,12 +479,18 @@ def load_model_config( metadata lookups must never execute a model repo's ``auto_map`` Python. Deliberate remote-code loads pass the flag explicitly through ``FastLanguageModel.from_pretrained`` with the user's own consent. + + ``local_files_only`` keeps the config read on the local HF cache (offline + export), so an offline probe never blocks on the network. """ from transformers import AutoConfig if token: return AutoConfig.from_pretrained( - model_name, trust_remote_code = trust_remote_code, token = token + model_name, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, ) if not use_auth: @@ -493,12 +500,14 @@ def load_model_config( model_name, trust_remote_code = trust_remote_code, token = None, + local_files_only = local_files_only, ) # Default auth (cached tokens) return AutoConfig.from_pretrained( model_name, trust_remote_code = trust_remote_code, + local_files_only = local_files_only, ) @@ -598,7 +607,7 @@ def _is_vlm(config) -> bool: def _raw_config_has_vision_config( - model_name: str, hf_token: Optional[str] = None + model_name: str, hf_token: Optional[str] = None, local_files_only: bool = False ) -> Optional[bool]: try: if is_local_path(model_name): @@ -610,6 +619,7 @@ def _raw_config_has_vision_config( repo_id = model_name, filename = "config.json", token = hf_token, + local_files_only = local_files_only, ) ) config = json.loads(config_path.read_text()) @@ -780,22 +790,29 @@ def _token_fingerprint(token: Optional[str]) -> Optional[str]: # Keyed by (normalized_model_name, token_fingerprint) to handle gated models. # Only definitive results are cached; transient failures (network, timeouts) # are NOT cached so they can be retried. -_vision_detection_cache: Dict[Tuple[str, Optional[str]], bool] = {} +_vision_detection_cache: Dict[Tuple[str, Optional[str], bool], bool] = {} _vision_cache_lock = threading.Lock() -def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: +def is_vision_model( + model_name: str, hf_token: Optional[str] = None, local_files_only: bool = False +) -> bool: """ Detect vision-language models (VLMs) via architecture in config. Works for fine-tuned models since they inherit the base architecture. Models needing transformers 5.x are checked in a .venv_t5/ subprocess. - Results are cached per (model_name, token_fingerprint) for the process - lifetime; transient failures are not cached so they can be retried. + Results are cached per (model_name, token_fingerprint, local_files_only) for + the process lifetime; transient failures are not cached so they can be + retried. local_files_only is part of the key because an offline probe reads + only the on-disk cache and can differ from an online probe (e.g. a + transformers-5 VLM the offline path cannot run the subprocess for), so the + two results must never share a cache entry. Args: model_name: Model identifier (HF repo or local path) hf_token: Optional HF token for gated/private models + local_files_only: Keep detection on the local HF cache (offline export) """ # Local GGUF models are served by llama-server. Their multimodal # capability comes from a companion mmproj, not a Transformers config. @@ -829,7 +846,7 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: exc, ) resolved_name = model_name - cache_key = (resolved_name, _token_fingerprint(hf_token)) + cache_key = (resolved_name, _token_fingerprint(hf_token), bool(local_files_only)) # Lock-free fast path for cache hits. Sentinel distinguishes "key not found" # from "value is False" in a single atomic dict.get() call. @@ -840,7 +857,9 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: # Compute outside the lock so long-running detection isn't serialized across # models. Two concurrent calls may both run, but produce the same result. - result = _is_vision_model_uncached(resolved_name, hf_token) + result = _is_vision_model_uncached( + resolved_name, hf_token, local_files_only = local_files_only + ) # Only cache definitive results; None is a transient failure, retry later. if result is not None: with _vision_cache_lock: @@ -849,7 +868,9 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: return False -def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]: +def _is_vision_model_uncached( + model_name: str, hf_token: Optional[str] = None, local_files_only: bool = False +) -> Optional[bool]: """Uncached vision detection; use is_vision_model() instead. Returns True/False for definitive results, or None on transient errors @@ -858,15 +879,19 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - # Try the raw-config reader FIRST (code-free, version-independent): it classifies # repo-code VLMs like DeepSeek-OCR via declarative vision_config with no remote-code # execution or transformers-5.x subprocess. - raw = _raw_config_has_vision_config(model_name, hf_token = hf_token) + raw = _raw_config_has_vision_config( + model_name, hf_token = hf_token, local_files_only = local_files_only + ) if raw is not None: return raw # Raw read failed transiently: fall back to AutoConfig with remote code DISABLED # (in a transformers-5.x subprocess when the main process can't parse the arch). + # Skip the subprocess offline: it does its own network probe and would diverge + # from the online path, so offline we stay on the local cache via load_model_config. from utils.transformers_version import needs_transformers_5 - if needs_transformers_5(model_name): + if not local_files_only and needs_transformers_5(model_name): logger.info( "Model '%s' needs transformers 5.x -- checking vision via subprocess", model_name, @@ -874,7 +899,12 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - return _is_vision_model_subprocess(model_name, hf_token = hf_token) try: - config = load_model_config(model_name, use_auth = True, token = hf_token) + config = load_model_config( + model_name, + use_auth = True, + token = hf_token, + local_files_only = local_files_only, + ) if _is_vlm(config): model_type = getattr(config, "model_type", None) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index f19e7ea616..d8b93c4381 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -860,6 +860,7 @@ class FastLanguageModel(FastLlamaModel): old_model_name, token = token, revision = revision, + local_files_only = local_files_only, is_trainable = True, trust_remote_code = trust_remote_code, ) @@ -1767,6 +1768,7 @@ class FastModel(FastBaseModel): old_model_name, token = token, revision = revision, + local_files_only = local_files_only, is_trainable = True, trust_remote_code = trust_remote_code, ) diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 35b867f2dc..4c4a23ff80 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -737,12 +737,25 @@ def _offline_aware_load(fn): try: return fn(*args, **kwargs) except Exception as e: - if _is_offline_related_error(e): - kwargs["local_files_only"] = True - with _force_hf_offline(): - return fn(*args, **kwargs) - raise - + if not _is_offline_related_error(e): + raise + # Retry OUTSIDE the except block: an except-scoped exception still holds its + # __traceback__, which pins the failed attempt's frame locals (e.g. a + # partially loaded model) until the block exits. Loading the model a second + # time while the first copy is still alive can OOM a large VLM. Letting the + # except block close drops the traceback; collect + empty the device cache + # so the partial load is freed before the forced-offline retry reallocates. + try: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() + except Exception: + pass + kwargs["local_files_only"] = True + with _force_hf_offline(): + return fn(*args, **kwargs) return _wrapper