unsloth/studio/backend/utils/models/model_config.py
Daniel Han 815f242970
Studio: offer the latest transformers release for brand-new architectures (#7056)
* Studio: offer the latest transformers release for brand-new architectures

When a model's config.json model_type is absent from every installed
transformers overlay (base 4.57.x and the .venv_t5_530/550/510 sidecars),
Studio now checks, unauthenticated and cached, whether the newest
transformers ships it:

- utils/transformers_latest.py fetches the latest release version from
  https://pypi.org/pypi/transformers/json and the CONFIG_MAPPING_NAMES
  sources for that tag and for main from raw.githubusercontent.com
  (never api.github.com), parsing them with the same AST extractor the
  static router uses (no code execution, no trust_remote_code). Results
  are cached in memory and in a JSON snapshot under studio_root()/cache
  with a one day ttl; fetches are bounded to 5s with one retry and a
  failure backoff, and offline mode or the new kill switch
  UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS=1 short-circuits to None.

- POST /api/inference/validate gains requires_transformers_upgrade plus
  a transformers_upgrade payload (model_type, pypi_version,
  supported_in_pypi, supported_in_main) so the frontend can raise the
  install consent dialog before /load, mirroring the existing
  remote-code consent flow. The check fires only when the model_type is
  unknown to all installed overlays and the hardcoded tier tables.

- POST /api/inference/install-latest-transformers provisions a new
  persistent .venv_t5_latest sidecar after user consent, pinned to the
  exact PyPI version (re-verified server-side) with the same
  --target/--no-deps recipe as the fixed sidecars. A JSON pin marker
  inside the dir records the installed package set, so restarts
  revalidate it and routing resolves the new highest-ranked tier
  automatically. A dependency preflight (compat_plan) compares the
  release's requires_dist against the running env: unsatisfied
  tokenizers/safetensors floors are shadow-installed as exact pins into
  the sidecar, anything else unsatisfied blocks the install with a
  clear message.

Routing for every already-supported model_type is unchanged: the
hardcoded lists and the 530/550/510 static resolver run first, the new
tier only participates once its venv exists, and the probe order gains
the latest sidecar only when provisioned. Verified against live PyPI
and GitHub (transformers 5.13.0: 674 model_types, 26 absent from all
installed overlays, e.g. cosmos3_omni; 4 dev-only on main) and with a
real sidecar install plus restart persistence. 64 new tests; the
existing 200-test transformers_version suite passes unchanged.

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

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

* Latest-transformers check: fetch outside the lock, serialize installs

Release the module lock during the network refresh so a slow fetch cannot
stall other threads in the ASGI pool; concurrent callers during a fetch get
None (the graceful fallthrough) via an in-flight flag instead of stacking
fetches. Serialize install_latest_transformers with an in-progress flag so
concurrent consents cannot race the sidecar delete and recreate; the loser
gets a structured already-in-progress refusal.

* Latest-transformers check: LoRA bases, pin-gated mapping, live reverify

Run the upgrade check over the [adapter, base] target set so a LoRA whose
base model is a brand-new architecture surfaces the prompt (the worker
activates transformers for the base, not the adapter).

Gate the latest overlay's mapping lookup on a valid pin marker, matching
activation and the probe order, so a partial or manual .venv_t5_latest dir
cannot be routed to and then refused at activation.

Re-verify the requested version against a live PyPI snapshot at install
time, falling back to the cached one on fetch failure, so a release
published inside the cache TTL is not silently missed.

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

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

* Latest-transformers check: nested config types and latest-tier vision probe

Collect every model_type in the config (top level plus each nested
sub-config) and signal on the first one missing from all installed
overlays, so a supported wrapper carrying a brand-new backbone still
surfaces the upgrade prompt; wrappers instantiate sub-configs through
CONFIG_MAPPING and would fail on the nested type.

Route the vision capability subprocess through the pinned latest sidecar
when the model resolves to the latest tier, so latest-only VLMs are not
misclassified as text-only; every other tier keeps the 5.5 sidecar used
today.

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

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

* Latest tier: nested routing, vision probe after raw miss, safe upgrades

Route by every model_type in the config: a nested sub-config type can raise
the tier (wrappers instantiate sub-configs through CONFIG_MAPPING), so a
supported wrapper with a latest-only backbone routes to latest once
installed instead of staying on default. An unknown nested type never
vetoes; the primary type keeps its previous semantics. The collector is
shared with the upgrade checker.

Vision detection: when the raw heuristics say False for a model that routes
to the latest tier, run the AutoConfig subprocess under the pinned latest
sidecar instead of trusting heuristics built from older transformers.

Provisioning: stage-and-swap. Build the new sidecar in .venv_t5_latest.staging
and swap it in only when the install and pin marker are complete, so a failed
upgrade never destroys a previously working sidecar; restore the old dir if
the final swap fails.

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

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

* Latest-transformers checker, vision subprocess, and cache fixes

Require the latest release to support every missing model_type (the
primary included) before prompting; a nested-only match cannot make the
model loadable, so no install is offered for it.

The vision-check subprocess now unions the active sidecar's own
registry mappings into the inlined parent-process detection sets, so
architectures only the sidecar knows classify correctly.

A successful sidecar install clears the tier probe cache, the latest
tier's model_type mapping, and the vision-detection cache so the new
venv takes effect without a restart. Tests for all three.

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

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

* Aggregate upgrade support flags and keep install off /v1

The upgrade signal now reports supported_in_pypi only when the latest
release covers every missing model_type; a mix with a main-only nested
type surfaces as dev-only so no PyPI install is offered that would
still fail at load. The consented install endpoint moves to
studio_router so it is not reachable through the OpenAI-compatible /v1
mount. Tests for both.

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

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

* Honor the latest-transformers kill switch in routing

With UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS set after the sidecar was
provisioned, the latest tier still joined mapping and probe routing
because only the pin was checked. Both admission points now also check
the kill switch, so operators can roll back a problematic sidecar
without deleting files.

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

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

* Repair the latest sidecar through stage-and-swap

The lazy repair path installed into the live .venv_t5_latest, which
_ensure_venv_dir wipes first, so a failed repair deleted the pinned
sidecar and its marker. Both the consented install and the repair now
share one stage-and-swap helper: the incomplete-but-pinned dir survives
any failure and a later attempt can still repair it.

* Tighten comments

* Remove the staging dir when a latest-sidecar install fails

A pip failure inside _ensure_venv_dir returns False without raising, so
the except cleanup never ran and the partial .venv_t5_latest.staging
leaked until a later attempt. Also note on the validate response fields
that frontend consumption ships in the follow-up PR.

* Add the transformers-upgrade consent dialog to the frontend

When /validate reports requires_transformers_upgrade, every explicit load
path (chat runtime and the compare composer) now pauses on a consent
dialog modeled on the remote-code one: it names the model_type and the
latest PyPI transformers version, and on Accept calls
/api/inference/install-latest-transformers itself, shows an installing
state, and resumes the original load automatically on success. Errors
surface in the dialog with a retry; Cancel aborts the load like the
trust dialog's deny path. Architectures shipped only on transformers
main get a dev-only notice with no install button. Background auto-load
skips upgrade-requiring candidates instead of prompting, mirroring the
trust_remote_code rule. The dialog mounts once in the root layout and
runs before the security dialogs, since no load can proceed without the
runtime.

* Route a non-installable new architecture to the custom-code consent as a last resort

When the upgrade dialog has no installable PyPI release (the architecture
is only on transformers main, which Studio never installs), the dialog now
says so explicitly, and when the model also declares custom (auto_map)
code it offers Continue with custom code: resolving the paused load into
the existing trust_remote_code consent gate instead of hard-aborting.
Models with no custom code keep the Cancel-only notice. The backend
returns no upgrade signal at all for architectures unknown to both PyPI
and main, so those still route straight to the unchanged security gate.

* Force a 16-bit load for models on the latest-transformers sidecar

Live validation with Zyphra/ZAYA1-8B (model_type zaya, shipped by
transformers 5.13.1 but unknown to every installed tier) surfaced a
generation crash when the consented sidecar load kept the default bnb
4-bit quantization: transformers' grouped-MoE kernels feed the packed
uint8 expert weights straight into torch._grouped_mm, and generation
dies (plain 16-bit works). New latest_tier_active_for() mirrors the
sidecar activation's tier resolution and never raises; the inference
worker flips load_in_4bit off when it reports true, and the load route
applies the same flip so the pre-load VRAM guard and the worker command
agree. Fixed tiers are untouched. With the guard, ZAYA1-8B loads and
generates correctly in Studio chat.

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

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

* Offer the custom-code fallback when a latest-sidecar install fails

* Fail remote mapping fetches wholesale and mirror the 16-bit flip in validate

A transient fetch or parse failure of one auto-mapping file no longer caches
a partial latest-release map for the TTL (a real 404 on pre-5.10 tags is
still tolerated), and validate_model now applies the same latest-sidecar
16-bit sizing flip as /load before the training guard so the two agree.

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

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

* Tighten comments in the latest-transformers changes

* Resolve remote LoRA bases, fold nested tiers, and guard the sidecar swap

latest_tier_active_for now resolves a remote adapter's base model the same
way worker pre-activation does (and returns early without a sidecar pin), a
hardcoded fast-path tier is raised when a nested sub-config's model_type
needs a higher sidecar, and the install route refuses to swap .venv_t5_latest
while training runs on it and unloads a latest-tier chat model first.

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

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

* Gate the sidecar install on worker liveness and size installable upgrades 16-bit

The install route now refuses while any training or export runs (tier
re-resolution without the load token is unreliable for gated repos), holds
the inference lifecycle gate across the unload and the swap so no load can
interleave, and passes the model name to unload_model. validate_model runs
the upgrade check before the training guard and sizes an installable
upgrade as 16-bit, matching what /load and the worker will force after the
consented install.

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

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

* Close the sidecar install races and honor the kill switch over cached mappings

Training starts and mutating export routes now refuse while a transformers
install is in progress (shared is_install_in_progress flag), the chat unload
and idle export-worker teardown moved into a before_swap hook that runs only
once the staged install succeeded, and _config_model_types checks the kill
switch before returning a cached latest mapping.

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

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

* Reserve the sidecar swap before the gate wait and abort it on failed teardown

The install-in-progress flag moved into a shared sidecar swap reservation in
transformers_version, taken by the install route before awaiting the
inference lifecycle gate (so training and export starts see it for the whole
window) and by the lazy .venv_t5_latest repair path. The before_swap hook
now raises when the chat unload or export teardown reports failure, leaving
the previous sidecar untouched.

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

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

* Back the sidecar swap reservation with a cross-process lock file

The lazy repair runs inside worker subprocesses, where a module-level flag
is invisible to the parent's route checks. The reservation now also creates
a lock file next to .venv_t5_latest (O_EXCL, owner-only removal, stale after
two hours for crashed owners), so is_install_in_progress sees a repair from
any Studio process.

* Hand the swap reservation to the installer thread and harden pre-swap teardown

A cancelled install request no longer releases the reservation while the
installer thread is still staging (the thread owns and releases it, shielded
from cancellation). The route refuses while another inference request is
generating, export teardown runs before the chat unload and is judged by
worker liveness rather than the cleanup return value, and a live inference
worker with no active model (failed load residue) is shut down before the
swap.

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

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

* Keep the lifecycle gate with the installer and recheck the swap at spawn time

The gate moved into the shielded install task so a cancelled POST cannot
release the guard /load honors while the installer still runs, cached latest
probe results are ignored while the kill switch is set, and the training and
export subprocess spawns recheck the sidecar swap reservation right before
spawning (the route-level guards are one-shot and validation can outlast an
install's start).

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

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

* Close the spawn-registration windows against the sidecar install

Training marks the spawn in progress before its reservation recheck and
is_training_active honors the flag, so the install route sees a start that
has passed proc.start() but not yet recorded _proc. Export load-checkpoint
rechecks the reservation after setting _export_active and before tearing
down the old worker, so losing the race keeps the loaded checkpoint instead
of surfacing a 500.

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

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

* Refine the install-window interleavings around worker teardown

The inference busy count is rechecked under the lifecycle gate (streams
start by taking that gate, so nothing slips past a held gate), the training
handshake moved ahead of the VRAM-freeing before_spawn hook so a lost race
leaves chat/export intact, the export spawn-time check is op-aware (inside
an active op the install is the side that aborts), and the Xet-stall respawn
waits out a transient reservation instead of stranding the run.

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

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

* Track the install's server-side unload and guard export ops against the swap

The upgrade dialog store records when its install actually ran (the server
unloads the active chat model before swapping), and the load flow then marks
the previous model as unloaded so a later cancelled gate still triggers
rollback; the custom-code fallback leaves the flag unset. _run_export gained
the same reservation handshake as load_checkpoint so an install cannot block
behind an hours-long export op instead of returning 409.

* Tighten comments in the install-guard and upgrade-consent changes

* Surface install-race refusals cleanly and roll back after a failed swap unload

/load refuses while the sidecar swap is reserved so a load cannot succeed
and immediately be unloaded by the pre-swap teardown, worker starts that
lose the install race raise a typed SidecarSwapInProgress mapped to 409
instead of a 500, the install response reports model_unloaded even on a
structured failure so the client can restore its state, and the compare
flow tracks the server-side unload like the primary load path and clears a
stale checkpoint on abort.

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

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

* Type the export install races, scope the lock release, and keep the unload signal

Export load-checkpoint and export ops raise SidecarSwapInProgress (mapped to
409 in every export route) instead of a 400-shaped failure, the export spawn
check distinguishes repair reservations (always refused) from install ones
(op-aware), the swap lock release only unlinks a lock this process wrote so
a stale-superseded owner cannot drop the new owner's live lock, and the
frontend unload signal survives a superseding consent via read-and-clear
consumption instead of a reset.

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

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

* Finalize a stalled run when the respawn loses the install race and latch the unload signal

The Xet-stall respawn timeout now finalizes the run as a failure instead of
raising into the pump's broad finalization catch (which stranded it in a
training state with no worker), and a successful install retry ORs the
model_unloaded signal with the latched value so a failed-after-unload first
attempt still triggers rollback.

* Recheck the swap under the load gate and latch the unload before resolver checks

/load rechecks the sidecar reservation after acquiring the lifecycle gate
(an install can reserve while the load queues on it), and the dialog store
latches model_unloaded as soon as the install response arrives, before any
resolver-identity guard, so a superseded consent's unload still reaches
whichever load consumes the signal next.

* Report cleared-state unload failures, guard queued installs, and fold name tiers

A failed chat unload that still cleared the orchestrator's model state now
reports model_unloaded so the client rolls back, the installer aborts with
a 409 when a model load completed while it waited on the lifecycle gate,
and the fixed-tier name fast path consults the config mapping when a latest
sidecar is pinned so an accepted upgrade routes to the sidecar it installed
(no I/O added to the unpinned path).

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

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

* Report cleared-state unload failures and harden the spawn handshake flag

The failed-unload branch in before_swap now detects that the orchestrator
cleared its model state and reports model_unloaded before aborting (the
earlier commit claimed this fix but a scripting error dropped the edit),
the installer's queued-load check compares a load generation counter so a
same-model reload is caught, and both training spawn sites wrap everything
after the handshake in a guard that resets _spawn_in_progress on any
exception so a failed start cannot wedge is_training_active.

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

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

* Bump the load generation when the load is published, not at load start

A start-time bump is already visible when the installer snapshots mid-load,
so a same-model reload completing after the snapshot looked unchanged and
could be unloaded by the swap. The counter now increments alongside the
active_model_name publish.

* Self-heal a broken pinned sidecar, guard lazy repairs, and refresh stale retries

A valid pin whose transformers source dir vanished now triggers the repair
from the routing path (with a five minute backoff after failures) instead of
silently routing latest-only models to older tiers, the lazy repair refuses
while parent-visible chat/training/export workers are active since it has no
teardown of its own, and a version-mismatch install failure carries the
superseding release so the dialog's Retry re-requests a version that can
succeed.

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

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

* Flip latest-tier loads to 16-bit outside chat and protect export state

Training and export workers now apply the same latest-sidecar 16-bit flip
as the chat worker so a brand-new grouped-MoE architecture cannot reach bnb
4-bit through those paths, the latest-tier vision override returns None on
an inconclusive probe so a transient failure is not cached as not-vision,
and the install route refuses while an idle export checkpoint is loaded
rather than discard it with no rollback signal on a failed swap.

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

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

* Address parallel-review findings on the sidecar guards and install checks

The training route sizes latest-tier jobs 16-bit before GPU selection, the
inference subprocess spawn rechecks the swap reservation like training and
export (covering the OpenAI auto-switch path) with the typed error mapped
to a retryable 409, compat_plan blocks the install when dependency metadata
cannot be fetched instead of proceeding unverified, snapshot model-type
lists must contain only strings, and pin-marker package specs are validated
against the sidecar's own package set before ever reaching pip.

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

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

* Parent-only repairs, live-owner locks, remote-base activation, pre-teardown recheck

Lazy sidecar repairs now refuse inside worker children (whose empty backend
singletons cannot see live siblings) and run only in the parent where the
active-worker guard is real, swap-lock staleness requires the owner pid to
be dead so a slow live install is never superseded, both activation entry
points resolve a remote adapter's base model like the inference worker and
latest_tier_active_for already do, and load_model rechecks the reservation
before tearing down the old worker so losing the race keeps the current
model loaded.

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

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

* Check workers under the repair reservation and keep state on refused swaps

The lazy repair now reserves first and checks workers under the reservation
(worker starts set their active markers before rechecking, so every
interleaving aborts one side), with export ops and in-flight inference loads
counted as active. The inference pre-teardown and spawn guards refuse only
repair reservations since an install shares the load's lifecycle gate and
aborts via its queued-load snapshot, a SidecarSwapInProgress raised before
teardown no longer clears the live model mirrors, and an export spawn abort
after teardown clears current_checkpoint so the page cannot claim a loaded
checkpoint with no worker.

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

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

* Repair a present-but-incomplete latest sidecar from routing

The routing self-heal only fired when the pinned sidecar's transformers/
dir was missing. A sidecar that kept transformers/ but lost another pinned
package still routed models to the latest tier, and workers refuse
parent-only repairs, so every load failed until a manual reinstall. Routing
now validates the full pin (via _venv_dir_is_valid) and repairs any
incomplete sidecar under the same swap reservation and 5-minute backoff.

* Treat an unrepaired latest sidecar as unavailable in routing

When the pinned sidecar is incomplete and the lazy repair fails (offline,
pip failure, workers active) or is inside the backoff window, routing
returned the source dir anyway, sending models to a tier whose worker
activation is known to fail. Return None instead so models an older tier
supports keep loading there until a repair succeeds, matching the behavior
when the sidecar dir is missing entirely.

* Harden sidecar swap and repair against crash, survivor, and 16-bit paths

Reclaim a swap lock as soon as its recorded owner PID is dead instead of
waiting out the two-hour cutoff, so a crash mid-install no longer wedges
/load, training, export, and repair for hours. A lock whose PID cannot be
read yet still uses the long cutoff so the create-before-write window is
never mistaken for dead.

Probe process liveness with OpenProcess on Windows: os.kill(pid, 0) there
is CTRL_C_EVENT (a real Ctrl+C via GenerateConsoleCtrlEvent), not a
harmless check, and psutil is not always present.

Return whether _shutdown_subprocess actually killed the worker and keep the
live handle when it survives terminate/kill (an uninterruptible CUDA
syscall can outlive SIGKILL). The pre-swap liveness guard now trusts that
result, so the destructive .venv_t5_latest rename cannot proceed while a
live worker still holds sidecar modules.

Recover a sidecar stranded at .old when a swap's activation rename and its
rollback both fail: reading the pin restores it when no swap holds the
reservation, so latest-tier models are not permanently broken.

Resolve the latest tier in the parent for export loads and for explicitly
16-bit training runs, not only 4-bit ones: tier resolution self-heals an
incomplete sidecar, and repairs are parent-only, so those paths could not
recover before. Sidecar integrity and quantization are independent.

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

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

* Revert the parent-side latest-tier repair probe on training and export loads

The probe ran before the route freed VRAM, so a resident chat or export worker
made _workers_active_for_repair() refuse the parent-only repair; the route then
tore that worker down and spawned a child that also cannot repair, so an
incomplete sidecar still failed to load. Repairing correctly requires running the
repair between the worker teardown and the child spawn, decoupled from VRAM
sizing, which is a larger change tracked separately. Restore the prior behavior
so these paths match the reviewed form and do not partially attempt a repair that
cannot complete while workers are resident.

* Honor failed worker shutdowns on load and revalidate the cached latest mapping

The fresh-load paths spawned a new worker straight after _shutdown_subprocess
without checking its result, so a worker that outlived terminate/kill (a wedged
CUDA syscall) had its handle overwritten by the replacement while it still held
GPU memory, and is_worker_alive/the pre-swap guard could no longer see it. Both
the inference load and the export checkpoint load now abort when the old worker
did not exit, so the load can be retried once it does.

_config_model_types returned a cached latest mapping without re-checking the
sidecar, so a sidecar deleted or broken in-process after its first parse was
never re-validated: routing kept sending latest-only models to the stale latest
tier while activation failed. The cached latest mapping is now dropped and
re-resolved (self-healing) when the sidecar is no longer intact.

* Drop cached latest mapping when the pin is gone; keep 4-bit for custom-code fallback

_latest_sidecar_intact now returns False when the pin marker itself is gone, not
just when a pinned package is missing. Otherwise a cached latest mapping outlived
a deleted pin: _config_model_types kept returning it, so routing sent latest-only
models to a tier whose worker activation then failed (no pinned version) until
restart. It now drops the cache and re-resolves to no latest tier. The
_overlay_transformers_dir caller already gates on a present pin, so it is
unaffected.

validate_model forced 16-bit sizing whenever a PyPI upgrade was merely offered,
even for a model that can fall back to its own auto_map code. /load loads such a
model 4-bit without the install, and the install route refuses while training is
active, so 16-bit sizing here returned a VRAM 409 for the only viable 4-bit path.
The offered-upgrade flip is now gated on the absence of a custom-code fallback;
an already-active latest sidecar still always sizes 16-bit.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-15 05:25:26 -07:00

2928 lines
110 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
# Union the ACTIVE sidecar's registry into the inlined parent-process sets
# so architectures only the sidecar knows still classify correctly.
try:
from transformers.models.auto import modeling_auto as _ma
for _attr in ("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES",
"MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES"):
_d = dict(getattr(_ma, _attr, None) or {})
_VLM_MODEL_TYPES |= set(_d)
_VLM_CLASS_NAMES |= set(_d.values())
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_ONLY_MODEL_TYPES |= set(dict(getattr(_ma, _attr, None) or {}))
except Exception:
pass
# 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 ""
# Latest-only architectures need the latest sidecar for AutoConfig;
# other tiers keep the 5.5 sidecar.
sidecar_dir = _VENV_T5_DIR
try:
from utils.transformers_version import _VENV_T5_LATEST_DIR, get_transformers_tier
if get_transformers_tier(model_name, hf_token, probe = False) == "latest":
sidecar_dir = _VENV_T5_LATEST_DIR
except Exception:
pass
try:
result = subprocess.run(
[
sys.executable,
"-c",
_VISION_CHECK_SCRIPT,
sidecar_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:
if raw is False and not local_files_only:
# Raw heuristics predate latest-only architectures; on the latest tier,
# trust that sidecar's AutoConfig probe over the heuristic False. An
# inconclusive probe (sidecar mid-repair, timeout) is transient: return
# None so the heuristic False is not cached and the model is re-probed.
try:
from utils.transformers_version import get_transformers_tier
if get_transformers_tier(model_name, hf_token, probe = False) == "latest":
return _is_vision_model_subprocess(model_name, hf_token = hf_token)
except Exception:
pass
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_dirs: list[Path] = []
try:
if not cache_dir.is_dir():
return
for entry in cache_dir.iterdir():
if entry.is_dir() and entry.name.lower() == target:
repo_dirs.append(entry)
except OSError:
return
if not repo_dirs:
return
snap_dirs: list[Path] = []
for repo_dir in repo_dirs:
snapshots = repo_dir / "snapshots"
try:
if snapshots.is_dir():
for snap_dir in snapshots.iterdir():
try:
if snap_dir.is_dir():
snap_dirs.append(snap_dir)
except OSError:
continue
except OSError:
continue
if not snap_dirs:
return
snap_dirs_with_mtime = []
for snap_dir in snap_dirs:
try:
snap_dirs_with_mtime.append((snap_dir.stat().st_mtime, snap_dir))
except OSError:
continue
snap_dirs_with_mtime.sort(key = lambda item: item[0], reverse = True)
yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime)
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.
A newer snapshot can hold only a companion file (for example a vision
projector fetched on demand) while the quant files live in an older
snapshot. Returning the first snapshot that merely reports a vision flag
would shadow those real variants, so keep scanning older snapshots for
actual variants and carry the vision flag across snapshots.
"""
any_vision = False
for snap in _iter_hf_cache_snapshots(repo_id):
variants, has_vision = list_local_gguf_variants(str(snap))
any_vision = any_vision or has_vision
if variants:
return variants, any_vision
if any_vision:
return [], True
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
)