From 2b15cc9d17d785f3ca2b4fff7bef14af0f0ae99d Mon Sep 17 00:00:00 2001
From: xiaoyuzhang <xiaoyuzhang@microsoft.com>
Date: Fri, 29 May 2026 16:23:50 -0700
Subject: [PATCH 1/2] introspection: add list_components for lightweight
 component discovery

Adds mobius.list_components(model_id, *, task=None, trust_remote_code=False) -> list[str] which returns the ONNX component names that mobius.build() would produce, without downloading weights or constructing any graphs.

For transformer models it downloads only the HuggingFace config.json (via transformers.AutoConfig) and dispatches via the same registry + task lookup as build(), including the qwen3_5_moe -> qwen3_5_moe_vl override. For diffusers pipelines it downloads only model_index.json and mirrors the iteration, filtering, and sub-package flattening rules of build_diffusers_pipeline so the returned names exactly match the keys of the resulting ModelPackage.

Component names are read statically from each task's ComponentSpec (or default to ['model'] for single-component tasks).

Tests cover: per-task component lookup, composite-config unwrapping, transformer + diffusers paths, explicit task override, and graceful errors when neither config.json nor model_index.json are usable.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/mobius/__init__.py       |   2 +
 src/mobius/_introspection.py | 197 +++++++++++++++++++++++++++
 tests/introspection_test.py  | 250 +++++++++++++++++++++++++++++++++++
 3 files changed, 449 insertions(+)
 create mode 100644 src/mobius/_introspection.py
 create mode 100644 tests/introspection_test.py

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..daab825d
--- /dev/null
+++ b/src/mobius/_introspection.py
@@ -0,0 +1,197 @@
+# 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",
+]
+
+import logging
+
+from mobius._registry import _detect_fallback_registration, registry
+from mobius.tasks import ModelTask, get_task
+
+logger = logging.getLogger(__name__)
+
+
+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__

From 4410e02ccdf699e607fb3d963a3a66c80b4c2657 Mon Sep 17 00:00:00 2001
From: Xiaoyu <85524621+xiaoyu-work@users.noreply.github.com>
Date: Fri, 29 May 2026 16:46:37 -0700
Subject: [PATCH 2/2] Potential fix for pull request finding 'Unused global
 variable'

Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
Signed-off-by: Xiaoyu <85524621+xiaoyu-work@users.noreply.github.com>
---
 src/mobius/_introspection.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/mobius/_introspection.py b/src/mobius/_introspection.py
index daab825d..0519222b 100644
--- a/src/mobius/_introspection.py
+++ b/src/mobius/_introspection.py
@@ -27,13 +27,9 @@
     "list_components",
 ]
 
-import logging
-
 from mobius._registry import _detect_fallback_registration, registry
 from mobius.tasks import ModelTask, get_task
 
-logger = logging.getLogger(__name__)
-
 
 def list_components(
     model_id: str,