unsloth/studio/backend/utils/models/model_config.py
Daniel Han 1396c01253
Fix offline checkpoint load/export: "tokenizer is weirdly not loaded" (#6554)
* Fix offline checkpoint load/export failing with "tokenizer is weirdly not loaded"

Loading a fine-tuned checkpoint with no internet (e.g. a Studio export) crashed
with "Unsloth: The tokenizer is weirdly not loaded? Please check if there is one."

For a LoRA adapter the loader reassigns model_name to the base model repo id and
only keeps the local checkpoint dir as tokenizer_name when it contains
tokenizer_config.json, tokenizer.json AND special_tokens_map.json. Modern
tokenizers (e.g. Gemma) store special tokens inside tokenizer_config.json and
omit special_tokens_map.json, so tokenizer_name fell back to the base repo id.
The tokenizer/processor loads in vision.py then hit the Hub with no
local_files_only, so with no network they failed (AutoProcessor) or hung for
minutes (AutoTokenizer) even though every file was already cached.

loader.py: keep the local checkpoint dir as tokenizer_name when it has a
tokenizer config plus the actual tokenizer files (tokenizer.json / tokenizer.model
/ vocab files); special_tokens_map.json is no longer required.

vision.py: compute an effective local_files_only (explicit kwarg plus the
HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE env vars, mirroring loader.py and
diffusion.py) and thread it through every AutoConfig, AutoProcessor,
AutoTokenizer and the manual VLM processor fallback, including the
hf_hub_download in that fallback (which now prefers a local file). When a load
fails and no offline env var is set, retry against the local cache. The retry
forces HF offline mode because local_files_only alone does not stop
AutoProcessor / AutoTokenizer from issuing a /api/models request during class
resolution. The final error now explains the offline/cache cause instead of the
misleading "weirdly not loaded" message.

studio export: probe Hub reachability once per checkpoint load and pass
local_files_only when offline so exports use the local checkpoint dir / cache
instead of hanging or crashing with no internet.

Online behavior is unchanged: the new flags default to off and the retry only
runs after a network related failure.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: safer offline forcing, cached fallback config, proxy-aware probe

Follow-up to the offline checkpoint load fix, addressing review feedback:

- vision.py: only flip the process-wide HF offline flag when offline is actually
  requested (local_files_only / env) or after a real network failure, never
  pre-emptively while we might be online. The flip is now guarded by a lock +
  depth counter so nested or concurrent windows restore the flag correctly
  (no stale value).
- vision.py: guard the get_auto_processor fallback so a network error there
  returns None and the local-cache retry still runs instead of escaping.
- vision.py: in the manual VLM processor fallback, read tokenizer_config.json
  via hf_hub_download(..., local_files_only=...) so a cached repo-id config is
  still resolved offline and the model-specific image/video tokens are restored.
- studio export: make the reachability probe proxy aware (probe the configured
  HTTP(S) proxy egress, honour NO_PROXY, use the endpoint port) so a proxy-only
  setup is not wrongly marked offline; allow UNSLOTH_OFFLINE_PROBE=0 to disable.
- studio export: run the audio/vision type-detection probes inside the
  forced-offline window when offline, so their config/tokenizer reads hit the
  local cache instead of waiting out connection timeouts.

Online behavior remains unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: gate offline retry, safer tokenizer_name pop, skip audio net probe offline

- vision.py: only force the process-wide HF offline flag on the tokenizer
  retry when offline was requested or the captured primary error is actually
  network related, so a permanent tokenizer error no longer toggles global
  offline mode for other concurrent loads.
- loader.py: always pop tokenizer_name out of kwargs and let a caller-supplied
  value win, avoiding a "multiple values for keyword argument 'tokenizer_name'"
  TypeError when it is also passed explicitly downstream.
- model_config.py / export.py: add local_files_only to detect_audio_type so the
  raw requests.get tokenizer_config fetch is skipped offline (it ignores the HF
  offline flag), and pass it from the export probe.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: classify LocalEntryNotFoundError as offline-related

huggingface_hub's LocalEntryNotFoundError subclasses FileNotFoundError, so the
"not isinstance(cur, FileNotFoundError)" guard in _is_offline_related_error was
swallowing it and it could never be recognised as offline, despite being listed
in the network error types. It means "not in cache and the Hub is unreachable",
which is genuinely offline. Capture the class into an isinstance-checkable tuple
(empty, hence a no-op, if the import is unavailable) and exclude it from the
FileNotFoundError guard, so a real offline failure now triggers the local-cache
retry while a plain missing-file error still propagates.

* Address review: require merges.txt for BPE, status-gate HTTP errors, isolate local-only audio cache

- loader.py: a local dir with vocab.json but no merges.txt (and no tokenizer.json)
  is not a loadable BPE tokenizer, so do not treat it as self-sufficient; require
  merges.txt alongside vocab.json in both gate blocks, otherwise fall back to the
  base model tokenizer as before.
- vision.py: _is_offline_related_error no longer buckets every HfHubHTTPError /
  requests HTTPError as offline. HTTP errors are judged by status code: only a
  transient 5xx triggers the forced local-cache retry, while 401/403 (auth/gated)
  and 404 (missing) propagate as the real error instead of being masked. Hard
  signals (connection/timeout/OfflineModeIsEnabled/LocalEntryNotFoundError) still
  classify as offline.
- model_config.py: include local_files_only in the audio-detection cache key so a
  local-only (offline) negative result cannot be reused by a later online probe,
  which would otherwise route an audio model through the text loader until restart.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address re-review: fix studio test stubs, force offline env in probe window, drop redundant retry

- studio/backend/tests/test_vision_cache.py: the three _detect_audio_from_tokenizer
  stubs were called with the new local_files_only kwarg and raised TypeError, failing
  Backend CI. Add local_files_only to the stub signatures and add a test that a
  local-only negative does not poison a later online audio probe.
- export.py: the type-detection probe window now also sets HF_HUB_OFFLINE /
  TRANSFORMERS_OFFLINE env vars (saved/restored), not just the in-process flag.
  transformers_version._load_config_json / _check_tokenizer_config_needs_v5 gate
  their urllib fetches on the env vars, and is_vision_model may spawn a subprocess
  that inherits os.environ but not the in-process flag; without the env vars a
  probe-detected offline export could still block on a network timeout.
- vision.py: only retry the processor load when the first attempt was online and
  failed with a network error. When local_files_only was already requested the first
  attempt was forced offline, so the previous retry just repeated identical failing
  work before the last-resort path.
- model_config.py: correct the _audio_detection_cache type annotation to the 3-tuple
  key (name, token_fingerprint, local_files_only).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: thread-safe probe-offline env window, clear error for local dir without config

- export.py: guard the HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE mutation in
  _force_offline_probe_window with a lock + depth counter (mirrors _force_hf_offline),
  so concurrent / nested export probes only flip on first entry and restore on last
  exit. This prevents overlapping export requests from permanently poisoning those
  env vars or restoring a stale value.
- vision.py: in the VLM processor fallback, when tokenizer_name is a local directory,
  read its tokenizer_config.json directly and raise a clear FileNotFoundError if it is
  absent, instead of handing the local path to hf_hub_download (which would treat it as
  a repo id and raise a confusing HFValidationError / RepositoryNotFoundError).
  hf_hub_download is now only used for actual repo ids.

* Address review: classify raw socket.gaierror DNS failures as offline

Add the platform-specific getaddrinfo / DNS-resolution wording to the offline
detection list in _is_offline_related_error so a bare socket.gaierror (an OSError
subclass) is recovered from the local cache: "Name or service not known" and
"Temporary failure in name resolution" (Linux) and "nodename nor servname
provided" (macOS). Genuine non-network OSErrors (disk full, permission denied)
and plain FileNotFoundError still propagate.

* Address review: retry degraded VLM offline, force offline for text export + patch-tokenizer fallback

- vision.py: a degraded VLM processor (text-only, no image_processor) whose manual
  fallback fails offline used to be kept, so image inputs broke even with cached
  files. _construct_vlm_processor_fallback now returns its failure error;
  _acquire_processor surfaces it, and the caller retries forced-offline when the
  result is None OR a degraded VLM and the failure was network related, keeping the
  original result if the retry is not strictly better (never regress). The retry is
  still gated on an online first attempt + offline-related error so a permanent
  error never flips the global offline flag.
- vision.py: wrap the patch_tokenizer except-branch AutoTokenizer.from_pretrained in
  the same forced-offline-on-network-error pattern as the primary / last-resort
  loads, so an offline export where patch_tokenizer raises does not hang or fail.
- export.py: force HF offline around the two FastLanguageModel loads (text and SNAC)
  when the probe detected offline. Their text tokenizer path (load_correct_tokenizer
  -> AutoTokenizer) does not forward local_files_only, so without this a text export
  could still contact the Hub. Added a small _offline_window_if helper reused by the
  probe and load windows.

* Consolidate offline loading into one entry-point decision

Decide offline once per entry point instead of at every HF call site. The
prior approach threaded local_files_only into ~15 scattered config / tokenizer
/ processor / weight loads, each wrapped in its own try-online, classify-error,
retry-forced-offline dance, which is what kept surfacing "another call site you
missed", "another error shape misclassified", and global-flag thread-safety in
review.

FastLanguageModel / FastModel / FastBaseModel.from_pretrained now share an
@_offline_aware_load decorator: when offline (explicit local_files_only kwarg or
HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE env) it sets local_files_only and runs the
whole load inside one _force_hf_offline() window so every nested HF call inherits
it; when online it runs normally and, only if the load fails with a genuinely
network-related error, retries once forced-offline. The online path is unchanged
(no probe added) and 401 / 403 / 404 / permanent errors still propagate.

Centralise the offline helpers in loader_utils.py as the single source of truth
(shared by loader.py, re-exported from vision.py, and reused by the Studio
exporter):
- _force_hf_offline now sets the HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE env vars
  AND the in-process huggingface_hub / transformers flags, refcounted under one
  lock so nested / concurrent windows restore correctly. Setting the env vars
  covers env-gated urllib probes and spawned subprocesses too.
- _get_effective_local_files_only, _is_offline_related_error (unchanged
  classifier, retains the 5xx-vs-4xx, LocalEntryNotFound and gaierror handling),
  _offline_aware_load, and _resolve_checkpoint_tokenizer_name.

loader.py: wrap both entry points; drop the two duplicated env-var fallback
blocks and the two byte-identical local-tokenizer-gate blocks (now
_resolve_checkpoint_tokenizer_name).

vision.py: drop the per-site force_offline params and the three retry gates
(processor, patch_tokenizer fallback, last-resort). They now just surface the
underlying error so the single entry-point safety net retries forced-offline. A
network fallback error now takes precedence over a permanent primary error so the
offline retry still fires when the manual VLM fallback needs cached repo files.

studio/backend export.py: reuse the unified core _force_hf_offline (env + flags)
and drop the duplicate probe-window primitive; the snac / text branches no longer
need their own window. model_config.py: also gate the raw requests.get audio
fallback on the HF offline env vars so it is covered even without the kwarg.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address 10-reviewer P1 findings: vision cache split, PEFT offline, retry OOM

Split the Studio vision-detection cache by local_files_only, mirroring the audio
cache fix. is_vision_model / _is_vision_model_uncached / _raw_config_has_vision_config
/ load_model_config now thread local_files_only, the cache key includes it, and the
exporter passes it. Offline detection also skips the transformers-5 network
subprocess and stays on the local cache, so an offline negative can no longer be
keyed under the online entry and poison a later online probe. Adds a regression
test mirroring the audio poison test.

Forward local_files_only to both PeftModel.from_pretrained adapter-attach sites in
loader.py so a cached remote LoRA adapter resolves from the local cache under
explicit local-only / offline loads (defence-in-depth alongside the forced-offline
window).

_offline_aware_load: run the forced-offline retry OUTSIDE the except block and
collect + empty the device cache first. An except-scoped exception keeps its
__traceback__, which pins the failed attempt's frame locals (a partially loaded
model) until the block exits; loading the model again while that copy is still
alive could OOM a large VLM. Letting the except block close drops the traceback so
the partial load is freed before the retry reallocates.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address Codex review: env-offline cache key + rebuild HF sessions in offline window

Key the Studio audio and vision detection caches on the EFFECTIVE offline state
(local_files_only OR the HF offline env vars), not just the kwarg. detect_audio_type
and is_vision_model both skip the remote fetch / network subprocess when
HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE is set even with the default
local_files_only=False, so the result reflects offline; storing it under the online
(False) key let an env-offline negative poison a later online lookup once the env var
was cleared. Both now compute effective_offline once and use it for the cache key and
the downstream call. Adds a regression test for the env-offline dimension.

_force_hf_offline now rebuilds huggingface_hub's cached sessions on enter and exit
(best-effort _reset_hf_sessions). On hub 0.x the offline adapter is baked into the
per-thread requests.Session at creation, so flipping the constant alone leaves an
already-cached online session able to hit the network inside the window (and an
offline one stuck offline after restore); resetting forces the next get_session() to
match the current flag. On hub 1.x offline is checked dynamically per request, so
reset_sessions does not exist and the helper is a safe no-op.

The third review point (release the failed load before retrying) was already fixed in
af0f58a: the forced-offline retry now runs outside the except block and frees the
device cache first, so the failed attempt's traceback-pinned partial model is
released before the retry reallocates.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Align Studio _env_offline parsing with the canonical offline helper

