diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 513fb13b..9d681b58 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -39,6 +39,7 @@ "ep_registry", "get_build_dtype", "get_ep", + "list_components", "models", "optimize_model", "register_ep", @@ -77,6 +78,7 @@ from mobius._constants import OPSET_VERSION from mobius._diffusers_builder import build_diffusers_pipeline from mobius._execution_providers import EpCapabilities, ep_registry, get_ep, register_ep +from mobius._introspection import list_components from mobius._model_package import ModelPackage from mobius._optimizations import optimize_model from mobius._registry import ( diff --git a/src/mobius/_introspection.py b/src/mobius/_introspection.py new file mode 100644 index 00000000..0519222b --- /dev/null +++ b/src/mobius/_introspection.py @@ -0,0 +1,193 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Lightweight introspection of mobius-buildable models. + +These helpers answer questions such as *"which ONNX component names will +``mobius.build()`` produce for this model?"* without actually building the +graph or downloading the model weights. + +They are useful for tooling that needs to plan or schedule work on a per +component basis before paying the cost of a full build — for example, an +optimization pipeline that wants to assign different passes to the +``encoder`` and ``decoder`` components of an encoder-decoder model. + +Only metadata files are downloaded: + +- For transformer models: HuggingFace ``config.json`` (via + ``transformers.AutoConfig``). +- For diffusers pipelines: ``model_index.json``. + +No PyTorch weights are downloaded and no ONNX graphs are constructed. +""" + +from __future__ import annotations + +__all__ = [ + "list_components", +] + +from mobius._registry import _detect_fallback_registration, registry +from mobius.tasks import ModelTask, get_task + + +def list_components( + model_id: str, + *, + task: str | ModelTask | None = None, + trust_remote_code: bool = False, +) -> list[str]: + """Return the ONNX component names :func:`mobius.build` would produce. + + The result matches the keys of the :class:`~mobius.ModelPackage` + returned by ``build(model_id)``. For single-component models this is + ``["model"]``; for multi-component models (encoder-decoder, VAE, + diffusers pipelines, multimodal) it is the list of sub-model names. + + Only HuggingFace metadata files (``config.json`` or ``model_index.json``) + are downloaded; no weights are fetched and no graphs are built. The + cost is roughly one HTTP request. + + Args: + model_id: HuggingFace model repository ID (e.g. + ``"meta-llama/Llama-3-8B"`` or + ``"stabilityai/stable-diffusion-3-medium-diffusers"``). + task: Override the auto-detected task. Either a task name string + (e.g. ``"text-generation"``) or a :class:`ModelTask` instance. + Ignored for diffusers pipelines. + trust_remote_code: Whether to trust remote code when loading the + HuggingFace config. Forwarded to ``transformers.AutoConfig``. + + Returns: + A list of component names. Order matches the order + :func:`mobius.build` produces (sorted by spec for transformer + models, by ``model_index.json`` iteration order for diffusers). + + Raises: + ValueError: If *model_id* is neither a supported transformer model + nor a diffusers pipeline. + + Example:: + + >>> import mobius + >>> mobius.list_components("meta-llama/Llama-3-8B") + ['model'] + >>> mobius.list_components("stabilityai/stable-diffusion-3-medium-diffusers") + ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'transformer', 'vae'] + """ + try: + import transformers + except ImportError as e: + raise ImportError( + "transformers is required for list_components(); " + "install with `pip install transformers`." + ) from e + + from mobius._config_resolver import _default_task_for_model, _try_load_config_json + + try: + hf_config = transformers.AutoConfig.from_pretrained( + model_id, trust_remote_code=trust_remote_code + ) + except (ValueError, KeyError, OSError): + hf_config = _try_load_config_json(model_id) + if hf_config is None or hf_config.model_type not in registry: + return _list_diffusers_components(model_id) + + model_type = _resolve_effective_model_type(hf_config) + + if model_type in registry: + if task is None: + task = _default_task_for_model(model_type) + return _components_for_task(task) + + fallback = _detect_fallback_registration(hf_config) + if fallback is not None: + if task is None: + task = fallback.task or "text-generation" + return _components_for_task(task) + + return _list_diffusers_components(model_id) + + +def _resolve_effective_model_type(hf_config) -> str: + """Return the model_type :func:`mobius.build` would actually dispatch on. + + Mirrors the composite-config unwrapping in :func:`mobius._builder.build`: + for configs that wrap a text-only sub-config (``talker_config``, + ``thinker_config``, ``text_config``), the dispatched model_type may + differ from ``hf_config.model_type``. This is currently meaningful only + for the ``qwen3_5_moe`` → ``qwen3_5_moe_vl`` override, but mirroring + the full logic keeps the two paths in sync as ``build()`` evolves. + """ + model_type = hf_config.model_type + if hasattr(hf_config, "talker_config") or hasattr(hf_config, "thinker_config"): + return model_type + if hasattr(hf_config, "text_config"): + if ( + model_type == "qwen3_5_moe" + and getattr(hf_config, "vision_config", None) is not None + ): + return "qwen3_5_moe_vl" + return model_type + + +def _components_for_task(task: str | ModelTask) -> list[str]: + """Return the component names declared by *task*. + + Single-component tasks (with no :class:`ComponentSpec`) always + produce a single ``"model"`` entry. + """ + resolved = get_task(task) + if resolved.components is None: + return ["model"] + return list(resolved.components.keys()) + + +def _list_diffusers_components(model_id: str) -> list[str]: + """List the component names :func:`build_diffusers_pipeline` would emit. + + Mirrors the iteration and flattening logic in + :func:`mobius._diffusers_builder.build_diffusers_pipeline` so that the + returned names exactly match the keys of the :class:`ModelPackage` it + would produce. Components that ``build_diffusers_pipeline`` would skip + (private entries, non-tuple values, unregistered classes) are skipped + here too. + """ + from mobius._diffusers_builder import ( + _DIFFUSERS_CLASS_MAP, + _init_diffusers_class_map, + _load_diffusers_pipeline_index, + ) + + pipeline_index = _load_diffusers_pipeline_index(model_id) + if pipeline_index is None: + raise ValueError( + f"'{model_id}' is not a supported transformer model and is not " + f"a diffusers pipeline (no model_index.json found)." + ) + + _init_diffusers_class_map() + + components: list[str] = [] + for component_name, component_info in pipeline_index.items(): + if component_name.startswith("_"): + continue + if not isinstance(component_info, list) or len(component_info) != 2: + continue + + _, class_name = component_info + if class_name not in _DIFFUSERS_CLASS_MAP: + continue + + _, _, task_name = _DIFFUSERS_CLASS_MAP[class_name] + sub_components = _components_for_task(task_name) + + if sub_components == ["model"]: + components.append(component_name) + else: + components.extend(f"{component_name}_{sub_name}" for sub_name in sub_components) + + if not components: + raise ValueError(f"No supported components found in diffusers pipeline '{model_id}'.") + return components diff --git a/tests/introspection_test.py b/tests/introspection_test.py new file mode 100644 index 00000000..484063dc --- /dev/null +++ b/tests/introspection_test.py @@ -0,0 +1,250 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Unit tests for :func:`mobius.list_components`. + +These tests mock the network calls so they run offline. They verify the +component-name discovery logic against the real ``mobius`` task registry, +diffusers class map, and component specs. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import ClassVar +from unittest import mock + +import pytest + +import mobius +from mobius._introspection import ( + _components_for_task, + _resolve_effective_model_type, + list_components, +) + + +class TestComponentsForTask: + """``_components_for_task`` returns the names declared by each task.""" + + def test_single_component_task_returns_model(self): + assert _components_for_task("text-generation") == ["model"] + + def test_seq2seq_returns_encoder_and_decoder(self): + assert sorted(_components_for_task("seq2seq")) == ["decoder", "encoder"] + + def test_vae_returns_encoder_and_decoder(self): + assert sorted(_components_for_task("vae")) == ["decoder", "encoder"] + + def test_multimodal_task_returns_all_components(self): + # Phi4MM has vision_encoder, audio_encoder, embedding, decoder + names = set(_components_for_task("phi4mm-multimodal")) + assert {"vision_encoder", "audio_encoder", "embedding", "decoder"} <= names + + def test_unknown_task_raises(self): + with pytest.raises(ValueError, match="Unknown task"): + _components_for_task("does-not-exist") + + +class TestResolveEffectiveModelType: + """``_resolve_effective_model_type`` mirrors ``build()`` dispatch.""" + + def test_plain_model_type_returned_as_is(self): + cfg = SimpleNamespace(model_type="llama") + assert _resolve_effective_model_type(cfg) == "llama" + + def test_talker_config_does_not_override_model_type(self): + cfg = SimpleNamespace(model_type="qwen2_5_omni", talker_config=object()) + assert _resolve_effective_model_type(cfg) == "qwen2_5_omni" + + def test_qwen3_5_moe_with_vision_overrides_to_vl(self): + cfg = SimpleNamespace( + model_type="qwen3_5_moe", + text_config=object(), + vision_config=object(), + ) + assert _resolve_effective_model_type(cfg) == "qwen3_5_moe_vl" + + def test_qwen3_5_moe_without_vision_stays_text_only(self): + cfg = SimpleNamespace( + model_type="qwen3_5_moe", + text_config=object(), + vision_config=None, + ) + assert _resolve_effective_model_type(cfg) == "qwen3_5_moe" + + +class TestListComponentsTransformer: + """Transformer-model code path of ``list_components``.""" + + def _patch_autoconfig(self, hf_config): + return mock.patch( + "transformers.AutoConfig.from_pretrained", + return_value=hf_config, + ) + + def test_causal_lm_returns_single_model(self): + cfg = SimpleNamespace(model_type="llama") + with self._patch_autoconfig(cfg): + assert list_components("any/llama-id") == ["model"] + + def test_seq2seq_returns_encoder_decoder(self): + cfg = SimpleNamespace(model_type="t5") + with self._patch_autoconfig(cfg): + assert sorted(list_components("any/t5-id")) == ["decoder", "encoder"] + + def test_explicit_task_overrides_default(self): + cfg = SimpleNamespace(model_type="llama") + with self._patch_autoconfig(cfg): + assert set(list_components("any/llama-id", task="seq2seq")) == { + "encoder", + "decoder", + } + + def test_qwen3_5_moe_with_vision_returns_vl_components(self): + cfg = SimpleNamespace( + model_type="qwen3_5_moe", + text_config=object(), + vision_config=object(), + ) + with self._patch_autoconfig(cfg): + names = set(list_components("any/qwen-vl-id")) + assert {"decoder", "vision_encoder", "embedding"} <= names + + def test_unsupported_model_type_falls_back_to_diffusers(self): + # AutoConfig succeeds with a model_type that isn't in registry + # and isn't a CausalLM fallback candidate; we should attempt the + # diffusers path (which then raises since model_index.json is + # absent). + cfg = SimpleNamespace( + model_type="something-unsupported", + is_encoder_decoder=False, + architectures=[], + ) + with ( + self._patch_autoconfig(cfg), + mock.patch( + "mobius._diffusers_builder._load_diffusers_pipeline_index", + return_value=None, + ), + pytest.raises(ValueError, match="not a diffusers pipeline"), + ): + list_components("any/unknown-id") + + +class TestListComponentsDiffusers: + """Diffusers-pipeline code path of ``list_components``.""" + + SD3_INDEX: ClassVar[dict] = { + "_class_name": "StableDiffusion3Pipeline", + "_diffusers_version": "0.30.0", + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + "text_encoder": ["transformers", "CLIPTextModelWithProjection"], + "text_encoder_2": ["transformers", "CLIPTextModelWithProjection"], + "text_encoder_3": ["transformers", "T5EncoderModel"], + "tokenizer": ["transformers", "CLIPTokenizer"], + "tokenizer_2": ["transformers", "CLIPTokenizer"], + "tokenizer_3": ["transformers", "T5TokenizerFast"], + "transformer": ["diffusers", "SD3Transformer2DModel"], + "vae": ["diffusers", "AutoencoderKL"], + } + + def test_sd3_pipeline_returns_supported_components(self): + with ( + mock.patch( + "transformers.AutoConfig.from_pretrained", + side_effect=OSError("no config.json"), + ), + mock.patch( + "mobius._config_resolver._try_load_config_json", + return_value=None, + ), + mock.patch( + "mobius._diffusers_builder._load_diffusers_pipeline_index", + return_value=self.SD3_INDEX, + ), + ): + names = list_components("stabilityai/sd3-mock") + + # text_encoders and tokenizers from transformers are not in + # _DIFFUSERS_CLASS_MAP and therefore skipped; only the transformer + # and vae (single-component each) remain. + assert "transformer" in names + # VAE task is multi-component (encoder/decoder) so it gets flattened + # to vae_encoder + vae_decoder. + assert {"vae_encoder", "vae_decoder"} <= set(names) + assert "text_encoder" not in names # CLIPTextModelWithProjection not registered + assert "scheduler" not in names + + def test_diffusers_skips_private_entries(self): + index = { + "_class_name": "TestPipeline", + "_skip_me": ["x", "y"], + "transformer": ["diffusers", "SD3Transformer2DModel"], + } + with ( + mock.patch( + "transformers.AutoConfig.from_pretrained", + side_effect=OSError("no config.json"), + ), + mock.patch( + "mobius._config_resolver._try_load_config_json", + return_value=None, + ), + mock.patch( + "mobius._diffusers_builder._load_diffusers_pipeline_index", + return_value=index, + ), + ): + assert list_components("test/mock") == ["transformer"] + + def test_diffusers_pipeline_with_no_supported_components_raises(self): + index = { + "_class_name": "EmptyPipeline", + "scheduler": ["diffusers", "DDIMScheduler"], + "tokenizer": ["transformers", "CLIPTokenizer"], + } + with ( + mock.patch( + "transformers.AutoConfig.from_pretrained", + side_effect=OSError("no config.json"), + ), + mock.patch( + "mobius._config_resolver._try_load_config_json", + return_value=None, + ), + mock.patch( + "mobius._diffusers_builder._load_diffusers_pipeline_index", + return_value=index, + ), + pytest.raises(ValueError, match="No supported components"), + ): + list_components("test/empty") + + def test_no_model_index_and_not_a_supported_model_raises(self): + with ( + mock.patch( + "transformers.AutoConfig.from_pretrained", + side_effect=OSError("no config.json"), + ), + mock.patch( + "mobius._config_resolver._try_load_config_json", + return_value=None, + ), + mock.patch( + "mobius._diffusers_builder._load_diffusers_pipeline_index", + return_value=None, + ), + pytest.raises(ValueError, match="not a diffusers pipeline"), + ): + list_components("nonexistent/model") + + +class TestPublicAPI: + """``list_components`` is exported at the package root.""" + + def test_exposed_from_top_level(self): + assert mobius.list_components is list_components + + def test_in_dunder_all(self): + assert "list_components" in mobius.__all__