wehub-resource-sync
1383 lines
50 KiB
Python
1383 lines
50 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Model discovery for oMLX multi-model serving.
|
|
|
|
This module scans a model directory and discovers available models,
|
|
estimating memory usage for each.
|
|
|
|
Supports:
|
|
- LLM models: Use BatchedEngine for continuous batching with paged KV cache
|
|
- VLM models: Use VLMBatchedEngine for vision-language model inference
|
|
- Embedding models: Use EmbeddingEngine for batch embedding generation
|
|
- Reranker models: Use RerankerEngine for document reranking
|
|
- Audio STT models: Use STTEngine for speech-to-text (Whisper, Qwen3-ASR, ...)
|
|
- Audio TTS models: Use TTSEngine for text-to-speech (Qwen3-TTS, Kokoro, ...)
|
|
"""
|
|
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import re
|
|
import tempfile
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ModelType = Literal["llm", "vlm", "embedding", "reranker", "audio_stt", "audio_tts", "audio_sts"]
|
|
EngineType = Literal["batched", "vlm", "embedding", "reranker", "audio_stt", "audio_tts", "audio_sts"]
|
|
|
|
# Known VLM (Vision-Language Model) types from mlx-vlm
|
|
VLM_MODEL_TYPES = {
|
|
"qwen2_vl",
|
|
"qwen2_5_vl",
|
|
"qwen3_vl",
|
|
"qwen3_vl_moe",
|
|
"qwen3_5_moe",
|
|
"gemma3",
|
|
"gemma4",
|
|
"gemma4_unified",
|
|
"diffusion_gemma",
|
|
"llava",
|
|
"llava_next",
|
|
"llava-qwen2",
|
|
"llava_qwen2", # underscore form — matches FastVLM checkpoints on disk
|
|
"mllama",
|
|
"idefics3",
|
|
"internvl_chat",
|
|
"phi3_v",
|
|
"paligemma",
|
|
"mistral3",
|
|
"pixtral",
|
|
"molmo",
|
|
"molmo2",
|
|
"bunny_llama",
|
|
"multi_modality",
|
|
"florence2",
|
|
"deepseekocr",
|
|
"deepseekocr_2",
|
|
"dots_ocr",
|
|
"glm_ocr",
|
|
"minimax_m3_vl",
|
|
"minicpmv",
|
|
"phi4_siglip",
|
|
"phi4mm",
|
|
"youtu_vl",
|
|
}
|
|
|
|
# Text-only model families that are implemented in mlx-vlm rather than
|
|
# mlx-lm. They still use the VLM engine because that path loads mlx-vlm
|
|
# models and adapts their language model to oMLX's scheduler.
|
|
VLM_NATIVE_TEXT_MODEL_TYPES = {
|
|
"cohere2_moe",
|
|
"minimax_m3",
|
|
}
|
|
|
|
# Speculative-decoding "helper" checkpoints (dFlash / MTP / assistant drafters)
|
|
# are never meant to be served as standalone chat models. Some declare a
|
|
# distinctive top-level model_type — an ``*_assistant`` (e.g. gemma4_assistant)
|
|
# or ``*_mtp`` (e.g. qwen3_5_mtp) marker — but DFlash draft checkpoints declare
|
|
# a plain model_type (e.g. ``qwen3``) and are only distinguishable by their
|
|
# architecture name (``DFlashDraftModel``) or a drafter-only config block
|
|
# (``dflash_config``). Keep these in sync with the drafter resolution in
|
|
# engine_pool.py (~1498) and the dflash gate in engine/dflash.py when new
|
|
# drafter families are added.
|
|
HELPER_CONFIG_MODEL_TYPE_SUFFIXES = ("_assistant", "_mtp")
|
|
_HELPER_ARCH_TOKENS = ("draft", "assistant", "mtp")
|
|
_HELPER_CONFIG_KEYS = ("dflash_config",)
|
|
|
|
|
|
def is_helper_config_model_type(config_model_type: str | None) -> bool:
|
|
"""True when ``config_model_type`` marks a speculative-decoding drafter.
|
|
|
|
These are the raw top-level ``model_type`` values from a checkpoint's
|
|
config.json (e.g. ``gemma4_assistant``, ``qwen3_5_mtp``). Note this misses
|
|
DFlash drafts, whose model_type is a plain ``qwen3`` — use
|
|
:func:`is_helper_model_config` when the full config dict is available.
|
|
"""
|
|
if not isinstance(config_model_type, str) or not config_model_type:
|
|
return False
|
|
mt = config_model_type.lower()
|
|
return mt.endswith(HELPER_CONFIG_MODEL_TYPE_SUFFIXES)
|
|
|
|
|
|
def is_helper_model_config(config: dict) -> bool:
|
|
"""True when a parsed config.json marks a speculative-decoding drafter.
|
|
|
|
Catches three intrinsic signals, any of which is definitive:
|
|
an ``*_assistant`` / ``*_mtp`` model_type, a drafter architecture name
|
|
(e.g. ``DFlashDraftModel``), or a drafter-only config block
|
|
(e.g. ``dflash_config``). These are helper checkpoints backing
|
|
dFlash / MTP / assistant speculative decoding, not chat models.
|
|
"""
|
|
if not isinstance(config, dict):
|
|
return False
|
|
if is_helper_config_model_type(config.get("model_type")):
|
|
return True
|
|
if any(key in config for key in _HELPER_CONFIG_KEYS):
|
|
return True
|
|
architectures = config.get("architectures") or []
|
|
if isinstance(architectures, str):
|
|
architectures = [architectures]
|
|
elif not isinstance(architectures, list | tuple | set):
|
|
return False
|
|
for arch in architectures:
|
|
arch_lower = str(arch).lower()
|
|
if any(token in arch_lower for token in _HELPER_ARCH_TOKENS):
|
|
return True
|
|
return False
|
|
|
|
|
|
# Known VLM architectures
|
|
VLM_ARCHITECTURES = {
|
|
"LlavaForConditionalGeneration",
|
|
"LlavaNextForConditionalGeneration",
|
|
"Qwen2VLForConditionalGeneration",
|
|
"Qwen2_5_VLForConditionalGeneration",
|
|
"MllamaForConditionalGeneration",
|
|
"Gemma3ForConditionalGeneration",
|
|
"Gemma4ForConditionalGeneration",
|
|
"InternVLChatModel",
|
|
"Idefics3ForConditionalGeneration",
|
|
"PaliGemmaForConditionalGeneration",
|
|
"Phi3VForCausalLM",
|
|
"Pixtral",
|
|
"MolmoForCausalLM",
|
|
"Molmo2ForConditionalGeneration",
|
|
"LlavaQwen2ForCausalLM", # apple/FastVLM (all sizes)
|
|
"Florence2ForConditionalGeneration",
|
|
}
|
|
|
|
# Known embedding model types from mlx-embeddings
|
|
EMBEDDING_MODEL_TYPES = {
|
|
"bert",
|
|
"xlm-roberta",
|
|
"xlm_roberta",
|
|
"modernbert",
|
|
"siglip",
|
|
"colqwen2_5",
|
|
"colqwen2-5",
|
|
}
|
|
|
|
# Model types that have both embedding and LLM variants.
|
|
# These require architecture-based disambiguation via EMBEDDING_ARCHITECTURES.
|
|
AMBIGUOUS_EMBEDDING_MODEL_TYPES = {
|
|
"qwen3",
|
|
"gemma3-text",
|
|
"gemma3_text",
|
|
"lfm2",
|
|
}
|
|
|
|
# Known embedding architectures
|
|
EMBEDDING_ARCHITECTURES = {
|
|
"BertModel",
|
|
"BertForMaskedLM",
|
|
"XLMRobertaModel",
|
|
"XLMRobertaForMaskedLM",
|
|
"ModernBertModel",
|
|
"ModernBertForMaskedLM",
|
|
"Qwen3ForTextEmbedding",
|
|
"SiglipModel",
|
|
"SiglipVisionModel",
|
|
"SiglipTextModel",
|
|
}
|
|
|
|
# Supported reranker architectures
|
|
SUPPORTED_RERANKER_ARCHITECTURES = {
|
|
"ModernBertForSequenceClassification", # via mlx-embeddings
|
|
"XLMRobertaForSequenceClassification", # omlx native implementation
|
|
"JinaForRanking", # Jina v3 listwise reranker
|
|
}
|
|
|
|
# CausalLM-based reranker architectures.
|
|
# These are standard CausalLM models fine-tuned for reranking via yes/no logit scoring.
|
|
# Detected by architecture + heuristic (model name or tokenizer hints).
|
|
CAUSAL_LM_RERANKER_ARCHITECTURES = {
|
|
"Qwen3ForCausalLM",
|
|
}
|
|
|
|
# CausalLM-based embedding architectures.
|
|
# These use a standard CausalLM architecture but are fine-tuned for embeddings
|
|
# (no lm_head weights). Detected by architecture + directory name heuristic.
|
|
CAUSAL_LM_EMBEDDING_ARCHITECTURES = {
|
|
"Qwen3ForCausalLM", # Qwen3-Embedding uses CausalLM arch without lm_head
|
|
"Qwen2ForCausalLM", # jina-code-embeddings & similar; only treated as an
|
|
# embedding when the dir-name heuristic (_is_causal_lm_embedding) also
|
|
# matches, so Qwen2/Qwen2.5 chat models are unaffected.
|
|
}
|
|
|
|
# Multimodal (VLM-based) reranker architectures.
|
|
# These share an architecture with VLM chat models but are fine-tuned for
|
|
# reranking. Loaded via mlx-embeddings' model.process() API. Distinguished
|
|
# from VLM chat by directory name heuristic.
|
|
MULTIMODAL_RERANKER_ARCHITECTURES = {
|
|
"Qwen3VLForConditionalGeneration", # Qwen3-VL-Reranker
|
|
}
|
|
|
|
# Multimodal (VLM-based) embedding architectures.
|
|
# Same arch as the reranker variant; distinguished by directory name hint.
|
|
MULTIMODAL_EMBEDDING_ARCHITECTURES = {
|
|
"Qwen3VLForConditionalGeneration", # Qwen3-VL-Embedding
|
|
}
|
|
|
|
# Unsupported reranker architectures (future support)
|
|
UNSUPPORTED_RERANKER_ARCHITECTURES = {
|
|
"BertForSequenceClassification",
|
|
"Qwen3ForSequenceClassification",
|
|
}
|
|
|
|
# All known reranker architectures (for model type detection)
|
|
RERANKER_ARCHITECTURES = SUPPORTED_RERANKER_ARCHITECTURES | UNSUPPORTED_RERANKER_ARCHITECTURES
|
|
|
|
# Unsupported model types — detected and skipped during discovery.
|
|
# Only top-level config fields are checked; nested audio_config/tts_config in
|
|
# multimodal models (e.g., MiniCPM-o) won't trigger this.
|
|
# Note: "whisper" and "qwen3_tts" were previously listed here but are now
|
|
# handled as audio types (audio_stt / audio_tts) — see AUDIO_* sets below.
|
|
UNSUPPORTED_MODEL_TYPES: set[str] = set()
|
|
|
|
UNSUPPORTED_ARCHITECTURES: set[str] = set()
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Audio model detection — dynamically loaded from mlx-audio when available
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# mlx-audio maintains MODEL_REMAPPING dicts (model_type → module directory)
|
|
# and model directories under mlx_audio/{stt,tts,sts}/models/. We read these
|
|
# at import time so oMLX automatically recognises new audio model families
|
|
# when mlx-audio is updated. Falls back to static sets when mlx-audio is
|
|
# not installed.
|
|
#
|
|
# Some base LLM model_types (qwen3, llama, …) collide with mlx-audio TTS
|
|
# model directory names because mlx-audio extends these architectures for
|
|
# audio. We exclude them so a plain Qwen3 LLM is not misdetected as TTS.
|
|
|
|
_LLM_TYPE_COLLISIONS = {"qwen3", "llama", "dense"}
|
|
|
|
|
|
def _build_audio_detection_sets():
|
|
"""Build STT/TTS/STS model-type sets from mlx-audio at import time.
|
|
|
|
Returns (stt_types, tts_types, sts_types) where each is a set of
|
|
model_type strings that should trigger audio detection.
|
|
"""
|
|
try:
|
|
from pathlib import Path as _P
|
|
|
|
import mlx_audio as _mla
|
|
|
|
_base = _P(_mla.__file__).parent
|
|
|
|
def _dir_names(subdir: str) -> set:
|
|
d = _base / subdir / "models"
|
|
if d.is_dir():
|
|
return {p.name for p in d.iterdir()
|
|
if p.is_dir() and not p.name.startswith("__")}
|
|
return set()
|
|
|
|
# TTS: MODEL_REMAPPING keys + model dir names
|
|
from mlx_audio.tts.utils import MODEL_REMAPPING as _tts_remap
|
|
tts = set(_tts_remap.keys()) | _dir_names("tts")
|
|
|
|
# STT: MODEL_REMAPPING keys + model dir names
|
|
from mlx_audio.stt.utils import MODEL_REMAPPING as _stt_remap
|
|
stt = set(_stt_remap.keys()) | _dir_names("stt")
|
|
|
|
# STS: model dir names only (no unified utils/remapping)
|
|
sts = _dir_names("sts")
|
|
|
|
# Strip base-LLM names that collide with audio model dirs
|
|
tts -= _LLM_TYPE_COLLISIONS
|
|
stt -= _LLM_TYPE_COLLISIONS
|
|
|
|
logger.debug(
|
|
"Audio detection sets loaded from mlx-audio: "
|
|
"STT=%d, TTS=%d, STS=%d", len(stt), len(tts), len(sts),
|
|
)
|
|
return stt, tts, sts
|
|
|
|
except Exception:
|
|
logger.debug("mlx-audio not available — using static audio detection sets")
|
|
# Static fallback so model discovery still works without mlx-audio
|
|
_stt = {"whisper", "qwen3_asr", "parakeet", "qwen2_audio"}
|
|
_tts = {"qwen3_tts", "kokoro", "chatterbox", "vibevoice", "vibevoice_streaming", "kugelaudio", "audiodit"}
|
|
_sts = {"deepfilternet", "mossformer2_se", "sam_audio", "lfm_audio"}
|
|
return _stt, _tts, _sts
|
|
|
|
|
|
AUDIO_STT_MODEL_TYPES, AUDIO_TTS_MODEL_TYPES, AUDIO_STS_MODEL_TYPES = (
|
|
_build_audio_detection_sets()
|
|
)
|
|
|
|
# Architecture-based detection — these are checked before model_type and
|
|
# are always static because architecture strings are stable identifiers.
|
|
AUDIO_STT_ARCHITECTURES = {
|
|
"WhisperForConditionalGeneration",
|
|
"Qwen3ASRForConditionalGeneration",
|
|
"ParakeetForCTC",
|
|
"Qwen2AudioForConditionalGeneration",
|
|
}
|
|
|
|
AUDIO_TTS_ARCHITECTURES = {
|
|
"KokoroForConditionalGeneration",
|
|
"Qwen3TTSForConditionalGeneration",
|
|
"ChatterboxForConditionalGeneration",
|
|
"VibeVoiceForConditionalGeneration",
|
|
"VibeVoiceStreamingForConditionalGenerationInference",
|
|
"KugelAudioForConditionalGeneration",
|
|
}
|
|
|
|
AUDIO_STS_ARCHITECTURES = {
|
|
"DeepFilterNetModel",
|
|
"MossFormer2SEModel",
|
|
"SAMAudio",
|
|
"LFM2AudioModel",
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class DiscoveredModel:
|
|
"""Information about a discovered model."""
|
|
|
|
model_id: str # Directory name (e.g., "llama-3b")
|
|
model_path: str # Full path to model directory
|
|
model_type: ModelType # "llm", "vlm", "embedding", or "reranker"
|
|
engine_type: EngineType # "batched", "vlm", "embedding", or "reranker"
|
|
estimated_size: int # Estimated memory usage in bytes
|
|
config_model_type: str = "" # Raw model_type from config.json (e.g., "deepseekocr_2")
|
|
thinking_default: bool | None = None # True if model thinks by default, False if not, None if unknown
|
|
preserve_thinking_default: bool | None = None # True when template supports preserve_thinking (Qwen 3.6+)
|
|
model_context_length: int | None = None # Declared context length from config.json (None if unknown)
|
|
source_type: str = "local" # "local" or "hf_cache"
|
|
source_repo_id: str | None = None # HuggingFace repo id for cache-backed models
|
|
is_helper: bool = False # Speculative-decoding drafter (dFlash/Assistant/MTP)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HfCacheEntry:
|
|
"""Resolved HuggingFace Hub cache entry."""
|
|
|
|
snapshot_path: Path
|
|
model_id: str
|
|
source_repo_id: str
|
|
|
|
|
|
def _is_unsupported_model(model_path: Path) -> bool:
|
|
"""
|
|
Check if model is an unsupported type that should be skipped during discovery.
|
|
|
|
Audio models (STT/TTS) are NOT unsupported — they are detected as
|
|
"audio_stt" or "audio_tts" by detect_model_type() and served via
|
|
their own engine types.
|
|
|
|
Only checks top-level config fields. Multimodal models with nested
|
|
audio_config/tts_config (e.g., MiniCPM-o) are not affected.
|
|
"""
|
|
config_path = model_path / "config.json"
|
|
if not config_path.exists():
|
|
return False
|
|
|
|
try:
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
except (json.JSONDecodeError, IOError):
|
|
return False
|
|
|
|
architectures = config.get("architectures", [])
|
|
for arch in architectures:
|
|
if arch in UNSUPPORTED_ARCHITECTURES:
|
|
return True
|
|
|
|
model_type = config.get("model_type", "")
|
|
normalized = model_type.lower().replace("-", "_")
|
|
return normalized in UNSUPPORTED_MODEL_TYPES or model_type in UNSUPPORTED_MODEL_TYPES
|
|
|
|
|
|
def _is_causal_lm_reranker(model_path: Path) -> bool:
|
|
"""
|
|
Heuristic check for CausalLM models fine-tuned as rerankers.
|
|
|
|
CausalLM rerankers (e.g., Qwen3-Reranker) use the same architecture as
|
|
their base LLMs but are fine-tuned to output yes/no logits for relevance
|
|
scoring. We detect them by checking the model directory name for "reranker"
|
|
or "rerank" keywords, since config.json is identical to a standard LLM.
|
|
"""
|
|
name_lower = model_path.name.lower()
|
|
return "reranker" in name_lower or "rerank" in name_lower
|
|
|
|
|
|
def _is_causal_lm_embedding(model_path: Path) -> bool:
|
|
"""
|
|
Heuristic check for CausalLM models fine-tuned as embedding models.
|
|
|
|
CausalLM embeddings (e.g., Qwen3-Embedding) use the same architecture as
|
|
their base LLMs but are fine-tuned for embeddings and ship without lm_head
|
|
weights. We detect them by checking the model directory name for "embedding"
|
|
or "embed" keywords, since config.json is identical to a standard LLM.
|
|
"""
|
|
name_lower = model_path.name.lower()
|
|
return "embedding" in name_lower or "embed" in name_lower
|
|
|
|
|
|
def _has_sentence_transformers_embedding_pipeline(model_path: Path) -> bool:
|
|
"""
|
|
Detect sentence-transformers style embedding exports via modules.json.
|
|
|
|
This allows oMLX to recognize embedding exports whose base transformer
|
|
architecture is ambiguous (for example gemma3_text) but which include
|
|
sentence-transformers pooling/normalization modules.
|
|
"""
|
|
modules_path = model_path / "modules.json"
|
|
if not modules_path.exists():
|
|
return False
|
|
|
|
try:
|
|
with open(modules_path) as f:
|
|
modules = json.load(f)
|
|
except (json.JSONDecodeError, IOError):
|
|
return False
|
|
|
|
if not isinstance(modules, list):
|
|
return False
|
|
|
|
module_types = {
|
|
module.get("type", "")
|
|
for module in modules
|
|
if isinstance(module, dict)
|
|
}
|
|
if "sentence_transformers.models.Transformer" not in module_types:
|
|
return False
|
|
|
|
return any(
|
|
module_type.startswith("sentence_transformers.models.")
|
|
and module_type != "sentence_transformers.models.Transformer"
|
|
for module_type in module_types
|
|
)
|
|
|
|
|
|
def _looks_like_kokoro_config(config: dict) -> bool:
|
|
"""Return True for Kokoro exports that omit HF ``model_type``.
|
|
|
|
mlx-community Kokoro conversions (e.g. Kokoro-82M-bf16) keep the original
|
|
Kokoro config — top-level ``istftnet`` + ``plbert`` sections and a
|
|
``vocab`` table — with no HF-style ``model_type``/``architectures``.
|
|
mlx-audio loads them fine, but oMLX must classify them as TTS during
|
|
discovery or they fall through to the LLM engine, whose loader only
|
|
matches ``model*.safetensors`` and fails with a misleading
|
|
"No safetensors found" error.
|
|
"""
|
|
if not isinstance(config, dict):
|
|
return False
|
|
return (
|
|
isinstance(config.get("istftnet"), dict)
|
|
and isinstance(config.get("plbert"), dict)
|
|
and isinstance(config.get("vocab"), dict)
|
|
)
|
|
|
|
|
|
def _looks_like_nemo_asr_config(config: dict) -> bool:
|
|
"""Return True for NeMo ASR exports that omit HF ``model_type``.
|
|
|
|
NVIDIA Parakeet TDT/CTC MLX conversions keep the original NeMo ASR
|
|
training config instead of a HuggingFace-style ``model_type`` or
|
|
``architectures`` field. mlx-audio can load these models by name, but
|
|
oMLX must still classify them as STT during discovery or they fall through
|
|
to the LLM engine and fail with a misleading ``'model_type'`` error.
|
|
|
|
Only top-level NeMo ASR module targets are considered so multimodal models
|
|
with nested ``audio_config`` sections are not misclassified.
|
|
"""
|
|
if not isinstance(config, dict):
|
|
return False
|
|
|
|
module_targets: list[str] = []
|
|
for key in ("preprocessor", "encoder", "decoder", "joint"):
|
|
value = config.get(key)
|
|
if isinstance(value, dict):
|
|
target = value.get("_target_")
|
|
if isinstance(target, str):
|
|
module_targets.append(target.lower())
|
|
|
|
if not any("nemo.collections.asr" in target for target in module_targets):
|
|
return False
|
|
|
|
# NeMo ASR configs include an audio preprocessor plus tokenizer/decoder
|
|
# metadata. Requiring these keeps the heuristic narrow while covering
|
|
# Parakeet TDT exports whose config has no model_type at all.
|
|
preprocessor = config.get("preprocessor")
|
|
has_audio_preprocessor = isinstance(preprocessor, dict) and (
|
|
"audio" in str(preprocessor.get("_target_", "")).lower()
|
|
or "melspectrogram" in str(preprocessor.get("_target_", "")).lower()
|
|
)
|
|
has_asr_head = isinstance(config.get("decoder"), dict) or isinstance(
|
|
config.get("joint"), dict
|
|
)
|
|
has_tokenizer = isinstance(config.get("tokenizer"), dict)
|
|
return has_audio_preprocessor and has_asr_head and has_tokenizer
|
|
|
|
|
|
def _has_vision_subconfig(config: dict) -> bool:
|
|
"""
|
|
Return True if ``config`` carries evidence of a vision sub-config.
|
|
|
|
Three keys cover the conventions in the wild:
|
|
|
|
- ``vision_config`` — most VLMs (Qwen2-VL, Gemma3, LLaVA-Next, ...).
|
|
- ``vit_config`` — Molmo / Molmo2 family.
|
|
- ``mm_vision_tower`` — older LLaVA family including FastVLM's
|
|
``llava_qwen2``. The check is non-empty-only: a config-stub text-only
|
|
quant could in principle declare a tower path it doesn't ship weights
|
|
for, but in practice bf16 FastVLM ships a real path string.
|
|
|
|
Used by the VLM classifier in :func:`detect_model_type` and by other
|
|
paths (``oq``, admin model info) that need to ask "is this a VLM?".
|
|
"""
|
|
return (
|
|
"vision_config" in config
|
|
or "vit_config" in config
|
|
or bool(config.get("mm_vision_tower"))
|
|
)
|
|
|
|
|
|
def _architecture_indicates_causal_lm(architectures: list[str]) -> bool:
|
|
"""True when ``architectures`` describe a text causal LM (not mlx-audio STS).
|
|
|
|
Liquid LFM text checkpoints (LFM2, LFM2.5 MoE, etc.) use ``lfm*`` model
|
|
types and ``*ForCausalLM`` classes. mlx-audio LFM STS uses ``LFM2AudioModel``
|
|
and is handled earlier via :data:`AUDIO_STS_ARCHITECTURES`.
|
|
"""
|
|
return any("causallm" in arch.lower() for arch in architectures)
|
|
|
|
|
|
def detect_model_type(model_path: Path) -> ModelType:
|
|
"""
|
|
Detect model type from config.json.
|
|
|
|
Checks:
|
|
1. architectures field for reranker-specific classes (SequenceClassification)
|
|
2. CausalLM-based reranker/embedding detection (architecture + directory name)
|
|
3. sentence-transformers pipeline detection via modules.json
|
|
4. architectures field for embedding-specific classes
|
|
5. model_type field against known embedding types (unambiguous only)
|
|
6. VLM detection via architectures, model_type, or vision sub-config
|
|
presence (``vision_config`` / ``vit_config`` / non-empty
|
|
``mm_vision_tower`` — see :func:`_has_vision_subconfig`)
|
|
7. Audio model detection (STT/TTS/STS)
|
|
|
|
Args:
|
|
model_path: Path to model directory
|
|
|
|
Returns:
|
|
Model type: "llm", "vlm", "embedding", "reranker", "audio_stt", "audio_tts", or "audio_sts"
|
|
"""
|
|
config_path = model_path / "config.json"
|
|
if not config_path.exists():
|
|
return "llm"
|
|
|
|
try:
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
except (json.JSONDecodeError, IOError):
|
|
return "llm"
|
|
|
|
# Check architectures field for reranker first (more specific)
|
|
architectures = config.get("architectures", [])
|
|
for arch in architectures:
|
|
if arch in RERANKER_ARCHITECTURES:
|
|
return "reranker"
|
|
|
|
# Check for CausalLM-based rerankers (e.g., Qwen3-Reranker).
|
|
# These use a standard CausalLM architecture but are fine-tuned for reranking
|
|
# via yes/no logit scoring. Detected by architecture + model directory name hint.
|
|
for arch in architectures:
|
|
if arch in CAUSAL_LM_RERANKER_ARCHITECTURES:
|
|
if _is_causal_lm_reranker(model_path):
|
|
return "reranker"
|
|
|
|
# Check for CausalLM-based embeddings (e.g., Qwen3-Embedding).
|
|
# These use a standard CausalLM architecture but are fine-tuned for embeddings
|
|
# and ship without lm_head weights. Detected by architecture + directory name hint.
|
|
for arch in architectures:
|
|
if arch in CAUSAL_LM_EMBEDDING_ARCHITECTURES:
|
|
if _is_causal_lm_embedding(model_path):
|
|
return "embedding"
|
|
|
|
# Check for multimodal (VLM-based) rerankers and embeddings.
|
|
# Same architecture string as VLM chat models; distinguished by the
|
|
# directory name heuristic. Must come before VLM detection below so
|
|
# the reranker/embedding hint wins over default VLM classification.
|
|
for arch in architectures:
|
|
if arch in MULTIMODAL_RERANKER_ARCHITECTURES and _is_causal_lm_reranker(model_path):
|
|
return "reranker"
|
|
if arch in MULTIMODAL_EMBEDDING_ARCHITECTURES and _is_causal_lm_embedding(model_path):
|
|
return "embedding"
|
|
|
|
if _has_sentence_transformers_embedding_pipeline(model_path):
|
|
return "embedding"
|
|
|
|
# Check architectures field for embedding (before model_type to avoid
|
|
# false positives from ambiguous model types like qwen3, gemma3-text)
|
|
for arch in architectures:
|
|
if arch in EMBEDDING_ARCHITECTURES:
|
|
return "embedding"
|
|
|
|
# Check model_type field for unambiguous embedding types
|
|
model_type = config.get("model_type", "")
|
|
# Normalize: replace hyphens with underscores and lowercase
|
|
normalized_type = model_type.lower().replace("-", "_")
|
|
|
|
if normalized_type in EMBEDDING_MODEL_TYPES or model_type in EMBEDDING_MODEL_TYPES:
|
|
return "embedding"
|
|
|
|
# Ambiguous embedding types (have both embedding and LLM variants):
|
|
# only classified as embedding if architecture matched above
|
|
if (
|
|
normalized_type in AMBIGUOUS_EMBEDDING_MODEL_TYPES
|
|
or model_type in AMBIGUOUS_EMBEDDING_MODEL_TYPES
|
|
):
|
|
logger.info(
|
|
f"Model type '{model_type}' has both embedding and LLM variants, "
|
|
f"but architecture {architectures} is not an embedding architecture "
|
|
"— treating as LLM"
|
|
)
|
|
|
|
if normalized_type in VLM_NATIVE_TEXT_MODEL_TYPES:
|
|
logger.info(
|
|
f"{model_type} detected as mlx-vlm native text model"
|
|
)
|
|
return "vlm"
|
|
|
|
# Check for VLM: architectures field
|
|
# Some text-only quants (e.g., unsloth/gemma-4-31b-it-MLX-8bit) keep the VLM
|
|
# architecture name but strip vision_config and vision weights.
|
|
# For model families known to have text-only variants, require evidence
|
|
# of a vision sub-config — see :func:`_has_vision_subconfig` for the
|
|
# three keys we accept (``vision_config``, ``vit_config``,
|
|
# ``mm_vision_tower``).
|
|
for arch in architectures:
|
|
if arch in VLM_ARCHITECTURES:
|
|
if normalized_type in VLM_MODEL_TYPES and not _has_vision_subconfig(config):
|
|
logger.info(
|
|
f"Architecture '{arch}' is a VLM architecture but no "
|
|
"vision_config / vit_config / mm_vision_tower found — "
|
|
"treating as LLM (text-only quant)"
|
|
)
|
|
break
|
|
return "vlm"
|
|
|
|
# Check for VLM: model_type field (only if vision capabilities are present)
|
|
# Some model families (e.g., qwen3_5_moe) have both VLM and text-only variants.
|
|
# Text-only quants won't carry a vision sub-config. gemma4_unified and
|
|
# diffusion_gemma are exceptions: they are served by mlx-vlm regardless of
|
|
# vision_config presence in config.json.
|
|
if normalized_type in VLM_MODEL_TYPES:
|
|
if normalized_type in {"gemma4_unified", "diffusion_gemma"}:
|
|
logger.info(
|
|
f"{model_type} detected as VLM (mlx-vlm native model)"
|
|
)
|
|
return "vlm"
|
|
if _has_vision_subconfig(config):
|
|
return "vlm"
|
|
logger.info(
|
|
f"Model type '{model_type}' is in VLM_MODEL_TYPES but no "
|
|
"vision_config / vit_config / mm_vision_tower found — "
|
|
"treating as LLM (text-only quant)"
|
|
)
|
|
|
|
# Check for VLM: presence of a vision sub-config (fallback heuristic).
|
|
# Catch-all for VLMs that aren't yet listed in VLM_MODEL_TYPES.
|
|
if _has_vision_subconfig(config):
|
|
return "vlm"
|
|
|
|
# Check for audio models — architectures take priority over model_type.
|
|
# Only top-level architectures/model_type are inspected; nested audio_config
|
|
# inside multimodal models (e.g., MiniCPM-o) does not trigger this path.
|
|
#
|
|
# Architecture check first (unambiguous):
|
|
for arch in architectures:
|
|
if arch in AUDIO_STT_ARCHITECTURES:
|
|
return "audio_stt"
|
|
|
|
# NeMo ASR exports such as mlx-community/parakeet-tdt-0.6b-v3 ship a
|
|
# NeMo training config without HF model_type/architectures. They are
|
|
# still STT models and mlx-audio can load them by directory/repo name.
|
|
if _looks_like_nemo_asr_config(config):
|
|
return "audio_stt"
|
|
# Kokoro exports similarly ship a bare original config (istftnet/plbert).
|
|
if _looks_like_kokoro_config(config):
|
|
return "audio_tts"
|
|
for arch in architectures:
|
|
if arch in AUDIO_TTS_ARCHITECTURES:
|
|
return "audio_tts"
|
|
for arch in architectures:
|
|
if arch in AUDIO_STS_ARCHITECTURES:
|
|
return "audio_sts"
|
|
|
|
# model_type check (dynamically loaded from mlx-audio when available).
|
|
# Check TTS before STT because some model_type values (e.g. "vibevoice")
|
|
# appear in both sets — TTS is the more common category for these.
|
|
if normalized_type in AUDIO_TTS_MODEL_TYPES or model_type in AUDIO_TTS_MODEL_TYPES:
|
|
return "audio_tts"
|
|
if normalized_type in AUDIO_STT_MODEL_TYPES or model_type in AUDIO_STT_MODEL_TYPES:
|
|
return "audio_stt"
|
|
if normalized_type in AUDIO_STS_MODEL_TYPES or model_type in AUDIO_STS_MODEL_TYPES:
|
|
return "audio_sts"
|
|
# mlx-audio LFM STS may use an "lfm*" model_type without a known architecture
|
|
# string yet. Liquid LFM *text* checkpoints share that prefix — disambiguate
|
|
# with CausalLM architecture names (LFM2 / LFM2.5 MoE, future lfm* LMs).
|
|
if normalized_type.startswith("lfm") and normalized_type not in EMBEDDING_MODEL_TYPES:
|
|
if _architecture_indicates_causal_lm(architectures):
|
|
return "llm"
|
|
return "audio_sts"
|
|
|
|
return "llm"
|
|
|
|
|
|
def detect_thinking_default(model_path: Path) -> bool | None:
|
|
"""Detect whether a model's chat template enables thinking by default.
|
|
|
|
Inspects the Jinja chat template for ``enable_thinking`` references and
|
|
determines the default behaviour:
|
|
|
|
* **True** — model thinks by default (e.g. Qwen 3.x: only suppresses
|
|
thinking when ``enable_thinking is false``).
|
|
* **False** — model suppresses thinking by default (e.g. Gemma 4: only
|
|
enables thinking when ``enable_thinking`` is truthy,
|
|
``default(false)``).
|
|
* **None** — template does not reference ``enable_thinking`` (model has
|
|
no thinking toggle).
|
|
"""
|
|
# Try standalone Jinja file first, then tokenizer_config.json
|
|
template_text = None
|
|
jinja_path = model_path / "chat_template.jinja"
|
|
if jinja_path.exists():
|
|
with contextlib.suppress(OSError):
|
|
template_text = jinja_path.read_text(encoding="utf-8")
|
|
|
|
if template_text is None:
|
|
tc_path = model_path / "tokenizer_config.json"
|
|
if tc_path.exists():
|
|
try:
|
|
with open(tc_path) as f:
|
|
tc = json.load(f)
|
|
template_text = tc.get("chat_template")
|
|
except Exception:
|
|
pass
|
|
|
|
if not template_text or "enable_thinking" not in template_text:
|
|
return None
|
|
|
|
# Heuristic: if the template only disables thinking when explicitly
|
|
# ``enable_thinking is false``, then thinking is ON by default.
|
|
# If the template requires ``enable_thinking`` to be truthy or uses
|
|
# ``default(false)``, then thinking is OFF by default.
|
|
if "enable_thinking is false" in template_text:
|
|
return True # ON by default (Qwen pattern)
|
|
if "default(false)" in template_text or "enable_thinking)" in template_text:
|
|
return False # OFF by default (Gemma pattern)
|
|
|
|
return None
|
|
|
|
|
|
# Context-length keys, in priority order. Order mirrors HuggingFace
|
|
# Transformers conventions: ``max_position_embeddings`` is the canonical
|
|
# field for decoder-only LLMs; ``max_seq_len`` / ``seq_length`` show up on
|
|
# Llama / Mistral / Qwen forks; ``n_positions`` is the GPT-2 lineage.
|
|
_CONTEXT_LENGTH_KEYS = (
|
|
"max_position_embeddings",
|
|
"max_seq_len",
|
|
"max_seq_length",
|
|
"seq_length",
|
|
"n_positions",
|
|
)
|
|
|
|
# tokenizer_config.json's ``model_max_length`` is Transformers' fallback
|
|
# field. Transformers seeds it with ``int(1e30)`` when the tokenizer has
|
|
# no real cap, and downstream code distinguishes that sentinel from a
|
|
# real long context. Anything above ~1e18 is treated as the sentinel.
|
|
_TOKENIZER_MAX_LENGTH_SENTINEL = 10**18
|
|
|
|
|
|
def _read_model_context_length(model_path: Path) -> int | None:
|
|
"""Discover the declared context length from a model's config files.
|
|
|
|
Resolution order:
|
|
|
|
1. Top-level ``config.json`` keys (``max_position_embeddings`` first,
|
|
then the rest of :data:`_CONTEXT_LENGTH_KEYS`).
|
|
2. Nested ``text_config`` / ``language_config`` keys (used by VLM
|
|
wrappers and Qwen-style MoE configs that put the language head in
|
|
a sub-object).
|
|
3. ``tokenizer_config.json``'s ``model_max_length`` — but only when
|
|
it is a finite positive integer, since Transformers seeds it with
|
|
``int(1e30)`` as a "no cap" sentinel.
|
|
|
|
Returns:
|
|
Positive integer context length, or ``None`` when no usable
|
|
value was found in any of the above.
|
|
"""
|
|
config_path = model_path / "config.json"
|
|
if config_path.exists():
|
|
try:
|
|
with open(config_path, encoding="utf-8") as f:
|
|
model_config = json.load(f)
|
|
except (OSError, json.JSONDecodeError) as e:
|
|
logger.debug(f"Failed to read config.json for {model_path}: {e}")
|
|
model_config = {}
|
|
|
|
for key in _CONTEXT_LENGTH_KEYS:
|
|
value = model_config.get(key)
|
|
if isinstance(value, int) and value > 0:
|
|
return value
|
|
|
|
for nest_key in ("text_config", "language_config"):
|
|
nested = model_config.get(nest_key)
|
|
if isinstance(nested, dict):
|
|
for key in _CONTEXT_LENGTH_KEYS:
|
|
value = nested.get(key)
|
|
if isinstance(value, int) and value > 0:
|
|
return value
|
|
|
|
tc_path = model_path / "tokenizer_config.json"
|
|
if tc_path.exists():
|
|
try:
|
|
with open(tc_path, encoding="utf-8") as f:
|
|
tc = json.load(f)
|
|
except (OSError, json.JSONDecodeError) as e:
|
|
logger.debug(f"Failed to read tokenizer_config.json for {model_path}: {e}")
|
|
tc = {}
|
|
|
|
value = tc.get("model_max_length")
|
|
if isinstance(value, int) and 0 < value < _TOKENIZER_MAX_LENGTH_SENTINEL:
|
|
return value
|
|
|
|
return None
|
|
|
|
|
|
def detect_preserve_thinking(model_path: Path) -> bool | None:
|
|
"""Detect whether a model's chat template supports ``preserve_thinking``.
|
|
|
|
Qwen 3.6+ templates strip ``<think>`` blocks from historical assistant
|
|
turns by default and only keep them when ``preserve_thinking`` is true.
|
|
Stripping breaks KV prefix cache reuse, so we default to True when the
|
|
template supports this flag.
|
|
|
|
Returns:
|
|
True if the template references ``preserve_thinking`` (should be
|
|
enabled), None otherwise (template has no such flag).
|
|
"""
|
|
template_text = None
|
|
jinja_path = model_path / "chat_template.jinja"
|
|
if jinja_path.exists():
|
|
with contextlib.suppress(OSError):
|
|
template_text = jinja_path.read_text(encoding="utf-8")
|
|
|
|
if template_text is None:
|
|
tc_path = model_path / "tokenizer_config.json"
|
|
if tc_path.exists():
|
|
try:
|
|
with open(tc_path) as f:
|
|
tc = json.load(f)
|
|
template_text = tc.get("chat_template")
|
|
except Exception:
|
|
pass
|
|
|
|
if not template_text or "preserve_thinking" not in template_text:
|
|
return None
|
|
|
|
return True
|
|
|
|
|
|
def estimate_model_size(model_path: Path) -> int:
|
|
"""
|
|
Estimate model memory usage from safetensors/bin file sizes.
|
|
|
|
MLX keeps quantized weights in compressed form, so file size ≈ memory usage.
|
|
|
|
Args:
|
|
model_path: Path to model directory
|
|
|
|
Returns:
|
|
Estimated memory usage in bytes
|
|
"""
|
|
total_size = 0
|
|
|
|
# Primary: safetensors files
|
|
safetensors_files = list(model_path.glob("*.safetensors"))
|
|
for f in safetensors_files:
|
|
total_size += f.stat().st_size
|
|
|
|
# Fallback: .bin files (older PyTorch format)
|
|
if total_size == 0:
|
|
for f in model_path.glob("*.bin"):
|
|
# Filter out non-weight files
|
|
name_lower = f.name.lower()
|
|
if "optimizer" in name_lower or "training" in name_lower:
|
|
continue
|
|
total_size += f.stat().st_size
|
|
|
|
# Also check in subdirectories (some models store weights in subfolders)
|
|
if total_size == 0:
|
|
for f in model_path.glob("**/*.safetensors"):
|
|
total_size += f.stat().st_size
|
|
|
|
if total_size == 0:
|
|
raise ValueError(f"No model weights found in {model_path}")
|
|
|
|
# Add overhead for runtime buffers (~5%)
|
|
overhead_factor = 1.05
|
|
|
|
return int(total_size * overhead_factor)
|
|
|
|
|
|
def _is_adapter_dir(path: Path) -> bool:
|
|
"""Check if a directory contains a LoRA/PEFT adapter (has adapter_config.json)."""
|
|
return (path / "adapter_config.json").exists()
|
|
|
|
|
|
def _is_model_dir(path: Path) -> bool:
|
|
"""Check if a directory contains a valid model (has config.json)."""
|
|
return (path / "config.json").exists() and not _is_adapter_dir(path)
|
|
|
|
|
|
def model_directory_access_error(path: Path) -> str | None:
|
|
"""Return a user-facing error if a model directory cannot be scanned."""
|
|
try:
|
|
if not path.exists():
|
|
return f"Model directory does not exist: {path}"
|
|
if not path.is_dir():
|
|
return f"Model directory is not a directory: {path}"
|
|
next(path.iterdir(), None)
|
|
except OSError as e:
|
|
return (
|
|
f"Model directory is not readable: {path} "
|
|
f"({type(e).__name__}: {e})"
|
|
)
|
|
return None
|
|
|
|
|
|
def model_directory_write_error(path: Path, *, create: bool = False) -> str | None:
|
|
"""Return a user-facing error if a model directory cannot be written."""
|
|
try:
|
|
if not path.exists():
|
|
if create:
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
else:
|
|
return f"Model directory does not exist: {path}"
|
|
if not path.is_dir():
|
|
return f"Model directory is not a directory: {path}"
|
|
except OSError as e:
|
|
return (
|
|
f"Model directory is not writable: {path} "
|
|
f"({type(e).__name__}: {e})"
|
|
)
|
|
|
|
access_error = model_directory_access_error(path)
|
|
if access_error is not None:
|
|
return access_error
|
|
|
|
try:
|
|
with tempfile.NamedTemporaryFile(
|
|
prefix=".omlx-write-test-",
|
|
dir=path,
|
|
delete=True,
|
|
) as f:
|
|
f.write(b"")
|
|
f.flush()
|
|
except OSError as e:
|
|
return (
|
|
f"Model directory is not writable: {path} "
|
|
f"({type(e).__name__}: {e})"
|
|
)
|
|
|
|
return None
|
|
|
|
|
|
def _iter_readable_entries(path: Path, context: str) -> list[Path]:
|
|
"""Return sorted directory entries, or an empty list when scanning fails."""
|
|
try:
|
|
return sorted(path.iterdir())
|
|
except OSError as e:
|
|
logger.warning(
|
|
"Skipping unreadable %s %s: %s: %s",
|
|
context,
|
|
path,
|
|
type(e).__name__,
|
|
e,
|
|
)
|
|
return []
|
|
|
|
|
|
def _is_readable_dir(path: Path, context: str) -> bool:
|
|
try:
|
|
return path.is_dir()
|
|
except OSError as e:
|
|
logger.warning(
|
|
"Skipping inaccessible %s %s: %s: %s",
|
|
context,
|
|
path,
|
|
type(e).__name__,
|
|
e,
|
|
)
|
|
return False
|
|
|
|
|
|
def _decode_hf_cache_model_id(path: Path) -> tuple[str, str] | None:
|
|
"""Decode models--org--repo into (route_safe_id, repo_id)."""
|
|
name = path.name
|
|
if not name.startswith("models--"):
|
|
return None
|
|
|
|
encoded = name[len("models--"):]
|
|
if not encoded:
|
|
return None
|
|
|
|
parts = encoded.split("--")
|
|
if len(parts) == 1:
|
|
return parts[0], parts[0]
|
|
|
|
repo_name = "--".join(parts[1:])
|
|
return f"{parts[0]}--{repo_name}", f"{parts[0]}/{repo_name}"
|
|
|
|
|
|
def _resolve_hf_cache_entry(path: Path) -> HfCacheEntry | None:
|
|
"""Resolve an HF Hub cache entry (models--Org--Name/) to its active snapshot.
|
|
|
|
Returns an HfCacheEntry or None if not a valid HF model cache entry.
|
|
"""
|
|
decoded = _decode_hf_cache_model_id(path)
|
|
if decoded is None:
|
|
return None
|
|
model_id, source_repo_id = decoded
|
|
|
|
snapshots_dir = path / "snapshots"
|
|
if not snapshots_dir.is_dir():
|
|
return None
|
|
|
|
for ref_name in ("main", "master"):
|
|
try:
|
|
commit_hash = (path / "refs" / ref_name).read_text().strip()
|
|
except OSError:
|
|
continue
|
|
snapshot = snapshots_dir / commit_hash
|
|
if snapshot.is_dir():
|
|
return HfCacheEntry(snapshot, model_id, source_repo_id)
|
|
|
|
snapshots = [
|
|
p
|
|
for p in _iter_readable_entries(snapshots_dir, "HF cache snapshots")
|
|
if _is_readable_dir(p, "HF cache snapshot")
|
|
]
|
|
if not snapshots:
|
|
return None
|
|
if len(snapshots) == 1:
|
|
return HfCacheEntry(snapshots[0], model_id, source_repo_id)
|
|
|
|
snapshot = max(snapshots, key=lambda p: p.stat().st_mtime)
|
|
return HfCacheEntry(snapshot, model_id, source_repo_id)
|
|
|
|
|
|
def _safetensors_has_mlx_metadata(path: Path) -> bool:
|
|
"""Return True if any model safetensors shard declares MLX format."""
|
|
try:
|
|
from safetensors import safe_open
|
|
except Exception as e:
|
|
logger.debug(f"safetensors import failed while checking {path}: {e}")
|
|
return False
|
|
|
|
for shard in sorted(path.glob("model*.safetensors")):
|
|
try:
|
|
with safe_open(str(shard), framework="numpy") as f:
|
|
metadata = f.metadata() or {}
|
|
except Exception as e:
|
|
logger.debug(f"Could not read safetensors metadata from {shard}: {e}")
|
|
continue
|
|
if str(metadata.get("format", "")).lower() == "mlx":
|
|
return True
|
|
return False
|
|
|
|
|
|
_MLX_NAME_RE = re.compile(r"(^|[-_/])mlx($|[-_/])", re.IGNORECASE)
|
|
|
|
# Speculative-decoding helper checkpoints (e.g. z-lab/Qwen3.6-27B-DFlash) are
|
|
# MLX-loadable drafts even though their safetensors are saved in PyTorch ("pt")
|
|
# format and the repo name carries no "mlx" token, so the HF-cache MLX
|
|
# heuristic must recognise them explicitly or they vanish from the draft-model
|
|
# picker (#1643). Detection delegates to is_helper_model_config so the drafter
|
|
# markers stay in sync with helper flagging and engine_pool's drafter
|
|
# resolution.
|
|
def _is_helper_checkpoint(model_path: Path) -> bool:
|
|
"""True if ``model_path`` is a speculative-decoding helper checkpoint.
|
|
|
|
Must never raise on an unreadable entry: unlike the config reads inside
|
|
_register_model, this runs outside discover_models' per-model guard, so
|
|
an escaped exception would abort the whole scan. UnicodeDecodeError from
|
|
a non-UTF-8 config.json is a ValueError, not a JSONDecodeError.
|
|
"""
|
|
config_path = model_path / "config.json"
|
|
if not config_path.exists():
|
|
return False
|
|
try:
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
except (json.JSONDecodeError, UnicodeDecodeError, OSError):
|
|
return False
|
|
return is_helper_model_config(config)
|
|
|
|
|
|
def _is_hf_cache_mlx_compatible(model_dir: Path, source_repo_id: str) -> bool:
|
|
"""Heuristic for HF cache entries that can be loaded without conversion."""
|
|
if not _is_model_dir(model_dir):
|
|
return False
|
|
if not list(model_dir.glob("model*.safetensors")):
|
|
logger.debug(f"Skipping HF cache model without model*.safetensors: {source_repo_id}")
|
|
return False
|
|
if _safetensors_has_mlx_metadata(model_dir):
|
|
return True
|
|
|
|
repo_lower = source_repo_id.lower()
|
|
if repo_lower.startswith("mlx-community/") or _MLX_NAME_RE.search(source_repo_id):
|
|
logger.info(
|
|
f"Treating HF cache model as MLX-compatible by repo name: {source_repo_id}"
|
|
)
|
|
return True
|
|
|
|
if _is_helper_checkpoint(model_dir):
|
|
logger.info(
|
|
f"Treating HF cache model as MLX-compatible speculative-decoding "
|
|
f"helper: {source_repo_id}"
|
|
)
|
|
return True
|
|
|
|
logger.debug(f"Skipping non-MLX HF cache model: {source_repo_id}")
|
|
return False
|
|
|
|
|
|
def _register_model(
|
|
models: dict[str, DiscoveredModel],
|
|
model_dir: Path,
|
|
model_id: str,
|
|
*,
|
|
source_type: str = "local",
|
|
source_repo_id: str | None = None,
|
|
) -> None:
|
|
"""Try to register a single model directory into the models dict."""
|
|
if model_id in models:
|
|
logger.warning(
|
|
f"Duplicate model_id '{model_id}' found in {model_dir}, "
|
|
f"keeping version from {models[model_id].model_path}"
|
|
)
|
|
return
|
|
try:
|
|
if _is_unsupported_model(model_dir):
|
|
logger.info(f"Skipping unsupported model: {model_id}")
|
|
return
|
|
|
|
model_type = detect_model_type(model_dir)
|
|
if model_type == "embedding":
|
|
engine_type: EngineType = "embedding"
|
|
elif model_type == "reranker":
|
|
engine_type = "reranker"
|
|
elif model_type == "vlm":
|
|
engine_type = "vlm"
|
|
elif model_type == "audio_stt":
|
|
engine_type = "audio_stt"
|
|
elif model_type == "audio_tts":
|
|
engine_type = "audio_tts"
|
|
elif model_type == "audio_sts":
|
|
engine_type = "audio_sts"
|
|
else:
|
|
engine_type = "batched"
|
|
estimated_size = estimate_model_size(model_dir)
|
|
|
|
# Read raw config model_type for sub-type detection (e.g., OCR models)
|
|
# and flag speculative-decoding drafters (dFlash/Assistant/MTP).
|
|
config_model_type = ""
|
|
is_helper = False
|
|
try:
|
|
import json
|
|
with open(model_dir / "config.json") as f:
|
|
_config = json.load(f)
|
|
config_model_type = _config.get("model_type", "")
|
|
is_helper = is_helper_model_config(_config)
|
|
except Exception:
|
|
pass
|
|
|
|
thinking_default = detect_thinking_default(model_dir)
|
|
preserve_thinking_default = detect_preserve_thinking(model_dir)
|
|
model_context_length = _read_model_context_length(model_dir)
|
|
|
|
models[model_id] = DiscoveredModel(
|
|
model_id=model_id,
|
|
model_path=str(model_dir),
|
|
model_type=model_type,
|
|
engine_type=engine_type,
|
|
estimated_size=estimated_size,
|
|
config_model_type=config_model_type,
|
|
thinking_default=thinking_default,
|
|
preserve_thinking_default=preserve_thinking_default,
|
|
model_context_length=model_context_length,
|
|
source_type=source_type,
|
|
source_repo_id=source_repo_id,
|
|
is_helper=is_helper,
|
|
)
|
|
|
|
size_gb = estimated_size / (1024**3)
|
|
logger.info(
|
|
f"Discovered model: {model_id} "
|
|
f"(type: {model_type}, engine: {engine_type}, size: {size_gb:.2f}GB)"
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to discover model {model_id}: {e}")
|
|
|
|
|
|
def discover_models(model_dir: Path) -> dict[str, DiscoveredModel]:
|
|
"""
|
|
Scan model directory with two-level discovery.
|
|
|
|
Supports both flat and organized directory layouts:
|
|
|
|
Flat (one level):
|
|
model_dir/
|
|
├── llama-3b/ → model_id: "llama-3b"
|
|
│ ├── config.json
|
|
│ └── *.safetensors
|
|
└── qwen-7b/ → model_id: "qwen-7b"
|
|
|
|
Organized (two levels):
|
|
model_dir/
|
|
├── mlx-community/
|
|
│ ├── llama-3b/ → model_id: "llama-3b"
|
|
│ └── qwen-7b/ → model_id: "qwen-7b"
|
|
└── Qwen/
|
|
└── Qwen3-8B/ → model_id: "Qwen3-8B"
|
|
|
|
If a first-level subdirectory has config.json, it's treated as a model.
|
|
Otherwise, its children are scanned for models (organization folder).
|
|
|
|
Args:
|
|
model_dir: Path to directory containing model subdirectories
|
|
|
|
Returns:
|
|
Dictionary mapping model_id to DiscoveredModel
|
|
"""
|
|
access_error = model_directory_access_error(model_dir)
|
|
if access_error is not None:
|
|
if "not readable" in access_error:
|
|
logger.warning("Skipping directory %s: %s", model_dir, access_error)
|
|
return {}
|
|
raise ValueError(access_error)
|
|
|
|
models: dict[str, DiscoveredModel] = {}
|
|
|
|
for subdir in _iter_readable_entries(model_dir, "model directory"):
|
|
if not _is_readable_dir(subdir, "model entry") or subdir.name.startswith("."):
|
|
continue
|
|
|
|
if _is_adapter_dir(subdir):
|
|
logger.info(
|
|
f"Skipping LoRA adapter: {subdir.name} "
|
|
"(oMLX does not support LoRA/PEFT adapters)"
|
|
)
|
|
elif _is_model_dir(subdir):
|
|
# Level 1: direct model folder
|
|
_register_model(models, subdir, subdir.name)
|
|
else:
|
|
# HF Hub cache entry: models--Org--Name/snapshots/<hash>/
|
|
hf_resolved = _resolve_hf_cache_entry(subdir)
|
|
if hf_resolved is not None:
|
|
if _is_hf_cache_mlx_compatible(
|
|
hf_resolved.snapshot_path,
|
|
hf_resolved.source_repo_id,
|
|
):
|
|
_register_model(
|
|
models,
|
|
hf_resolved.snapshot_path,
|
|
hf_resolved.model_id,
|
|
source_type="hf_cache",
|
|
source_repo_id=hf_resolved.source_repo_id,
|
|
)
|
|
continue
|
|
|
|
# Level 2: organization folder — scan children
|
|
has_children = False
|
|
for child in _iter_readable_entries(subdir, "model group"):
|
|
if (
|
|
not _is_readable_dir(child, "model group entry")
|
|
or child.name.startswith(".")
|
|
):
|
|
continue
|
|
if _is_adapter_dir(child):
|
|
logger.info(
|
|
f"Skipping LoRA adapter: {child.name} "
|
|
"(oMLX does not support LoRA/PEFT adapters)"
|
|
)
|
|
elif _is_model_dir(child):
|
|
has_children = True
|
|
_register_model(models, child, child.name)
|
|
|
|
if not has_children:
|
|
logger.debug(
|
|
f"Skipping {subdir.name}: no config.json found "
|
|
f"(not a model or organization folder)"
|
|
)
|
|
|
|
# Fallback: if no models found and the directory itself is a model, register it.
|
|
# This supports pointing directly at a single model folder, e.g.:
|
|
# /Models/Qwen3.5-9B-MLX-4bit/ (contains config.json and weight files)
|
|
if not models and _is_model_dir(model_dir):
|
|
_register_model(models, model_dir, model_dir.name)
|
|
|
|
return models
|
|
|
|
|
|
def discover_models_from_dirs(
|
|
model_dirs: list[Path],
|
|
) -> dict[str, DiscoveredModel]:
|
|
"""
|
|
Scan multiple model directories and merge results.
|
|
|
|
Each directory is scanned using discover_models(). On model_id conflicts,
|
|
the first directory's model takes priority (earlier directory wins).
|
|
|
|
Args:
|
|
model_dirs: List of paths to directories containing model subdirectories
|
|
|
|
Returns:
|
|
Dictionary mapping model_id to DiscoveredModel
|
|
"""
|
|
merged: dict[str, DiscoveredModel] = {}
|
|
|
|
for model_dir in model_dirs:
|
|
access_error = model_directory_access_error(model_dir)
|
|
if access_error is not None:
|
|
logger.warning(f"Skipping directory {model_dir}: {access_error}")
|
|
continue
|
|
|
|
try:
|
|
discovered = discover_models(model_dir)
|
|
except ValueError as e:
|
|
logger.warning(f"Skipping directory {model_dir}: {e}")
|
|
continue
|
|
|
|
for model_id, info in discovered.items():
|
|
if model_id in merged:
|
|
logger.warning(
|
|
f"Duplicate model_id '{model_id}' found in {model_dir}, "
|
|
f"keeping version from {merged[model_id].model_path}"
|
|
)
|
|
continue
|
|
merged[model_id] = info
|
|
|
|
return merged
|
|
|
|
|
|
def format_size(size_bytes: int) -> str:
|
|
"""Format byte size as human-readable string."""
|
|
for unit in ["B", "KB", "MB", "GB", "TB"]:
|
|
if abs(size_bytes) < 1024.0:
|
|
return f"{size_bytes:.2f}{unit}"
|
|
size_bytes /= 1024.0
|
|
return f"{size_bytes:.2f}PB"
|