model_config._env_offline gates the raw requests.get tokenizer-config fallback in
detect_audio_type and the audio/vision detection cache keys, but it only accepted
unstripped "1"/"true"/"yes". unsloth's offline helpers (loader_utils._env_says_offline
and the from_pretrained env fallback) accept the canonical set {1,true,yes,on} after
strip + lowercase, so HF_HUB_OFFLINE=on or HF_HUB_OFFLINE=" 1 " was treated as offline
by the loaders but online here, leaving the raw network fetch reachable while
"offline". Use the same strip + lowercase {1,true,yes,on} set. Adds parsing tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix lint: drop dead offline-helper re-exports from vision.py

The import-hoist verifier (scripts/verify_import_hoist.py) flagged vision.py's
re-export block as HOISTED-IMPORT-UNUSED blockers: it imported eight offline
helpers from loader_utils but only used three internally
(_get_effective_local_files_only, _is_offline_related_error, _offline_aware_load).
The other five were imported purely to preserve `from unsloth.models.vision import
X`, but nothing imports four of them from vision, and loader.py already imports
_resolve_checkpoint_tokenizer_name straight from loader_utils.

Import only the three names vision.py actually uses, and point the Studio exporter
at the canonical source (from unsloth.models.loader_utils import _force_hf_offline)
instead of re-exporting it through vision. loader_utils stays the single source of
truth; no behaviour change.

* Address Opus review: chain probe errors, unify env-offline, status-less HTTP

Chain the original AutoConfig/PeftConfig probe exception into the combined
RuntimeError in both FastLanguageModel.from_pretrained and FastModel.from_pretrained
(`raise RuntimeError(combined_error) from (autoconfig_exc or peft_exc)`). The probes
caught every Exception and stringified it, so the re-raised RuntimeError had no
__cause__/__context__ and _is_offline_related_error could not classify it -- the
network-down-but-cached auto-retry never fired for these entry points. With the
cause chained, the decorator sees a ConnectionError/LocalEntryNotFoundError/5xx and
retries forced-offline from cache; a permanent cause (404 / bad config) is still not
offline-classified and propagates without a wasted retry.

Unify the third offline-env parser: studio/backend/utils/transformers_version._env_offline
now uses the canonical {1,true,yes,on} + strip + lowercase set (matching
loader_utils._env_says_offline and model_config._env_offline), so HF_HUB_OFFLINE=on
or " 1 " no longer leaks the direct urllib metadata fetches to the network.

_is_offline_related_error: a status-less HTTP error (no response / unparseable code)
now falls back to the network-wording check instead of being dropped, so a transient
HTTP failure with clear "couldn't connect" wording is treated as offline. HTTP errors
with a real status code still decide by code (4xx propagates, 5xx is offline).

* Condense offline-loading code comments, drop dead helper, dedupe import for PR #6554

* Add unit tests for offline-loading helpers for PR #6554

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Guard load cleanup with try/finally and add retry-contract tests for PR #6554

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add gc.collect retry-step test for PR #6554

* Tighten offline-loading comments and docstrings for PR #6554

* Raise the both-config-failed error before model-type lookup so offline retry fires for PR #6554

* Prefer offline cause for retry and bound export reachability probe for PR #6554

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Skip remote mapper while offline, harden text-load cleanup, and stop stacked offline retries for PR #6554

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Surface VLM fallback offline errors, probe offline before export version activation, and restore progress bars across retries for PR #6554

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Restore offline env after export version activation so the persistent worker re-decides per load for PR #6554

* Classify socket.gaierror and urllib URLError as offline by type for PR #6554

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Probe offline around export load preflights and never offline-retry TLS failures for PR #6554

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Force in-process offline for export preflights, verify proxy egress in probe, and skip caching offline version negatives for PR #6554

* Snapshot offline constants before forcing env and require local processor files for VLM checkpoints for PR #6554

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-25 23:16:53 -07:00

2864 lines
106 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Model and LoRA configuration handling."""
from dataclasses import dataclass
from typing import Optional, Dict, Any
from utils.paths import (
normalize_path,
is_local_path,
is_model_cached,
get_cache_path,
resolve_cached_repo_id_case,
outputs_root,
exports_root,
resolve_output_dir,
resolve_export_dir,
)
from utils.utils import without_hf_auth
from utils.models.gguf_metadata import (
is_mmproj_by_metadata,
pairing_score,
read_gguf_general_metadata,
)
import structlog
from loggers import get_logger
import os
import re
import subprocess
import sys
from pathlib import Path
from typing import List, Tuple
import hashlib
import json
import threading
import yaml
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
logger = get_logger(__name__)
_OFFLINE_TRUE_VALUES = {"1", "true", "yes", "on"}
def _env_offline() -> bool:
"""True if an HF offline env var is truthy (canonical strip+lower parse, on/true/yes/1)."""
return (
os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES
or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES
)
# ── Model size extraction ────────────────────────────────────
import re as _re
_MODEL_SIZE_RE = _re.compile(r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
# MoE active-parameter pattern: "A3B", "A3.5B", etc.
_ACTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
# Gemma 3n/4 effective-parameter pattern: "E2B", "E4B" -- the runtime
# footprint (MatFormer + per-layer embeddings), which is the size that
# matters for size-gated policies like sub-3B speculative-decoding fallback.
_EFFECTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])e(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
def extract_model_size_b(model_id: str) -> float | None:
"""Extract model size in billions from a model identifier.
Prefers MoE active-parameter notation (e.g. ``A3B`` in
``Qwen3.5-35B-A3B``), then Gemma effective-parameter notation
(e.g. ``E2B``), over total params. Handles ``B`` (billions) and
``M`` (millions) suffixes.
"""
mid = (model_id or "").lower()
# First match wins, in priority order: active > effective > total.
for pattern in (_ACTIVE_SIZE_RE, _EFFECTIVE_SIZE_RE, _MODEL_SIZE_RE):
m = pattern.search(mid)
if m:
val = float(m.group(1))
return val / 1000.0 if m.group(2).lower() == "m" else val
return None
# Maps equivalent model names to their canonical YAML config file.
# Format: "canonical_model_name.yaml": [equivalent model names].
# Canonical filename derives from the first model name in each list.
MODEL_NAME_MAPPING = {
# ── Embedding models ──
"unsloth_all-MiniLM-L6-v2.yaml": [
"unsloth/all-MiniLM-L6-v2",
"sentence-transformers/all-MiniLM-L6-v2",
],
"unsloth_bge-m3.yaml": [
"unsloth/bge-m3",
"BAAI/bge-m3",
],
"unsloth_embeddinggemma-300m.yaml": [
"unsloth/embeddinggemma-300m",
"google/embeddinggemma-300m",
],
"unsloth_gte-modernbert-base.yaml": [
"unsloth/gte-modernbert-base",
"Alibaba-NLP/gte-modernbert-base",
],
"unsloth_Qwen3-Embedding-0.6B.yaml": [
"unsloth/Qwen3-Embedding-0.6B",
"Qwen/Qwen3-Embedding-0.6B",
"unsloth/Qwen3-Embedding-4B",
"Qwen/Qwen3-Embedding-4B",
],
# ── Other models ──
"unsloth_answerdotai_ModernBERT-large.yaml": [
"answerdotai/ModernBERT-large",
],
"unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml": [
"unsloth/Qwen2.5-Coder-7B-Instruct-bnb-4bit",
"unsloth/Qwen2.5-Coder-7B-Instruct",
"Qwen/Qwen2.5-Coder-7B-Instruct",
],
"unsloth_codegemma-7b-bnb-4bit.yaml": [
"unsloth/codegemma-7b-bnb-4bit",
"unsloth/codegemma-7b",
"google/codegemma-7b",
],
"unsloth_ERNIE-4.5-21B-A3B-PT.yaml": [
"unsloth/ERNIE-4.5-21B-A3B-PT",
],
"unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml": [
"unsloth/ERNIE-4.5-VL-28B-A3B-PT",
],
"tiiuae_Falcon-H1-0.5B-Instruct.yaml": [
"tiiuae/Falcon-H1-0.5B-Instruct",
"unsloth/Falcon-H1-0.5B-Instruct",
],
"unsloth_functiongemma-270m-it.yaml": [
"unsloth/functiongemma-270m-it-unsloth-bnb-4bit",
"google/functiongemma-270m-it",
"unsloth/functiongemma-270m-it-unsloth-bnb-4bit",
],
"unsloth_gemma-2-2b.yaml": [
"unsloth/gemma-2-2b-bnb-4bit",
"google/gemma-2-2b",
],
"unsloth_gemma-2-27b-bnb-4bit.yaml": [
"unsloth/gemma-2-9b-bnb-4bit",
"unsloth/gemma-2-9b",
"google/gemma-2-9b",
"unsloth/gemma-2-27b",
"google/gemma-2-27b",
],
"unsloth_gemma-3-4b-pt.yaml": [
"unsloth/gemma-3-4b-pt-unsloth-bnb-4bit",
"google/gemma-3-4b-pt",
"unsloth/gemma-3-4b-pt-bnb-4bit",
],
"unsloth_gemma-3-4b-it.yaml": [
"unsloth/gemma-3-4b-it-unsloth-bnb-4bit",
"google/gemma-3-4b-it",
"unsloth/gemma-3-4b-it-bnb-4bit",
],
"unsloth_gemma-3-27b-it.yaml": [
"unsloth/gemma-3-27b-it-unsloth-bnb-4bit",
"google/gemma-3-27b-it",
"unsloth/gemma-3-27b-it-bnb-4bit",
],
"unsloth_gemma-3-270m-it.yaml": [
"unsloth/gemma-3-270m-it-unsloth-bnb-4bit",
"google/gemma-3-270m-it",
"unsloth/gemma-3-270m-it-bnb-4bit",
],
"unsloth_gemma-3n-E4B-it.yaml": [
"unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit",
"google/gemma-3n-E4B-it",
"unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit",
],
"unsloth_gemma-3n-E4B.yaml": [
"unsloth/gemma-3n-E4B-unsloth-bnb-4bit",
"google/gemma-3n-E4B",
],
"unsloth_gemma-4-31B-it.yaml": [
"unsloth/gemma-4-31B-it",
"google/gemma-4-31B-it",
],
"unsloth_gemma-4-26B-A4B-it.yaml": [
"unsloth/gemma-4-26B-A4B-it",
"google/gemma-4-26B-A4B-it",
],
"unsloth_gemma-4-E2B-it.yaml": [
"unsloth/gemma-4-E2B-it",
"google/gemma-4-E2B-it",
],
"unsloth_gemma-4-E4B-it.yaml": [
"unsloth/gemma-4-E4B-it",
"google/gemma-4-E4B-it",
],
"unsloth_gemma-4-31B.yaml": [
"unsloth/gemma-4-31B",
"google/gemma-4-31B",
],
"unsloth_gemma-4-26B-A4B.yaml": [
"unsloth/gemma-4-26B-A4B",
"google/gemma-4-26B-A4B",
],
"unsloth_gemma-4-E2B.yaml": [
"unsloth/gemma-4-E2B",
"google/gemma-4-E2B",
],
"unsloth_gemma-4-E4B.yaml": [
"unsloth/gemma-4-E4B",
"google/gemma-4-E4B",
],
"unsloth_gpt-oss-20b.yaml": [
"openai/gpt-oss-20b",
"unsloth/gpt-oss-20b-unsloth-bnb-4bit",
"unsloth/gpt-oss-20b-BF16",
],
"unsloth_gpt-oss-120b.yaml": [
"openai/gpt-oss-120b",
"unsloth/gpt-oss-120b-unsloth-bnb-4bit",
],
"unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml": [
"unsloth/granite-4.0-350m",
"ibm-granite/granite-4.0-350m",
"unsloth/granite-4.0-350m-bnb-4bit",
],
"unsloth_granite-4.0-h-micro.yaml": [
"ibm-granite/granite-4.0-h-micro",
"unsloth/granite-4.0-h-micro-bnb-4bit",
"unsloth/granite-4.0-h-micro-unsloth-bnb-4bit",
],
"unsloth_LFM2-1.2B.yaml": [
"unsloth/LFM2-1.2B",
],
"unsloth_llama-3-8b-bnb-4bit.yaml": [
"unsloth/llama-3-8b",
"meta-llama/Meta-Llama-3-8B",
],
"unsloth_llama-3-8b-Instruct-bnb-4bit.yaml": [
"unsloth/llama-3-8b-Instruct",
"meta-llama/Meta-Llama-3-8B-Instruct",
],
"unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml": [
"unsloth/Meta-Llama-3.1-8B-bnb-4bit",
"unsloth/Meta-Llama-3.1-8B-unsloth-bnb-4bit",
"meta-llama/Meta-Llama-3.1-8B",
"unsloth/Meta-Llama-3.1-70B-bnb-4bit",
"unsloth/Meta-Llama-3.1-8B",
"unsloth/Meta-Llama-3.1-70B",
"meta-llama/Meta-Llama-3.1-70B",
"unsloth/Meta-Llama-3.1-405B-bnb-4bit",
"meta-llama/Meta-Llama-3.1-405B",
],
"unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml": [
"unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
"meta-llama/Meta-Llama-3.1-8B-Instruct",
"unsloth/Meta-Llama-3.1-8B-Instruct",
"RedHatAI/Llama-3.1-8B-Instruct-FP8",
"unsloth/Llama-3.1-8B-Instruct-FP8-Block",
"unsloth/Llama-3.1-8B-Instruct-FP8-Dynamic",
],
"unsloth_Llama-3.2-3B-Instruct.yaml": [
"unsloth/Llama-3.2-3B-Instruct-unsloth-bnb-4bit",
"meta-llama/Llama-3.2-3B-Instruct",
"unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
"RedHatAI/Llama-3.2-3B-Instruct-FP8",
"unsloth/Llama-3.2-3B-Instruct-FP8-Block",
"unsloth/Llama-3.2-3B-Instruct-FP8-Dynamic",
],
"unsloth_Llama-3.2-1B-Instruct.yaml": [
"unsloth/Llama-3.2-1B-Instruct-unsloth-bnb-4bit",
"meta-llama/Llama-3.2-1B-Instruct",
"unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
"RedHatAI/Llama-3.2-1B-Instruct-FP8",
"unsloth/Llama-3.2-1B-Instruct-FP8-Block",
"unsloth/Llama-3.2-1B-Instruct-FP8-Dynamic",
],
"unsloth_Llama-3.2-11B-Vision-Instruct.yaml": [
"unsloth/Llama-3.2-11B-Vision-Instruct-unsloth-bnb-4bit",
"meta-llama/Llama-3.2-11B-Vision-Instruct",
"unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit",
],
"unsloth_Llama-3.3-70B-Instruct.yaml": [
"unsloth/Llama-3.3-70B-Instruct-unsloth-bnb-4bit",
"meta-llama/Llama-3.3-70B-Instruct",
"unsloth/Llama-3.3-70B-Instruct-bnb-4bit",
"RedHatAI/Llama-3.3-70B-Instruct-FP8",
"unsloth/Llama-3.3-70B-Instruct-FP8-Block",
"unsloth/Llama-3.3-70B-Instruct-FP8-Dynamic",
],
"unsloth_Llasa-3B.yaml": [
"HKUSTAudio/Llasa-1B",
"unsloth/Llasa-3B",
],
"unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml": [
"unsloth/Magistral-Small-2509",
"mistralai/Magistral-Small-2509",
"unsloth/Magistral-Small-2509-bnb-4bit",
],
"unsloth_Ministral-3-3B-Instruct-2512.yaml": [
"unsloth/Ministral-3-3B-Instruct-2512",
],
"unsloth_mistral-7b-v0.3-bnb-4bit.yaml": [
"unsloth/mistral-7b-v0.3-bnb-4bit",
"unsloth/mistral-7b-v0.3",
"mistralai/Mistral-7B-v0.3",
],
"unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml": [
"unsloth/Mistral-Nemo-Base-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Base-2407",
"mistralai/Mistral-Nemo-Base-2407",
"unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Instruct-2407",
"mistralai/Mistral-Nemo-Instruct-2407",
],
"unsloth_Mistral-Small-Instruct-2409.yaml": [
"unsloth/Mistral-Small-Instruct-2409-bnb-4bit",
"mistralai/Mistral-Small-Instruct-2409",
],
"unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml": [
"unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
"unsloth/mistral-7b-instruct-v0.3",
"mistralai/Mistral-7B-Instruct-v0.3",
],
"unsloth_Qwen2.5-1.5B-Instruct.yaml": [
"unsloth/Qwen2.5-1.5B-Instruct-unsloth-bnb-4bit",
"Qwen/Qwen2.5-1.5B-Instruct",
"unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit",
],
"unsloth_Nemotron-3-Nano-30B-A3B.yaml": [
"unsloth/Nemotron-3-Nano-30B-A3B",
],
"unsloth_orpheus-3b-0.1-ft.yaml": [
"unsloth/orpheus-3b-0.1-ft",
"unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit",
"canopylabs/orpheus-3b-0.1-ft",
"unsloth/orpheus-3b-0.1-ft-bnb-4bit",
],
"OuteAI_Llama-OuteTTS-1.0-1B.yaml": [
"OuteAI/Llama-OuteTTS-1.0-1B",
"unsloth/Llama-OuteTTS-1.0-1B",
"unsloth/llama-outetts-1.0-1b",
"OuteAI/OuteTTS-1.0-0.6B",
"unsloth/OuteTTS-1.0-0.6B",
"unsloth/outetts-1.0-0.6b",
],
"unsloth_PaddleOCR-VL.yaml": [
"unsloth/PaddleOCR-VL",
],
"unsloth_Phi-3-medium-4k-instruct.yaml": [
"unsloth/Phi-3-medium-4k-instruct-bnb-4bit",
"microsoft/Phi-3-medium-4k-instruct",
],
"unsloth_Phi-3.5-mini-instruct.yaml": [
"unsloth/Phi-3.5-mini-instruct-bnb-4bit",
"microsoft/Phi-3.5-mini-instruct",
],
"unsloth_Phi-4.yaml": [
"unsloth/phi-4-unsloth-bnb-4bit",
"microsoft/phi-4",
"unsloth/phi-4-bnb-4bit",
],
"unsloth_Pixtral-12B-2409.yaml": [
"unsloth/Pixtral-12B-2409-unsloth-bnb-4bit",
"mistralai/Pixtral-12B-2409",
"unsloth/Pixtral-12B-2409-bnb-4bit",
],
"unsloth_Qwen2-7B.yaml": [
"unsloth/Qwen2-7B-bnb-4bit",
"Qwen/Qwen2-7B",
],
"unsloth_Qwen2-VL-7B-Instruct.yaml": [
"unsloth/Qwen2-VL-7B-Instruct-unsloth-bnb-4bit",
"Qwen/Qwen2-VL-7B-Instruct",
"unsloth/Qwen2-VL-7B-Instruct-bnb-4bit",
],
"unsloth_Qwen2.5-7B.yaml": [
"unsloth/Qwen2.5-7B-unsloth-bnb-4bit",
"Qwen/Qwen2.5-7B",
"unsloth/Qwen2.5-7B-bnb-4bit",
],
"unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml": [
"unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit",
"Qwen/Qwen2.5-Coder-1.5B-Instruct",
],
"unsloth_Qwen2.5-Coder-14B-Instruct.yaml": [
"unsloth/Qwen2.5-Coder-14B-Instruct-bnb-4bit",
"Qwen/Qwen2.5-Coder-14B-Instruct",
],
"unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml": [
"unsloth/Qwen2.5-VL-7B-Instruct",
"Qwen/Qwen2.5-VL-7B-Instruct",
"unsloth/Qwen2.5-VL-7B-Instruct-unsloth-bnb-4bit",
],
"unsloth_Qwen3-0.6B.yaml": [
"unsloth/Qwen3-0.6B-unsloth-bnb-4bit",
"Qwen/Qwen3-0.6B",
"unsloth/Qwen3-0.6B-bnb-4bit",
"Qwen/Qwen3-0.6B-FP8",
"unsloth/Qwen3-0.6B-FP8",
],
"unsloth_Qwen3-4B-Instruct-2507.yaml": [
"unsloth/Qwen3-4B-Instruct-2507-unsloth-bnb-4bit",
"Qwen/Qwen3-4B-Instruct-2507",
"unsloth/Qwen3-4B-Instruct-2507-bnb-4bit",
"Qwen/Qwen3-4B-Instruct-2507-FP8",
"unsloth/Qwen3-4B-Instruct-2507-FP8",
],
"unsloth_Qwen3-4B-Thinking-2507.yaml": [
"unsloth/Qwen3-4B-Thinking-2507-unsloth-bnb-4bit",
"Qwen/Qwen3-4B-Thinking-2507",
"unsloth/Qwen3-4B-Thinking-2507-bnb-4bit",
"Qwen/Qwen3-4B-Thinking-2507-FP8",
"unsloth/Qwen3-4B-Thinking-2507-FP8",
],
"unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml": [
"unsloth/Qwen3-14B-Base",
"Qwen/Qwen3-14B-Base",
"unsloth/Qwen3-14B-Base-bnb-4bit",
],
"unsloth_Qwen3-14B.yaml": [
"unsloth/Qwen3-14B-unsloth-bnb-4bit",
"Qwen/Qwen3-14B",
"unsloth/Qwen3-14B-bnb-4bit",
"Qwen/Qwen3-14B-FP8",
"unsloth/Qwen3-14B-FP8",
],
"unsloth_Qwen3-32B.yaml": [
"unsloth/Qwen3-32B-unsloth-bnb-4bit",
"Qwen/Qwen3-32B",
"unsloth/Qwen3-32B-bnb-4bit",
"Qwen/Qwen3-32B-FP8",
"unsloth/Qwen3-32B-FP8",
],
"unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml": [
"Qwen/Qwen3-VL-8B-Instruct-FP8",
"unsloth/Qwen3-VL-8B-Instruct-FP8",
"unsloth/Qwen3-VL-8B-Instruct",
"Qwen/Qwen3-VL-8B-Instruct",
"unsloth/Qwen3-VL-8B-Instruct-bnb-4bit",
],
"sesame_csm-1b.yaml": [
"sesame/csm-1b",
"unsloth/csm-1b",
],
"Spark-TTS-0.5B_LLM.yaml": [
"Spark-TTS-0.5B/LLM",
"unsloth/Spark-TTS-0.5B",
],
"unsloth_tinyllama-bnb-4bit.yaml": [
"unsloth/tinyllama",
"TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
],
"unsloth_whisper-large-v3.yaml": [
"unsloth/whisper-large-v3",
"openai/whisper-large-v3",
],
}
# Reverse lookup: model_name -> canonical_filename
_REVERSE_MODEL_MAPPING = {}
for canonical_file, model_names in MODEL_NAME_MAPPING.items():
for model_name in model_names:
_REVERSE_MODEL_MAPPING[model_name.lower()] = canonical_file
def load_model_config(
model_name: str,
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.
``trust_remote_code`` defaults to ``False``: capability detection and
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,
local_files_only = local_files_only,
)
if not use_auth:
# No auth, for public model checks
with without_hf_auth():
return AutoConfig.from_pretrained(
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,
)
# Detection sets come from the installed transformers registry, unioned with a
# small curated set of auto_map VLMs (DeepSeek-OCR, Kimi, phi3_v) whose arch is
# repo-defined and absent from the registry. ForConditionalGeneration is NOT a
# vision signal (overloaded across text/audio/vision); ForVisionText2Text is.
_VLM_ARCH_SUFFIXES = ("ForVisionText2Text",)
_CURATED_REMOTE_VLM_TYPES = frozenset(
{
"phi3_v",
"llava",
"llava_next",
"llava_onevision",
"internvl_chat",
"cogvlm2",
"minicpmv",
"gemma4",
"deepseek_vl_v2",
"kimi_k25",
}
)
# Fallbacks used only if the transformers registry import fails.
_FALLBACK_AUDIO_MODEL_TYPES = frozenset({"csm", "whisper"})
def _build_detection_sets():
"""Return (vlm_model_types, vlm_class_names, audio_model_types) from the
installed transformers registry, unioned with the curated repo-code VLM
set. Reads only static name dicts -- no model is loaded, no code runs.
Falls back to curated/hardcoded values if transformers is unavailable.
"""
try:
from transformers.models.auto import modeling_auto as _ma
def _names(attr):
d = getattr(_ma, attr, None)
return dict(d) if d else {}
itt = _names("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES")
v2s = _names("MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES")
vlm_types = set(itt) | set(v2s) | set(_CURATED_REMOTE_VLM_TYPES)
vlm_classes = set(itt.values()) | set(v2s.values())
audio_types: set = set()
for attr in (
"MODEL_FOR_CTC_MAPPING_NAMES",
"MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES",
"MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES",
"MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES",
"MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES",
"MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES",
):
audio_types |= set(_names(attr))
audio_types |= set(_FALLBACK_AUDIO_MODEL_TYPES)
return frozenset(vlm_types), frozenset(vlm_classes), frozenset(audio_types)
except Exception as exc: # pragma: no cover - defensive
logger.warning("Could not build detection sets from transformers: %s", exc)
return (
frozenset(_CURATED_REMOTE_VLM_TYPES),
frozenset(),
frozenset(_FALLBACK_AUDIO_MODEL_TYPES),
)
_VLM_MODEL_TYPES, _VLM_CLASS_NAMES, _AUDIO_ONLY_MODEL_TYPES = _build_detection_sets()
# Pre-computed .venv_t5 paths and backend dir for subprocess version switching.
# Vision check uses the Gemma 4 5.5 sidecar for existing Gemma 4 architectures.
from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402
_VENV_T5_DIR = str(_studio_root() / ".venv_t5_550")
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent)
def _is_vlm(config) -> bool:
architectures = getattr(config, "architectures", None) or []
model_type = getattr(config, "model_type", None)
explicit_vision = (
hasattr(config, "vision_config")
or hasattr(config, "img_processor")
or hasattr(config, "image_token_index")
or hasattr(config, "projector_config")
)
# Audio-only models are vision only if they carry an explicit vision sub-config.
if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:
return False
return (
explicit_vision
or any(x in _VLM_CLASS_NAMES for x in architectures)
or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
or model_type in _VLM_MODEL_TYPES
)
def _raw_config_has_vision_config(
model_name: str,
hf_token: Optional[str] = None,
local_files_only: bool = False,
) -> Optional[bool]:
try:
if is_local_path(model_name):
config_path = Path(normalize_path(model_name)).expanduser() / "config.json"
else:
from huggingface_hub import hf_hub_download
config_path = Path(
hf_hub_download(
repo_id = model_name,
filename = "config.json",
token = hf_token,
local_files_only = local_files_only,
)
)
config = json.loads(config_path.read_text())
architectures = config.get("architectures") or []
model_type = config.get("model_type")
explicit_vision = (
"vision_config" in config
or "img_processor" in config
or "image_token_index" in config
or "projector_config" in config
)
# Audio-only models are vision only if they carry an explicit vision sub-config.
if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:
return False
return (
explicit_vision
or any(isinstance(x, str) and x in _VLM_CLASS_NAMES for x in architectures)
or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
or model_type in _VLM_MODEL_TYPES
)
except Exception as exc:
logger.warning("Could not read config.json for '%s': %s", model_name, exc)
return None
# why: inline _is_vlm and constants are prepended so the subprocess stays
# self-contained and does not import the parent backend module graph.
_VISION_CHECK_INLINE_HELPERS = (
"_VLM_ARCH_SUFFIXES = " + repr(tuple(_VLM_ARCH_SUFFIXES)) + "\n"
"_VLM_MODEL_TYPES = " + repr(set(_VLM_MODEL_TYPES)) + "\n"
"_VLM_CLASS_NAMES = " + repr(set(_VLM_CLASS_NAMES)) + "\n"
"_AUDIO_ONLY_MODEL_TYPES = " + repr(set(_AUDIO_ONLY_MODEL_TYPES)) + "\n"
"def _is_vlm(config):\n"
" architectures = getattr(config, 'architectures', None) or []\n"
" model_type = getattr(config, 'model_type', None)\n"
" explicit_vision = (\n"
" hasattr(config, 'vision_config')\n"
" or hasattr(config, 'img_processor')\n"
" or hasattr(config, 'image_token_index')\n"
" or hasattr(config, 'projector_config')\n"
" )\n"
" if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:\n"
" return False\n"
" return (\n"
" explicit_vision\n"
" or any(x in _VLM_CLASS_NAMES for x in architectures)\n"
" or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n"
" or model_type in _VLM_MODEL_TYPES\n"
" )\n"
)
# Subprocess script run with transformers 5.x active. Takes model_name and
# token via argv, prints JSON result to stdout.
_VISION_CHECK_SCRIPT = (
r"""
import sys, os, json
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Activate transformers 5.x
venv_t5 = sys.argv[1]
backend_dir = sys.argv[2]
model_name = sys.argv[3]
token = sys.argv[4] if len(sys.argv) > 4 and sys.argv[4] != "" else None
sys.path.insert(0, venv_t5)
if backend_dir not in sys.path:
sys.path.insert(0, backend_dir)
"""
+ _VISION_CHECK_INLINE_HELPERS
+ r"""
try:
from transformers import AutoConfig
# Capability detection never executes model repo code.
kwargs = {"trust_remote_code": False}
if token:
kwargs["token"] = token
config = AutoConfig.from_pretrained(model_name, **kwargs)
is_vlm = _is_vlm(config)
model_type = getattr(config, "model_type", None)
archs = getattr(config, "architectures", [])
print(json.dumps({"is_vision": is_vlm, "model_type": model_type,
"architectures": archs}))
except Exception as exc:
print(json.dumps({"error": str(exc)}))
sys.exit(1)
"""
)
def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
"""Run is_vision_model in a subprocess with transformers 5.x.
Spawns a clean subprocess with .venv_t5/ on sys.path so AutoConfig
recognizes newer architectures. Returns True/False for definitive results,
or None for transient failures (timeouts, subprocess errors), which are not
cached so they can be retried.
"""
token_arg = hf_token or ""
try:
result = subprocess.run(
[
sys.executable,
"-c",
_VISION_CHECK_SCRIPT,
_VENV_T5_DIR,
_BACKEND_DIR,
model_name,
token_arg,
],
capture_output = True,
text = True,
timeout = 60,
env = child_env_without_native_path_secret(),
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
stderr = result.stderr.strip()
logger.warning(
"Vision check subprocess failed for '%s': %s",
model_name,
stderr or result.stdout.strip(),
)
return None
data = json.loads(result.stdout.strip())
if "error" in data:
logger.warning(
"Vision check subprocess error for '%s': %s",
model_name,
data["error"],
)
return None
is_vlm = data["is_vision"]
logger.info(
"Vision check (subprocess, transformers 5.x) for '%s': "
"model_type=%s, architectures=%s, is_vision=%s",
model_name,
data.get("model_type"),
data.get("architectures"),
is_vlm,
)
return is_vlm
except subprocess.TimeoutExpired:
logger.warning("Vision check subprocess timed out for '%s'", model_name)
return None
except Exception as exc:
logger.warning("Vision check subprocess failed for '%s': %s", model_name, exc)
return None
def _token_fingerprint(token: Optional[str]) -> Optional[str]:
"""SHA256 digest of the token for use as a cache key (avoids storing the
raw bearer token in process memory)."""
if token is None:
return None
return hashlib.sha256(token.encode("utf-8")).hexdigest()
# Vision detection cache keyed by (name, token, local_files_only); only definitive results cached.
_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,
local_files_only: bool = False,
) -> bool:
"""Detect VLMs via the config architecture (works for fine-tunes); transformers-5.x
models are checked in a .venv_t5/ subprocess. Cached per (model_name, token,
local_files_only) minus transient failures; local_files_only is in the key so an
offline probe never shares an online entry."""
# Local GGUF models are served by llama-server. Their multimodal
# capability comes from a companion mmproj, not a Transformers config.
# Do not cache this lookup: a projector may be added beside an existing
# weight file after it was first inspected.
if is_local_path(model_name):
local_path = normalize_path(model_name)
gguf_file = detect_gguf_model(local_path)
if gguf_file:
companion_root = _local_gguf_companion_search_root(local_path, gguf_file)
mmproj_file = detect_mmproj_file(gguf_file, search_root = companion_root)
is_vision = mmproj_file is not None
logger.debug(
"Local GGUF vision check for '%s': mmproj=%s, is_vision=%s",
gguf_file,
mmproj_file,
is_vision,
)
return is_vision
# Normalize model name so different casings of the same repo share a key
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
# Key on effective offline (kwarg OR env) so an offline probe can't poison a later
# online lookup once the env var is cleared.
effective_offline = bool(local_files_only or _env_offline())
cache_key = (resolved_name, _token_fingerprint(hf_token), effective_offline)
# Lock-free fast path for cache hits. Sentinel distinguishes "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 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, local_files_only = effective_offline)
# Only cache definitive results; None is a transient failure, retry later.
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,
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
(network, timeout, subprocess failure) so the caller knows not to cache.
"""
# 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, local_files_only = local_files_only
)
if raw is not None:
return raw
# Raw read failed transiently: fall back to AutoConfig (remote code DISABLED), via a
# transformers-5.x subprocess if needed. Skip that subprocess offline (it probes the network).
from utils.transformers_version import needs_transformers_5
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,
)
return _is_vision_model_subprocess(model_name, hf_token = hf_token)
try:
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)
archs = getattr(config, "architectures", None) or []
logger.info(
"Model %s detected as VLM (model_type=%s, architectures=%s)",
model_name,
model_type,
archs,
)
return True
return False
except Exception as e:
logger.warning(f"Could not determine if {model_name} is vision model: {e}")
# Permanent failures (not found, gated, bad config) cache as False;
# transient ones (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")
# Keyed like the vision cache by (name, token, local_files_only) so an unauthenticated
# or offline miss cannot poison a later authenticated / online lookup.
_audio_detection_cache: Dict[Tuple[str, Optional[str], bool], Optional[str]] = {}
# Tokenizer token patterns → audio_type (all 6 types from tokenizer_config.json)
_AUDIO_TOKEN_PATTERNS = {
"csm": lambda tokens: "<|AUDIO|>" in tokens and "<|audio_eos|>" in tokens,
"whisper": lambda tokens: "<|startoftranscript|>" in tokens,
# Gemma 3n: <audio_soft_token>; Gemma 4: <|audio|> (not csm's <|AUDIO|>).
"audio_vlm": lambda tokens: "<audio_soft_token>" in tokens or "<|audio|>" in tokens,
"bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens),
"dac": lambda tokens: (
"<|audio_start|>" in tokens
and "<|audio_end|>" in tokens
and "<|text_start|>" in tokens
and "<|text_end|>" in tokens
),
"snac": lambda tokens: (sum(1 for t in tokens if t.startswith("<custom_token_")) > 10000),
}
def detect_audio_type(
model_name: str,
hf_token: Optional[str] = None,
local_files_only: bool = False,
) -> Optional[str]:
"""Detect if a model is an audio model and return its type.
Works for any model via tokenizer_config.json special tokens.
Returns an audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper',
'audio_vlm') or None.
When local_files_only is True (offline export) the remote HuggingFace fetch
is skipped so detection never blocks on a network read; only the local HF
cache is consulted.
"""
# Normalize casing + include the token fingerprint (mirrors is_vision_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:
resolved_name = model_name
# Key on effective offline (kwarg OR env), matching where the remote fetch is skipped,
# so an offline negative can't poison a later online probe.
effective_offline = bool(local_files_only or _env_offline())
cache_key = (resolved_name, _token_fingerprint(hf_token), effective_offline)
if cache_key in _audio_detection_cache:
return _audio_detection_cache[cache_key]
result, definitive = _detect_audio_from_tokenizer(
model_name, hf_token, local_files_only = effective_offline
)
# Cache only definitive results; a transient read failure stays None and retries.
if definitive:
_audio_detection_cache[cache_key] = result
if result:
logger.info(f"Model {model_name} detected as audio model: audio_type={result}")
return result
def _detect_audio_from_tokenizer(
model_name: str,
hf_token: Optional[str] = None,
local_files_only: bool = False,
) -> Tuple[Optional[str], bool]:
"""Detect audio type from tokenizer special tokens.
Checks local HF cache first, then (unless local_files_only) fetches
tokenizer_config.json from HF; examines added_tokens_decoder for distinctive
patterns.
Returns (audio_type_or_None, definitive). definitive is False only on a
transient read failure (network/timeout/5xx) so the caller skips caching and
retries; a successful read with no audio tokens is a definitive None.
"""
def _check_token_patterns(tok_config: dict) -> Optional[str]:
added = tok_config.get("added_tokens_decoder", {})
if not added:
return None
token_contents = [v.get("content", "") for v in added.values()]
for audio_type, check_fn in _AUDIO_TOKEN_PATTERNS.items():
if check_fn(token_contents):
return audio_type
return None
read_any = False # parsed at least one tokenizer_config -> a None is definitive
# 1) Local HF cache first (works for gated/offline models)
try:
repo_dir = get_cache_path(model_name)
if repo_dir is not None and repo_dir.exists():
snapshots_dir = repo_dir / "snapshots"
if snapshots_dir.exists():
for snapshot in snapshots_dir.iterdir():
for tok_path in [
"tokenizer_config.json",
"LLM/tokenizer_config.json",
]:
tok_file = snapshot / tok_path
if tok_file.exists():
tok_config = json.loads(tok_file.read_text())
read_any = True
result = _check_token_patterns(tok_config)
if result:
return result, True
except Exception as e:
logger.debug(f"Could not check local cache for {model_name}: {e}")
# 2) Fall back to the HuggingFace API. This raw requests.get ignores the HF offline
# flag, so gate it on local_files_only OR the env vars to skip the network offline.
if local_files_only or _env_offline():
return None, read_any
try:
import requests
import os
except Exception:
return None, read_any
paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
token = hf_token or os.environ.get("HF_TOKEN")
headers = {"Authorization": f"Bearer {token}"} if token else {}
transient = False # a fetch failed for a non-404 reason (network/5xx)
for tok_path in paths_to_try:
url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
try:
resp = requests.get(url, headers = headers, timeout = 15)
except Exception as e:
logger.debug(f"Could not fetch {tok_path} for {model_name}: {e}")
transient = True
continue
if resp.status_code == 404:
continue # genuinely absent on this path
if not resp.ok:
transient = True # 5xx/403/etc -- can't tell, don't cache
continue
try:
tok_config = resp.json()
except Exception as e:
logger.debug(f"Bad tokenizer_config for {model_name}/{tok_path}: {e}")
transient = True
continue
read_any = True
result = _check_token_patterns(tok_config)
if result:
return result, True
# No audio tokens: definitive unless every attempt failed transiently.
return None, (read_any or not transient)
def is_audio_input_type(audio_type: Optional[str]) -> bool:
"""True if an audio_type accepts audio input: whisper (ASR), audio_vlm (Gemma3n)."""
return audio_type in ("whisper", "audio_vlm")
def _is_mmproj(filename: str) -> bool:
"""Check if a GGUF filename is a vision projection (mmproj) file."""
return "mmproj" in filename.lower()
def _is_mtp_drafter(path: str) -> bool:
"""True for a separate-file MTP drafter (speculative head), a companion
to the main model rather than a selectable quant: the repo-root
``mtp-*.gguf`` or the ``MTP/`` subdir copies (Gemma 4).
Mirrors hub.utils.gguf.is_mtp_drafter_path (utils cannot import hub).
Must be excluded everywhere mmproj is, or the drafter leaks into variant
menus (a phantom quant) and quant-matched file lookups -- e.g. a ``Q8_0``
request must not resolve to ``MTP/...-Q8_0-MTP.gguf``, which sorts ahead
of the real weight.
"""
p = path.lower()
if not p.endswith(".gguf"):
return False
name = p.rsplit("/", 1)[-1]
return name.startswith("mtp-") or "/mtp/" in f"/{p}"
# Family tokens for #5347's filename fallback. Lowercase; order irrelevant.
_MODEL_FAMILY_TOKENS: tuple[str, ...] = (
"qwen",
"gemma",
"llama",
"mistral",
"ministral",
"magistral",
"devstral",
"phi",
"deepseek",
"internvl",
"minicpm",
"llava",
"glm",
"yi",
"command-r",
"molmo",
"pixtral",
"smolvlm",
"moondream",
"granite",
"ovis",
"nemotron",
"kimi",
"nanonets",
"cosmos",
"mimo",
"apriel",
"lfm",
)
# Word-bounded match: a letter on either side disqualifies (stops ``phi``
# matching ``sapphire``, ``yi`` matching ``tiny``).
_FAMILY_TOKEN_RE_CACHE: Dict[str, "_re.Pattern[str]"] = {}
def _family_token_re(token: str) -> "_re.Pattern[str]":
pat = _FAMILY_TOKEN_RE_CACHE.get(token)
if pat is None:
pat = _re.compile(rf"(?:^|[^a-z])({_re.escape(token)})(?:[^a-z]|$)")
_FAMILY_TOKEN_RE_CACHE[token] = pat
return pat
def _detect_family_token(filename: str) -> Optional[str]:
"""Leftmost-position match; ties prefer the longer token."""
name = filename.lower()
best: Optional[tuple[int, int, str]] = None # (start, -len, token)
for token in _MODEL_FAMILY_TOKENS:
m = _family_token_re(token).search(name)
if m is None:
continue
key = (m.start(1), -len(token), token)
if best is None or key < best:
best = key
return None if best is None else best[2]
def mmproj_matches_model_family(model_path: str, mmproj_path: str) -> bool:
"""Launcher guard: True unless both filenames carry recognised family
tokens that disagree."""
model_fam = _detect_family_token(Path(model_path).name)
mmproj_fam = _detect_family_token(Path(mmproj_path).name)
if model_fam is None or mmproj_fam is None:
return True
return model_fam == mmproj_fam
def _shared_prefix_len(a: str, b: str) -> int:
n = min(len(a), len(b))
for i in range(n):
if a[i] != b[i]:
return i
return n
def _is_gguf_filename(filename: str) -> bool:
return filename.lower().endswith(".gguf")
def _iter_gguf_files(directory: Path, recursive: bool = False):
if not directory.is_dir():
return
iterator = directory.rglob("*") if recursive else directory.iterdir()
for f in iterator:
if f.is_file() and _is_gguf_filename(f.name):
yield f
def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]:
"""Find the mmproj GGUF for a model.
``path``: directory or a .gguf file. ``search_root``: optional ancestor
to also walk (snapshot layouts where the weight is in ``snapshot/BF16/``
but the projector sits at ``snapshot/``). Returns the projector path or
``None``."""
p = Path(path)
start_dir = p.parent if p.is_file() else p
if not start_dir.is_dir():
return None
# Walk incrementally so a sibling subdir's mmproj cannot leak in.
seen: set[Path] = set()
scan_order: list[Path] = []
def _add(d: Path) -> None:
try:
resolved = d.resolve()
except OSError:
return
if resolved in seen or not resolved.is_dir():
return
seen.add(resolved)
scan_order.append(resolved)
_add(start_dir)
# Ollama's .studio_links/foo.gguf -> blobs/sha256-...: also scan target dir.
try:
if p.is_symlink() and p.is_file():
target_parent = p.resolve().parent
if target_parent.is_dir():
_add(target_parent)
except OSError:
pass
if search_root is not None:
try:
root_resolved = Path(search_root).resolve()
start_resolved = start_dir.resolve()
if root_resolved == start_resolved or (
start_resolved.is_relative_to(root_resolved)
if hasattr(start_resolved, "is_relative_to")
else str(start_resolved).startswith(str(root_resolved) + "/")
):
cur = start_resolved
while cur != root_resolved and cur.parent != cur:
cur = cur.parent
_add(cur)
if cur == root_resolved:
break
except OSError:
pass
candidates: list[Path] = []
seen_resolved: set[Path] = set()
for d in scan_order:
for f in _iter_gguf_files(d):
try:
resolved = f.resolve()
except OSError:
continue
if resolved in seen_resolved:
continue
# Prefer ``general.type=='mmproj'``, else filename.
meta = read_gguf_general_metadata(str(resolved))
by_meta = is_mmproj_by_metadata(meta)
if by_meta is True or (by_meta is None and _is_mmproj(f.name)):
seen_resolved.add(resolved)
candidates.append(resolved)
if not candidates:
return None
# Directory path: no model name to compare against; legacy behaviour.
if not p.is_file():
return str(candidates[0])
# Stage 1: GGUF metadata. Stage 2: filename family token (#5347).
model_stem = p.stem.lower()
model_family = _detect_family_token(p.name)
weight_meta = read_gguf_general_metadata(str(p))
scored: list[tuple[int, Path]] = []
for c in candidates:
cand_meta = read_gguf_general_metadata(str(c))
meta_score = pairing_score(weight_meta, cand_meta)
if meta_score == -1:
logger.info(f"detect_mmproj_file: dropped {c.name} (metadata mismatch)")
continue
if meta_score == 0 and model_family is not None:
# Unrecognised candidate family is a wildcard (``mmproj-F16.gguf``).
cand_family = _detect_family_token(c.name)
if cand_family is not None and cand_family != model_family:
logger.info(
f"detect_mmproj_file: dropped {c.name} "
f"(filename family {cand_family!r} vs model {model_family!r})"
)
continue
scored.append((meta_score, c))
if not scored:
return None
# Score first, then longest shared prefix, then shorter stem.
best = max(
scored,
key = lambda sc: (
sc[0],
_shared_prefix_len(model_stem, sc[1].stem.lower()),
-len(sc[1].stem),
),
)
return str(best[1])
def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[str]:
"""Find the separate MTP drafter (``mtp-*.gguf``) for a local GGUF model.
The drafter that pairs with the main weights sits at the repo/snapshot
root (Gemma 4); the weight itself may be at the root or in a quant subdir,
so scan the weight's directory and ``search_root``. Matches by the
``mtp-`` filename prefix unsloth uses for ``-hf`` auto-discovery -- the
same signal as the HF download path. Repos that bake the head into the
main GGUF (Qwen) have no such sibling, so this returns None.
Pairs by name so a multi-model folder can't attach a foreign drafter:
unsloth names the drafter ``mtp-<model>.gguf`` where ``<model>`` prefixes
the weight filename across all Gemma 4 repos (e.g.
``mtp-gemma-4-12B-it.gguf`` next to ``gemma-4-12B-it-qat-Q4_0.gguf``).
An unmatched drafter is skipped (fail-safe: no MTP).
"""
p = Path(path)
weight_name = p.name.lower() if p.suffix.lower() == ".gguf" else None
start_dir = p.parent if p.is_file() else p
dirs = [start_dir]
if search_root is not None:
dirs.append(Path(search_root))
for d in dirs:
try:
entries = sorted(d.iterdir())
except OSError:
continue
for f in entries:
name = f.name.lower()
if not (name.startswith("mtp-") and name.endswith(".gguf")):
continue
stem = name[len("mtp-") : -len(".gguf")]
if not stem or (weight_name is not None and not weight_name.startswith(stem)):
continue
try:
if f.is_file():
return str(f.resolve())
except OSError:
continue
return None
def detect_gguf_model(path: str) -> Optional[str]:
"""Check if a local path is or contains a GGUF model file.
Handles a direct .gguf path or a directory of .gguf files. Skips mmproj
files (pass those via ``--mmproj``; see :func:`detect_mmproj_file`). Returns
the .gguf path or None. For HF repos, use detect_gguf_model_remote().
"""
p = Path(path)
# Case 1: direct .gguf file
if p.suffix.lower() == ".gguf":
# Companions are not models: rejecting a drafter here also keeps
# detect_mtp_file from pairing the same file with itself
# (-m drafter --model-draft drafter). Include the immediate parent
# dir so the MTP/ subdir copies are caught -- the basename alone
# (...-MTP.gguf) doesn't match the predicate's mtp- prefix.
rel = f"{p.parent.name}/{p.name}"
quant = _extract_quant_label(rel)
if _is_mmproj(p.name) or _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
return None
# Extension is authoritative: don't gate on is_file()/exists(), which
# can fail in the Windows lock window after llama-server is killed.
try:
is_dir = p.is_dir()
except OSError:
is_dir = False # stat() unavailable in the lock window
if not is_dir:
return str(p.absolute()) # absolute() keeps symlink names readable
# Directory named "*.gguf": fall through to the dir scan below.
# Case 2: directory containing .gguf files (skip mmproj / MTP drafter)
if p.is_dir():
gguf_files = []
for f in _iter_gguf_files(p):
context_rel = f"{f.parent.name}/{f.name}"
quant = _extract_quant_label(context_rel)
if (
_is_mmproj(f.name)
or _is_mtp_drafter(context_rel)
or _is_big_endian_gguf_path(context_rel, quant)
):
continue
gguf_files.append(f)
gguf_files.sort(key = lambda f: f.stat().st_size, reverse = True)
if gguf_files:
return str(gguf_files[0].resolve())
return None
# Preferred GGUF quant levels, descending priority. UD (Unsloth Dynamic)
# variants beat standard quants on quality per bit; repos without UD fall back
# to standard quants. Ordered by size/quality tradeoff, not raw quality.
_GGUF_QUANT_PREFERENCE = [
# UD variants (best quality per bit) -- Q4 is the sweet spot
"UD-Q4_K_XL",
"UD-Q4_K_L",
"UD-Q5_K_XL",
"UD-Q3_K_XL",
"UD-Q6_K_XL",
"UD-Q6_K_S",
"UD-Q8_K_XL",
"UD-Q2_K_XL",
"UD-IQ4_NL",
"UD-IQ4_XS",
"UD-IQ3_S",
"UD-IQ3_XXS",
"UD-IQ2_M",
"UD-IQ2_XXS",
"UD-IQ1_M",
"UD-IQ1_S",
# Standard quants (fallback for non-Unsloth repos)
"Q4_K_M",
"Q4_K_S",
"Q5_K_M",
"Q5_K_S",
"Q6_K",
"Q8_0",
"Q3_K_M",
"Q3_K_L",
"Q3_K_S",
"Q2_K",
"Q2_K_L",
"IQ4_NL",
"IQ4_XS",
"IQ3_M",
"IQ3_XXS",
"IQ2_M",
"IQ1_M",
"F16",
"BF16",
"F32",
]
def _pick_best_gguf(filenames: list[str]) -> Optional[str]:
"""Pick the best GGUF file: quant levels in _GGUF_QUANT_PREFERENCE order, else first .gguf."""
gguf_files = [f for f in filenames if f.lower().endswith(".gguf")]
if not gguf_files:
return None
for quant in _GGUF_QUANT_PREFERENCE:
for f in gguf_files:
if quant in f:
return f
return gguf_files[0]
@dataclass
class GgufVariantInfo:
"""A single GGUF quantization variant from a HuggingFace repo."""
filename: str # e.g., "gemma-3-4b-it-Q4_K_M.gguf"
quant: str # e.g., "Q4_K_M" (extracted from filename)
size_bytes: int # file size
def _extract_quant_label(filename: str) -> str:
"""
Extract quant label like Q4_K_M, IQ4_XS, BF16 from a GGUF filename.
Examples:
"gemma-3-4b-it-Q4_K_M.gguf""Q4_K_M"
"model-IQ4_NL.gguf""IQ4_NL"
"model-BF16.gguf""BF16"
"model-UD-IQ1_S.gguf""UD-IQ1_S"
"model-UD-TQ1_0.gguf""UD-TQ1_0"
"MXFP4_MOE/model-MXFP4_MOE-0001.gguf""MXFP4_MOE"
"Qwen3.6-IQ4_XS-3.53bpw.gguf""IQ4_XS-3.53bpw"
"""
import re
basename = filename.rsplit("/", 1)[-1]
# Strip .gguf and any shard suffix (-00001-of-00010)
stem = re.sub(r"-\d{3,}-of-\d{3,}", "", basename.rsplit(".", 1)[0])
quant_re = (
r"(UD-)?" # Optional UD- prefix (Ultra Discrete)
r"(MXFP[0-9]+(?:_[A-Z0-9]+)*" # MXFP variants: MXFP4, MXFP4_MOE
r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
r"|TQ[0-9]+_[0-9]+" # Ternary quant: TQ1_0, TQ2_0
r"|Q[0-9]+_K_[A-Z]+" # K-quant: Q4_K_M, Q3_K_S
r"|Q[0-9]+_[0-9]+" # Standard: Q8_0, Q5_1
r"|Q[0-9]+_K" # Short K-quant: Q6_K
r"|BF16|F16|F32)" # Full precision
# Optional bits-per-weight modifier so repos that ship multiple
# files at the same base quant (e.g. byteshape's IQ4_XS at 3.53,
# 3.97, 4.19 bpw) don't collapse into a single merged variant.
r"(-[0-9]+(?:\.[0-9]+)?bpw)?"
)
match = re.search(quant_re, stem, re.IGNORECASE)
# Subdir layouts like ``BF16/foo.gguf`` keep the quant in the directory,
# not the basename. Check parent dirs too so the label matches the
# snapshot-relative path produced elsewhere.
if not match and "/" in filename:
parents = filename.rsplit("/", 1)[0]
for segment in reversed(parents.split("/")):
m = re.search(quant_re, segment, re.IGNORECASE)
if m:
match = m
break
if match:
prefix = match.group(1) or ""
bpw = match.group(3) or ""
return f"{prefix}{match.group(2)}{bpw}"
# Fallback: last hyphen-separated segment
return stem.split("-")[-1]
_BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE)
_GGUF_KNOWN_QUANT_RE = re.compile(
r"(UD-)?"
r"(MXFP[0-9]+(?:_[A-Z0-9]+)*"
r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?"
r"|TQ[0-9]+_[0-9]+"
r"|Q[0-9]+_K_[A-Z]+"
r"|Q[0-9]+_[0-9]+"
r"|Q[0-9]+_K"
r"|BF16|F16|F32)",
re.IGNORECASE,
)
def _is_big_endian_gguf_path(path: str, quant: str = "") -> bool:
normalized = path.replace("\\", "/")
name = normalized.rsplit("/", 1)[-1]
stem = name.rsplit(".", 1)[0].lower()
quant_key = quant.strip().lower()
quant_index = stem.find(quant_key) if quant_key else -1
parent = normalized.rsplit("/", 1)[0].lower() if "/" in normalized else ""
quant_in_parent_only = (
bool(parent)
and quant_index < 0
and (
(quant_key and quant_key in parent)
or (not quant_key and _GGUF_KNOWN_QUANT_RE.search(parent) is not None)
)
)
for match in _BIG_ENDIAN_GGUF_FILENAME_RE.finditer(stem):
if quant_index >= 0 and quant_index < match.start():
return True
tail = stem[match.end() :].lstrip("._-")
if not tail or _GGUF_KNOWN_QUANT_RE.search(tail) is None:
return not quant_in_parent_only
return False
def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str:
"""Directory to scan upward from for local GGUF companion files."""
import re
selected = Path(selected_path)
gguf_path = Path(gguf_file)
if selected.suffix.lower() != ".gguf":
return selected_path
gguf_dir = gguf_path.parent
if not gguf_dir.name:
return str(gguf_dir)
quant_dir_re = (
r"(UD-)?("
r"MXFP[0-9]+(?:_[A-Z0-9]+)*"
r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?"
r"|TQ[0-9]+_[0-9]+"
r"|Q[0-9]+_K_[A-Z]+"
r"|Q[0-9]+_[0-9]+"
r"|Q[0-9]+_K"
r"|BF16|F16|F32"
r")"
)
if re.fullmatch(quant_dir_re, gguf_dir.name, re.IGNORECASE):
return str(gguf_dir.parent)
return str(gguf_dir)
def _iter_hf_cache_snapshots(repo_id: str):
"""Yield HF cache snapshot dirs for *repo_id*, newest first.
Empty if HF_HUB_CACHE is missing, the repo isn't cached, or has no
snapshots. Repo name match is case-insensitive to handle casing drift
between download time and lookup.
"""
try:
from huggingface_hub import constants as hf_constants
except Exception:
return
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
repo_dir: Optional[Path] = None
try:
if not cache_dir.is_dir():
return
for entry in cache_dir.iterdir():
if entry.is_dir() and entry.name.lower() == target:
repo_dir = entry
break
except OSError:
return
if repo_dir is None:
return
snapshots = repo_dir / "snapshots"
try:
if not snapshots.is_dir():
return
snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()]
except OSError:
return
snap_dirs.sort(key = lambda s: s.stat().st_mtime, reverse = True)
yield from snap_dirs
def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]:
"""Variants from the local HF cache snapshot, or None if not cached."""
for snap in _iter_hf_cache_snapshots(repo_id):
variants, has_vision = list_local_gguf_variants(str(snap))
if variants or has_vision:
return variants, has_vision
return None
def list_gguf_variants(
repo_id: str, hf_token: Optional[str] = None
) -> tuple[list[GgufVariantInfo], bool]:
"""List all GGUF quant variants in a HF repo.
Separates main model files from mmproj (vision projection) files; mmproj
presence flags a vision-capable model.
Returns:
(variants, has_vision): non-mmproj GGUF variants + vision flag.
"""
from huggingface_hub import model_info as hf_model_info
# Offline: skip the API and serve from cache
if _env_offline():
cached = _list_gguf_variants_from_hf_cache(repo_id)
if cached is not None:
return cached
try:
info = hf_model_info(repo_id, token = hf_token, files_metadata = True)
except Exception as e:
# Permanent errors (deleted/gated/bad revision) must surface to the
# caller; serving stale cache would mask the real cause. Matches the
# early-return in ``detect_gguf_model_remote``.
if type(e).__name__ in (
"RepositoryNotFoundError",
"GatedRepoError",
"RevisionNotFoundError",
"EntryNotFoundError",
):
raise
# API failed transiently; fall back to local snapshot if fully downloaded.
cached = _list_gguf_variants_from_hf_cache(repo_id)
if cached is not None:
logger.warning(
"HF API unreachable for %s (%s); using local cache snapshot.",
repo_id,
e.__class__.__name__,
)
return cached
raise
variants: list[GgufVariantInfo] = []
has_vision = False
quant_totals: dict[str, int] = {} # quant -> total bytes
quant_first_file: dict[str, str] = {} # quant -> first filename (display)
for sibling in info.siblings:
fname = sibling.rfilename
if not fname.lower().endswith(".gguf"):
continue
size = sibling.size or 0
# mmproj files are vision projections, not main model files
if "mmproj" in fname.lower():
has_vision = True
continue
# MTP drafters are speculative-decoding companions, not quants.
if _is_mtp_drafter(fname):
continue
quant = _extract_quant_label(fname)
if _is_big_endian_gguf_path(fname, quant):
continue
quant_totals[quant] = quant_totals.get(quant, 0) + size
if quant not in quant_first_file:
quant_first_file[quant] = fname
for quant, total_size in quant_totals.items():
variants.append(
GgufVariantInfo(
filename = quant_first_file[quant],
quant = quant,
size_bytes = total_size,
)
)
# Sort by size descending (largest = best quality first); pinning and OOM
# demotion happen client-side where GPU VRAM info exists.
variants.sort(key = lambda v: -v.size_bytes)
return variants, has_vision
def _resolve_gguf_dir(p: Path) -> Optional[Path]:
"""Resolve a path to the directory containing GGUF variants.
Directory *p* returns directly. A ``.gguf`` file whose parent dir has
model metadata (``config.json`` or ``adapter_config.json``) returns the
parent -- all GGUFs there 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()
or (parent / "export_metadata.json").exists()
):
return parent
return None
def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], bool]:
"""List GGUF quant variants in a local directory.
Like :func:`list_gguf_variants` but reads the filesystem. Aggregates shard
sizes by quant label so split GGUFs appear as one variant.
Returns:
(variants, has_vision): non-mmproj GGUF variants + vision flag.
"""
p = _resolve_gguf_dir(Path(directory))
if p is None:
return [], False
quant_totals: dict[str, int] = {}
quant_first_file: dict[str, str] = {}
has_vision = False
# Recurse so variant-specific subdirs (e.g. ``BF16/...gguf`` used by
# some HF GGUF repos for the largest quants) are picked up. Result
# filenames keep the relative subpath so ``_find_local_gguf_by_variant``
# can locate the file again.
for f in sorted(_iter_gguf_files(p, recursive = True)):
if _is_mmproj(f.name):
has_vision = True
continue
try:
size = f.stat().st_size
except OSError:
size = 0
# Use the relative path so ``BF16/foo.gguf`` and ``Q4_K_M/foo.gguf``
# get distinct quant labels instead of collapsing on basename.
rel = f.relative_to(p).as_posix()
if _is_mtp_drafter(rel):
continue
quant = _extract_quant_label(rel)
if _is_big_endian_gguf_path(rel, quant):
continue
quant_totals[quant] = quant_totals.get(quant, 0) + size
if quant not in quant_first_file:
quant_first_file[quant] = rel
variants = [
GgufVariantInfo(
filename = quant_first_file[q],
quant = q,
size_bytes = s,
)
for q, s in quant_totals.items()
]
variants.sort(key = lambda v: -v.size_bytes)
return variants, has_vision
def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
"""Find the GGUF file in *directory* matching a quantization *variant*.
For sharded GGUFs (multiple files sharing a quant label), returns the
first shard (sorted by name), which is what ``llama-server -m`` expects.
Returns the resolved absolute path, or ``None`` if no match.
"""
p = _resolve_gguf_dir(Path(directory))
if p is None:
return None
# Recurse so variants under a quant-named subdir (e.g.
# ``BF16/foo-BF16-00001-of-00002.gguf``) are found. Match the relative
# path so the quant label can come from the dir name when the basename
# omits it.
matches = []
for f in _iter_gguf_files(p, recursive = True):
rel = f.relative_to(p).as_posix()
if _is_mmproj(f.name) or _is_mtp_drafter(rel):
continue
quant = _extract_quant_label(rel)
if quant != variant or _is_big_endian_gguf_path(rel, quant):
continue
matches.append(f)
matches.sort()
if matches:
return str(matches[0].resolve())
return None
def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]:
"""Best GGUF filename for *repo_id* from the local HF cache, or None.
Excludes mmproj (vision projector) files so a partial cache holding only
the projector cannot route it as the main model.
"""
for snap in _iter_hf_cache_snapshots(repo_id):
rel_files = []
for f in _iter_gguf_files(snap, recursive = True):
rel = f.relative_to(snap).as_posix()
quant = _extract_quant_label(rel)
if _is_mmproj(f.name) or _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
continue
rel_files.append(rel)
if rel_files:
return _pick_best_gguf(rel_files)
return None
def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Optional[str]:
"""Return the best GGUF filename in a HF repo, or None.
Retries (3 attempts, 1s/2s/4s backoff) on transient HF Hub failures: a
silent None would make the caller treat a GGUF-only repo as non-GGUF and
fall through to MLX on Apple Silicon. Offline falls back to the local cache.
"""
import time
from huggingface_hub import model_info as hf_model_info
if _env_offline():
cached = _detect_gguf_from_hf_cache(repo_id)
if cached is not None:
return cached
last_err: Optional[Exception] = None
for attempt in range(3):
try:
info = hf_model_info(repo_id, token = hf_token)
repo_files = []
for sibling in info.siblings:
fname = sibling.rfilename
if not fname.lower().endswith(".gguf"):
continue
quant = _extract_quant_label(fname)
if (
_is_mmproj(fname)
or _is_mtp_drafter(fname)
or _is_big_endian_gguf_path(fname, quant)
):
continue
repo_files.append(fname)
return _pick_best_gguf(repo_files)
except Exception as e:
last_err = e
# 404 / RepoNotFound is permanent -- don't retry
err_name = type(e).__name__
if err_name in (
"RepositoryNotFoundError",
"GatedRepoError",
"RevisionNotFoundError",
"EntryNotFoundError",
):
logger.debug(f"Could not check GGUF files for '{repo_id}': {e}")
return None
if attempt < 2:
time.sleep(2**attempt)
# All attempts failed; fall back to local cache for offline users.
cached = _detect_gguf_from_hf_cache(repo_id)
if cached is not None:
logger.warning(
"HF API unreachable for '%s' (%s); using local cache to detect GGUF.",
repo_id,
type(last_err).__name__ if last_err else "unknown",
)
return cached
logger.warning(f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}")
return None
def download_gguf_file(
repo_id: str,
filename: str,
hf_token: Optional[str] = None,
) -> str:
"""Download a specific GGUF file from a HF repo; returns the local path."""
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
repo_id = repo_id,
filename = filename,
token = hf_token,
)
return local_path
# Cache embedding detection per session to avoid repeated HF API calls
_embedding_detection_cache: Dict[tuple, bool] = {}
def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
"""Detect embedding/sentence-transformer models via HF metadata.
Combines three signals: "sentence-transformers" or "feature-extraction" in
tags, or pipeline_tag in {"sentence-similarity", "feature-extraction"}.
Catches models like gte-modernbert whose library_name is "transformers".
Args:
model_name: Model identifier (HF repo or local path)
hf_token: Optional HF token for gated/private models
Returns:
True if embedding model, else False (default for local paths or errors).
"""
cache_key = (model_name, hf_token)
if cache_key in _embedding_detection_cache:
return _embedding_detection_cache[cache_key]
# Local paths: check for sentence-transformer marker (modules.json)
if is_local_path(model_name):
local_dir = normalize_path(model_name)
is_emb = os.path.isfile(os.path.join(local_dir, "modules.json"))
_embedding_detection_cache[cache_key] = is_emb
return is_emb
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(model_name, token = hf_token)
tags = set(info.tags or [])
pipeline_tag = info.pipeline_tag or ""
is_emb = (
"sentence-transformers" in tags
or "feature-extraction" in tags
or pipeline_tag in ("sentence-similarity", "feature-extraction")
)
_embedding_detection_cache[cache_key] = is_emb
if is_emb:
logger.info(
f"Model {model_name} detected as embedding model: "
f"pipeline_tag={pipeline_tag}, "
f"sentence-transformers in tags={('sentence-transformers' in tags)}, "
f"feature-extraction in tags={('feature-extraction' in tags)}"
)
return is_emb
except Exception as e:
logger.warning(f"Could not determine if {model_name} is embedding model: {e}")
_embedding_detection_cache[cache_key] = False
return False
def _has_model_weight_files(model_dir: Path) -> bool:
"""Return True when a directory contains loadable model weights."""
for item in model_dir.iterdir():
if not item.is_file():
continue
suffix = item.suffix.lower()
if suffix == ".safetensors":
return True
if suffix == ".gguf":
return "mmproj" not in item.name.lower()
if suffix == ".bin":
name = item.name.lower()
if (
name.startswith("pytorch_model")
or name.startswith("model")
or name.startswith("adapter_model")
or name.startswith("consolidated")
):
return True
return False
def _detect_training_output_type(model_dir: Path) -> Optional[str]:
"""Classify a Studio training output as LoRA or full finetune."""
adapter_config = model_dir / "adapter_config.json"
adapter_model = model_dir / "adapter_model.safetensors"
if adapter_config.exists() or adapter_model.exists():
return "lora"
config_file = model_dir / "config.json"
if config_file.exists() and _has_model_weight_files(model_dir):
return "merged"
return None
def _looks_like_lora_adapter(model_dir: Path) -> bool:
return model_dir.is_dir() and (
(model_dir / "adapter_config.json").exists()
or any(model_dir.glob("adapter_model*.safetensors"))
or any(model_dir.glob("adapter_model*.bin"))
)
def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str, str]]:
"""Scan outputs folder for trained Studio models.
Returns:
List of (display_name, model_path, model_type), where model_type is
"lora" for adapter runs or "merged" for full finetunes.
"""
trained_models = []
outputs_path = resolve_output_dir(outputs_dir)
if not outputs_path.exists():
logger.warning(f"Outputs directory not found: {outputs_dir}")
return trained_models
try:
for item in outputs_path.iterdir():
if item.is_dir():
model_type = _detect_training_output_type(item)
if model_type is None:
continue
display_name = item.name
model_path = str(item)
trained_models.append((display_name, model_path, model_type))
logger.debug("Found trained model: %s (%s)", display_name, model_type)
# Sort by mtime, newest first
trained_models.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
logger.info(
"Found %s trained models in %s",
len(trained_models),
outputs_dir,
)
return trained_models
except Exception as e:
logger.error(f"Error scanning outputs folder: {e}")
return []
def scan_exported_models(
exports_dir: str = str(exports_root()),
) -> List[Tuple[str, str, str, Optional[str]]]:
"""Scan exports folder for exported models (merged, LoRA, GGUF).
Supports two layouts: two-level {run}/{checkpoint}/ (merged & LoRA) and
flat {name}-finetune-gguf/ (GGUF).
Returns:
List of (display_name, model_path, export_type, base_model), where
export_type is "lora" | "merged" | "gguf".
"""
results = []
exports_path = resolve_export_dir(exports_dir)
if not exports_path.exists():
return results
try:
for run_dir in exports_path.iterdir():
if not run_dir.is_dir():
continue
# Flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/).
# Skip mmproj (vision projection) files — not loadable as main models.
gguf_files = [f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name)]
if gguf_files:
base_model = None
export_meta = run_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
base_model = meta.get("base_model")
except Exception:
pass
display_name = run_dir.name
model_path = str(gguf_files[0])
results.append((display_name, model_path, "gguf", base_model))
logger.debug(f"Found GGUF export: {display_name}")
continue
# Two-level: {run}/{checkpoint}/
for checkpoint_dir in run_dir.iterdir():
if not checkpoint_dir.is_dir():
continue
adapter_config = checkpoint_dir / "adapter_config.json"
config_file = checkpoint_dir / "config.json"
has_weights = any(checkpoint_dir.glob("*.safetensors")) or any(
checkpoint_dir.glob("*.bin")
)
has_gguf = any(_iter_gguf_files(checkpoint_dir))
base_model = None
export_type = None
if adapter_config.exists():
export_type = "lora"
try:
cfg = json.loads(adapter_config.read_text())
base_model = cfg.get("base_model_name_or_path")
except Exception:
pass
elif config_file.exists() and has_weights:
export_type = "merged"
export_meta = checkpoint_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
base_model = meta.get("base_model")
except Exception:
pass
elif has_gguf:
export_type = "gguf"
gguf_list = list(_iter_gguf_files(checkpoint_dir))
# checkpoint_dir first, then run_dir (export.py writes
# metadata to the top-level export dir)
for meta_dir in (checkpoint_dir, run_dir):
export_meta = meta_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
base_model = meta.get("base_model")
if base_model:
break
except Exception:
pass
display_name = f"{run_dir.name} / {checkpoint_dir.name}"
model_path = str(gguf_list[0]) if gguf_list else str(checkpoint_dir)
results.append((display_name, model_path, export_type, base_model))
logger.debug(f"Found GGUF export: {display_name}")
continue
else:
continue
# Fallback: base model from ./outputs/{run_name}/adapter_config.json
if not base_model:
outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json"
try:
if outputs_adapter_cfg.exists():
cfg = json.loads(outputs_adapter_cfg.read_text())
base_model = cfg.get("base_model_name_or_path")
except Exception:
pass
display_name = f"{run_dir.name} / {checkpoint_dir.name}"
model_path = str(checkpoint_dir)
results.append((display_name, model_path, export_type, base_model))
logger.debug(f"Found exported model: {display_name} ({export_type})")
results.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(results)} exported models in {exports_dir}")
return results
except Exception as e:
logger.error(f"Error scanning exports folder: {e}")
return []
def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]:
"""Read the base model name from a local training or checkpoint directory."""
try:
checkpoint_path_obj = Path(checkpoint_path)
adapter_config_path = checkpoint_path_obj / "adapter_config.json"
if adapter_config_path.exists():
with open(adapter_config_path, "r") as f:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
logger.info("Detected base model from adapter_config.json: %s", base_model)
return base_model
config_path = checkpoint_path_obj / "config.json"
if config_path.exists():
with open(config_path, "r") as f:
config = json.load(f)
for key in ("model_name", "_name_or_path"):
base_model = config.get(key)
if base_model and str(base_model) != str(checkpoint_path_obj):
logger.info(
"Detected base model from config.json (%s): %s",
key,
base_model,
)
return base_model
# TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; re-enable via safe_globals or weights_only=False once threat model allows.
# training_args_path = checkpoint_path_obj / "training_args.bin"
# if training_args_path.exists():
# try:
# import torch
#
# training_args = torch.load(training_args_path)
# if hasattr(training_args, "model_name_or_path"):
# base_model = training_args.model_name_or_path
# logger.info(
# "Detected base model from training_args.bin: %s", base_model
# )
# return base_model
# except Exception as e:
# logger.warning(f"Could not load training_args.bin: {e}")
dir_name = checkpoint_path_obj.name
if dir_name.startswith("unsloth_"):
parts = dir_name.split("_")
if len(parts) >= 2:
model_parts = parts[1:-1]
base_model = "unsloth/" + "_".join(model_parts)
logger.info("Detected base model from directory name: %s", base_model)
return base_model
logger.warning(f"Could not detect base model for checkpoint: {checkpoint_path}")
return None
except Exception as e:
logger.error(f"Error reading base model from checkpoint config: {e}")
return None
def get_base_model_from_lora(lora_path: str) -> Optional[str]:
"""Read the base model name from a LoRA adapter's config, or None."""
try:
lora_path_obj = Path(lora_path)
if not _looks_like_lora_adapter(lora_path_obj):
return None
# adapter_config.json first
adapter_config_path = lora_path_obj / "adapter_config.json"
if adapter_config_path.exists():
with open(adapter_config_path, "r") as f:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
logger.info(f"Detected base model from adapter_config.json: {base_model}")
return base_model
# Fallback: try training_args.bin (requires torch)
# TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an remote code execution sink for third-party LoRAs via this route, re-enable behind a trust check if needed.
# training_args_path = lora_path_obj / "training_args.bin"
# if training_args_path.exists():
# try:
# import torch
#
# training_args = torch.load(training_args_path)
# if hasattr(training_args, "model_name_or_path"):
# base_model = training_args.model_name_or_path
# logger.info(
# f"Detected base model from training_args.bin: {base_model}"
# )
# return base_model
# except Exception as e:
# logger.warning(f"Could not load training_args.bin: {e}")
# Last resort: parse from dir name (unsloth_<model>_<timestamp>)
dir_name = lora_path_obj.name
if dir_name.startswith("unsloth_"):
parts = dir_name.split("_")
if len(parts) >= 2:
model_parts = parts[1:-1] # Skip "unsloth" and timestamp
base_model = "unsloth/" + "_".join(model_parts)
logger.info(f"Detected base model from directory name: {base_model}")
return base_model
logger.warning(f"Could not detect base model for LoRA: {lora_path}")
return None
except Exception as e:
logger.error(f"Error reading base model from LoRA config: {e}")
return None
def get_base_model_from_lora_identifier(
identifier: str, hf_token: Optional[str] = None
) -> Optional[str]:
"""Resolve a LoRA adapter's base model for a LOCAL dir OR a REMOTE HF repo.
``get_base_model_from_lora`` only reads a local adapter directory (it requires
``is_dir()``). The SECURITY gates must also follow a *remote* adapter's base,
because the base model's code / weights are what execute on load: an attacker's
adapter repo can point ``base_model_name_or_path`` at a base carrying a poisoned
pickle or HIGH auto_map code. For a remote repo id we fetch ONLY the small
``adapter_config.json`` (metadata; never a weight file) and read the base. Use
this in the gate paths so a remote LoRA base is scanned, not just the adapter.
Returns the base model id, or ``None`` when the identifier is not a LoRA adapter
or the base cannot be determined (the caller still scans the identifier itself).
A genuine 404 (no ``adapter_config.json`` / repo absent) is distinguished from a
transient error: the latter is retried once, then logged as a WARNING (a missed
base would be scanned by neither gate), so a network blip does not silently and
invisibly skip the base.
"""
# Local path: reuse the existing directory reader (identical behavior).
try:
if is_local_path(identifier):
return get_base_model_from_lora(identifier)
except Exception:
return get_base_model_from_lora(identifier)
# Remote repo id: read base_model_name_or_path from adapter_config.json only.
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
last_exc = None
for _attempt in range(2): # one retry: a transient blip must not skip the base
try:
cfg_path = hf_hub_download(
identifier, "adapter_config.json", token = hf_token if hf_token else None
)
except (EntryNotFoundError, RepositoryNotFoundError):
# No adapter_config.json -> not a resolvable LoRA; caller scans the identifier.
return None
except Exception as exc: # transient / auth / network -> retry once
last_exc = exc
continue
try:
with open(cfg_path, "r") as f:
base_model = json.load(f).get("base_model_name_or_path")
except Exception as exc:
logger.warning("Could not parse adapter_config.json for '%s': %s", identifier, exc)
return None
if base_model:
logger.info(
"Detected base model from remote adapter_config.json (%s): %s",
identifier,
base_model,
)
return base_model # may be None if the key is absent (still a valid answer)
# Both attempts failed transiently: log loudly -- a missed base is gated by neither gate.
logger.warning(
"Could not resolve remote LoRA base for '%s' after retry (%s); its base, if "
"any, will not be added to the security scan targets.",
identifier,
type(last_exc).__name__ if last_exc else "unknown",
)
return None
# Status indicators that appear in UI dropdowns
UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", ""]
def load_model_defaults(model_name: str) -> Dict[str, Any]:
"""Load default training parameters for a model from a YAML file.
Looks in configs/model_defaults/ (incl. subfolders) by model name or its
MODEL_NAME_MAPPING aliases, else falls back to default.yaml. Returns the
parameter dict, or {} if none found.
"""
# No model selected yet (or a non-string id): nothing to load. Guard before
# the .lower() calls below so this doesn't raise and get logged as
# "Error loading model defaults for None: 'NoneType' object has no attribute
# 'lower'".
if not isinstance(model_name, str) or not model_name:
return {}
try:
script_dir = Path(__file__).parent.parent.parent
defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
# Check the mapping first
if model_name.lower() in _REVERSE_MODEL_MAPPING:
canonical_file = _REVERSE_MODEL_MAPPING[model_name.lower()]
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path} (via mapping)")
return config
# For local paths (e.g. /home/.../Spark-TTS-0.5B/LLM from
# adapter_config.json, or C:\Users\...\model on Windows), match the
# last 1-2 path components against the registry (e.g. "Spark-TTS-0.5B/LLM").
_is_local_path = is_local_path(model_name)
# Normalize Windows backslash paths so Path().parts splits correctly
# on POSIX/WSL hosts (pathlib treats backslashes as literals on Linux).
_normalized = normalize_path(model_name) if _is_local_path else model_name
if model_name.lower() not in _REVERSE_MODEL_MAPPING and _is_local_path:
parts = Path(_normalized).parts
for depth in [2, 1]:
if len(parts) >= depth:
suffix = "/".join(parts[-depth:])
if suffix.lower() in _REVERSE_MODEL_MAPPING:
canonical_file = _REVERSE_MODEL_MAPPING[suffix.lower()]
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(
f"Loaded model defaults from {config_path} (via path suffix '{suffix}')"
)
return config
# Exact model name match (backward compatibility). For local paths,
# use only the dir basename to avoid passing absolute paths (e.g.
# C:\...) into rglob, which raises "Non-relative patterns are
# unsupported" on Windows.
_lookup_name = Path(_normalized).name if _is_local_path else model_name
model_filename = _lookup_name.replace("/", "_") + ".yaml"
# Search subfolders and root
for config_path in defaults_dir.rglob(model_filename):
if config_path.is_file():
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path}")
return config
# Fall back to default.yaml
default_config_path = defaults_dir / "default.yaml"
if default_config_path.exists():
with open(default_config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded default model defaults from {default_config_path}")
return config
logger.warning(f"No default config found for model {model_name}")
return {}
except Exception as e:
logger.error(f"Error loading model defaults for {model_name}: {e}")
return {}
@dataclass
class ModelConfig:
"""Configuration for a model to load."""
identifier: str # Clean model identifier (org/name or path)
display_name: str # Original UI display name
path: str # Normalized filesystem path
is_local: bool # Local file vs HF model?
is_cached: bool # Already in HF cache?
is_vision: bool # Vision model?
is_lora: bool # LoRA adapter?
is_gguf: bool = False # GGUF model?
is_audio: bool = False # TTS audio model?
audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
has_audio_input: bool = False # Accepts audio input (ASR/speech understanding)
gguf_file: Optional[str] = None # Full path to the .gguf file (local mode)
gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection)
gguf_mtp_file: Optional[str] = None # Full path to the separate MTP drafter (local mode)
gguf_hf_repo: Optional[str] = (
None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
)
gguf_variant: Optional[str] = None # Quantization variant (e.g. "Q4_K_M")
base_model: Optional[str] = None # Base model (for LoRAs)
@classmethod
def from_lora_path(
cls,
lora_path: str,
hf_token: Optional[str] = None,
) -> Optional["ModelConfig"]:
"""Create ModelConfig from a local LoRA adapter path, auto-detecting the
base model from adapter config.
Args:
lora_path: Path to the LoRA adapter directory
hf_token: HF token for vision detection
"""
try:
lora_path_obj = Path(lora_path)
if not lora_path_obj.exists():
logger.error(f"LoRA path does not exist: {lora_path}")
return None
base_model = get_base_model_from_lora(lora_path)
if not base_model:
logger.error(f"Could not determine base model for LoRA: {lora_path}")
return None
is_vision = is_vision_model(base_model, hf_token = hf_token)
audio_type = detect_audio_type(base_model, hf_token = hf_token)
display_name = lora_path_obj.name
identifier = lora_path # path is the identifier for local LoRAs
return cls(
identifier = identifier,
display_name = display_name,
path = lora_path,
is_local = True,
is_cached = True, # local LoRAs are always cached
is_vision = is_vision,
is_lora = True,
is_audio = audio_type is not None and audio_type != "audio_vlm",
audio_type = audio_type,
has_audio_input = is_audio_input_type(audio_type),
base_model = base_model,
)
except Exception as e:
logger.error(f"Error creating ModelConfig from LoRA path: {e}")
return None
@classmethod
def from_identifier(
cls,
model_id: str,
hf_token: Optional[str] = None,
is_lora: bool = False,
gguf_variant: Optional[str] = None,
) -> Optional["ModelConfig"]:
"""Create ModelConfig from a clean model identifier (HF repo or local
path), for FastAPI routes that send sanitized paths.
Args:
model_id: Clean model identifier (HF repo name or local path)
hf_token: Optional HF token for vision detection on gated models
is_lora: Whether this is a LoRA adapter
gguf_variant: Optional GGUF quant variant (e.g. "Q4_K_M") to load
via -hf for remote repos; None auto-selects via _pick_best_gguf().
Returns:
ModelConfig or None if it cannot be created.
"""
if not model_id or not model_id.strip():
return None
identifier = model_id.strip()
is_local = is_local_path(identifier)
path = normalize_path(identifier) if is_local else identifier
# Add unsloth/ prefix for shorthand HF models
if not is_local and "/" not in identifier:
identifier = f"unsloth/{identifier}"
path = identifier
# Reuse a cached case-variant's exact repo_id spelling to avoid
# one-time re-downloads after #2592.
if not is_local:
resolved_identifier = resolve_cached_repo_id_case(identifier)
if resolved_identifier != identifier:
logger.info(
"Using cached repo_id casing '%s' for requested '%s'",
resolved_identifier,
identifier,
)
identifier = resolved_identifier
path = resolved_identifier
# Auto-detect GGUF models (check before LoRA/vision detection)
if is_local:
if gguf_variant:
gguf_file = _find_local_gguf_by_variant(path, gguf_variant)
else:
gguf_file = detect_gguf_model(path)
if gguf_file:
display_name = Path(gguf_file).stem
logger.info(f"Detected local GGUF model: {gguf_file}")
# Vision: check base model, then look for mmproj
mmproj_file = None
gguf_is_vision = False
gguf_dir = Path(gguf_file).parent
# Is this a vision model, per export metadata?
base_is_vision = False
meta_path = gguf_dir / "export_metadata.json"
if meta_path.exists():
try:
meta = json.loads(meta_path.read_text())
base = meta.get("base_model")
if base and is_vision_model(base, hf_token = hf_token):
base_is_vision = True
logger.info(f"GGUF base model '{base}' is a vision model")
except Exception as e:
logger.debug(f"Could not read export metadata: {e}")
# Direct file selections may point into a quant subdir while
# mmproj-*.gguf lives at the snapshot root.
companion_root = _local_gguf_companion_search_root(path, gguf_file)
mmproj_file = detect_mmproj_file(gguf_file, search_root = companion_root)
if mmproj_file:
gguf_is_vision = True
logger.info(f"Detected mmproj for vision: {mmproj_file}")
elif base_is_vision:
logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}")
# Separate MTP drafter sibling (Gemma 4), mirroring mmproj.
mtp_file = detect_mtp_file(gguf_file, search_root = companion_root)
if mtp_file:
logger.info(f"Detected MTP drafter: {mtp_file}")
return cls(
identifier = identifier,
display_name = display_name,
path = path,
is_local = True,
is_cached = True,
is_vision = gguf_is_vision,
is_lora = False,
is_gguf = True,
gguf_file = gguf_file,
gguf_mmproj_file = mmproj_file,
gguf_mtp_file = mtp_file,
)
else:
# Does the HF repo contain GGUF files?
gguf_filename = detect_gguf_model_remote(identifier, hf_token = hf_token)
if gguf_filename:
# Preflight: verify llama-server binary exists before a multi-GB
# download. include_denied: a transiently locked binary still
# exists (the lock clears long before the download finishes; the
# load itself reports a still-locked binary distinctly).
from core.inference.llama_cpp import (
LLAMA_SERVER_NOT_FOUND_DETAIL,
LlamaCppBackend,
LlamaServerNotFoundError,
)
if not LlamaCppBackend._find_llama_server_binary(include_denied = True):
raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL)
# list_gguf_variants() detects vision & resolves the variant
variants, has_vision = list_gguf_variants(identifier, hf_token = hf_token)
variant = gguf_variant
if not variant: # auto-select best quant
variant_filenames = [v.filename for v in variants]
best = _pick_best_gguf(variant_filenames)
if best:
variant = _extract_quant_label(best)
else:
variant = "Q4_K_M" # Fallback — llama-server's own default
display_name = f"{identifier.split('/')[-1]} ({variant})"
logger.info(
f"Detected remote GGUF repo '{identifier}', "
f"variant={variant}, vision={has_vision}"
)
return cls(
identifier = identifier,
display_name = display_name,
path = identifier,
is_local = False,
is_cached = False,
is_vision = has_vision,
is_lora = False,
is_gguf = True,
gguf_file = None,
gguf_hf_repo = identifier,
gguf_variant = variant,
)
# Auto-detect LoRA for local paths (adapter_config.json on disk)
if not is_lora and is_local:
detected_base = (
get_base_model_from_lora(path) if _looks_like_lora_adapter(Path(path)) else None
)
if detected_base:
is_lora = True
logger.info(f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})")
# Auto-detect LoRA for remote HF models. When offline, huggingface_hub
# raises OfflineModeIsEnabled in ~0ms; we fall through to the cache.
if not is_lora and not is_local:
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(identifier, token = hf_token)
repo_files = [s.rfilename for s in info.siblings]
if "adapter_config.json" in repo_files:
is_lora = True
logger.info(f"Auto-detected remote LoRA adapter: '{identifier}'")
except Exception as e:
logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}")
# API may have failed; adapter_config.json could still be cached.
if not is_lora:
for snap in _iter_hf_cache_snapshots(identifier):
if (snap / "adapter_config.json").is_file():
is_lora = True
logger.info(f"Auto-detected cached LoRA adapter: '{identifier}'")
break
# Handle LoRA adapters
base_model = None
if is_lora:
if is_local:
# Local LoRA: read adapter_config.json from disk
base_model = get_base_model_from_lora(path)
else:
# Remote LoRA: fetch adapter_config.json from HF
try:
from huggingface_hub import hf_hub_download
config_path = hf_hub_download(identifier, "adapter_config.json", token = hf_token)
with open(config_path, "r") as f:
adapter_config = json.load(f)
base_model = adapter_config.get("base_model_name_or_path")
if base_model:
logger.info(f"Resolved remote LoRA base model: '{base_model}'")
except Exception as e:
logger.warning(
f"Could not download adapter_config.json for '{identifier}': {e}"
)
if not base_model:
logger.warning(f"Could not determine base model for LoRA '{path}'")
return None
check_model = base_model
else:
check_model = identifier
vision = is_vision_model(check_model, hf_token = hf_token)
audio_type_val = detect_audio_type(check_model, hf_token = hf_token)
has_audio_in = is_audio_input_type(audio_type_val)
display_name = Path(path).name if is_local else identifier.split("/")[-1]
return cls(
identifier = identifier,
display_name = display_name,
path = path,
is_local = is_local,
is_cached = is_model_cached(identifier) if not is_local else True,
is_vision = vision,
is_lora = is_lora,
is_audio = audio_type_val is not None and audio_type_val != "audio_vlm",
audio_type = audio_type_val,
has_audio_input = has_audio_in,
base_model = base_model,
)
@classmethod
def from_ui_selection(
cls,
dropdown_value: Optional[str],
search_value: Optional[str],
local_models: list = None,
hf_token: Optional[str] = None,
is_lora: bool = False,
) -> Optional["ModelConfig"]:
"""Create a ModelConfig from UI dropdown/search selections (base models and LoRAs)."""
selected = None
if search_value and search_value.strip():
selected = search_value.strip()
elif dropdown_value:
selected = dropdown_value
if not selected:
return None
display_name = selected
# Resolve display names via the 'local_models' parameter
if " (Active)" in selected or " (Ready)" in selected:
clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "")
if local_models:
for local_display, local_path in local_models:
if local_display == clean_display_name:
selected = local_path
break
# Strip all UI status indicators to get the final identifier
identifier = selected
for status in UI_STATUS_INDICATORS:
identifier = identifier.replace(status, "")
identifier = identifier.strip()
is_local = is_local_path(identifier)
path = normalize_path(identifier) if is_local else identifier
# Add unsloth/ prefix for shorthand HF models
if not is_local and "/" not in identifier:
identifier = f"unsloth/{identifier}"
path = identifier
if not is_local:
resolved_identifier = resolve_cached_repo_id_case(identifier)
if resolved_identifier != identifier:
identifier = resolved_identifier
path = resolved_identifier
# Keep existing local GGUF selections on the llama-server path. This
# constructor is still used by older inference helpers and must not
# describe a .gguf weight file as loadable by FastVisionModel.
if is_local and not is_lora and detect_gguf_model(path):
gguf_config = cls.from_identifier(path, hf_token = hf_token)
if gguf_config is not None:
gguf_config.display_name = display_name
return gguf_config
# --- Base Model and Vision Detection ---
base_model = None
is_vision = False
if is_lora:
# A LoRA MUST have a base model.
base_model = get_base_model_from_lora(path)
if not base_model:
logger.warning(
f"Could not determine base model for LoRA '{path}'. Cannot create config."
)
return None # cannot proceed without a base model
# A LoRA's vision capability comes from its base model.
is_vision = is_vision_model(base_model, hf_token = hf_token)
else:
# Base model: check its own vision status.
is_vision = is_vision_model(identifier, hf_token = hf_token)
from utils.paths import is_model_cached
is_cached = is_model_cached(identifier) if not is_local else True
return cls(
identifier = identifier,
display_name = display_name,
path = path,
is_local = is_local,
is_cached = is_cached,
is_vision = is_vision,
is_lora = is_lora,
base_model = base_model, # None for base models, set for LoRAs
)