configs refactor 3/3: registry-based vision extractor hooks - #311
Merged
Conversation
_configs.py grew to 2591 lines (9 sub-configs, 19 per-model configs,
3 mega-extractor switch functions, a 530-line ArchitectureConfig.
from_transformers). This is the first of three mechanical refactors
that carve it into a scalable package layout. No behavior change in
this PR — every public name is still importable from mobius._configs.
src/mobius/_configs/
├── __init__.py # re-exports everything that was in _configs.py
├── _sub_configs.py # pure-data dataclasses (RoPE/Vision/Audio/Codec/TTS)
├── _quantization.py # QuantizationConfig + from_transformers
└── _base.py # BaseModelConfig, ArchitectureConfig, per-model
# subclasses, and the _extract_* helpers
Follow-up PRs in this series:
Part 2/3 — convert the _extract_audio_config / _extract_vision_config
model_type switches into a decorator-registered dispatch
so new models add a file under per_model/ instead of a
branch in the central function.
Part 3/3 — move per-model config subclasses (Gemma2Config,
MllamaConfig, NemotronHConfig, ...) into per_model/ and
carve up ArchitectureConfig.from_transformers.
Tests: 2769 passed (full src/ + tests/build_graph_test.py + cli_test.py).
Ruff: clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Replaces the 135-line model_type switch in _extract_audio_config with
a tiny plugin registry. Each model now contributes a hook in its own
file under src/mobius/_configs/per_model/, mutating an audio_fields
dict or short-circuiting with a fully-formed sub-config. Adding a new
audio-capable architecture no longer requires editing _base.py.
New module:
src/mobius/_configs/_extractors.py # register_audio_hook + dispatch
Per-model files (one per former branch):
per_model/_audio_default.py # audio_processor, embd_layer, speech_lora
per_model/_phi4mm_audio.py # phi4mm audio_token_id
per_model/_qwen3_asr_audio.py # thinker_config.audio_config + token ids
per_model/_gemma4_audio.py # short-circuit to Gemma4AudioConfig
per_model/_sensevoice_audio.py # encoder_conf + frontend_conf mapping
_extract_audio_config in _base.py shrinks to a 5-line shim that
triggers the per_model side-effect import and calls the dispatcher.
No behavior change: every existing audio-capable model still produces
the same AudioConfig from the same HF config.
Vision-side conversion (also a mega-switch) is intentionally deferred
to a follow-up so reviewers can verify the registry pattern on the
smaller surface first.
Tests:
pytest src/ tests/build_graph_test.py tests/cli_test.py -n auto
2769 passed, 41 skipped
pytest tests/arch_validation_test.py -k 'sensevoice_small or phi4mm
or qwen3_asr or gemma4'
15 passed
Ruff: clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Before: every hook body had to open with
if model_type != "phi4mm":
return None
That's noisy and easy to get wrong when copy-pasting a hook for a new
model. Make the decorator accept an optional filter and have the
dispatcher skip hooks whose filter doesn't match the current
model_type:
@register_audio_hook # runs for every model_type
def _default(...): ...
@register_audio_hook("phi4mm") # phi4mm only
def _phi4mm(...): ...
@register_audio_hook("gemma4", "gemma4_text")
def _gemma4_simple_case(...): ...
Hooks that also need to look at parent_config (Gemma4 audio,
Qwen3-ASR via thinker_config) keep the manual if-guard inside the
body, since the decorator filter can't reach beyond model_type.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
force-pushed
the
refactor-configs-3-vision-extractor-registry
branch
from
May 20, 2026 14:20
54b6726 to
4f30881
Compare
Two concerns covered: 1. Mechanism — register_audio_hook works as a bare decorator and as a parameterised factory; the dispatcher applies model_type filters; a hook can short-circuit by returning a dict. 2. Cross-contamination — every per-model audio hook is verified to NOT fire for unrelated model_types. For filtered hooks (phi4mm, sensevoice), the dispatcher's filter must skip non-matching types even when the input config has matching shape (e.g. an audio_config dict that would normally trigger phi4mm). Bare hooks (default, qwen3_asr, gemma4) must return cleanly on a vanilla text-model config. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Mirrors part 2/3 for the vision side: replaces the 185-line
_extract_vision_config switch with a per-model hook registry. The
vision dispatcher carries a bit more logic than audio because vision
also lifts a fixed set of shared fields (image_token_id, mrope_section,
spatial_merge_size, ...) up to the top-level of the returned dict.
New per_model files:
_vision_default.py # canonical HF vision_config + shared bits
_phi4mm_vision.py # hard-coded SigLIP encoder dims
_hunyuan_vl_mot_vision.py # InternViT-style ViT, flat config.json
_internvl_vision.py # default image_token_id for InternVL chain
After this PR, every model-type switch in the original
_extract_*_config functions has been moved into discoverable per-model
files. New audio- or vision-capable models add a single file under
src/mobius/_configs/per_model/ instead of editing _base.py.
Tests:
pytest src/ tests/build_graph_test.py tests/cli_test.py -n auto
2769 passed, 41 skipped
pytest tests/arch_validation_test.py -k 'phi4mm or hunyuan_vl_mot
or qwen2_vl or qwen3_vl
or gemma4'
30 passed
Ruff: clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
…) filters Apply the same decorator-filter pattern PR2 introduced for audio. phi4mm and hunyuan_vl_mot match cleanly on model_type alone, so they move to the parameterised form and drop the manual model_type guard. The internvl hook also matches on parent_config.model_type (composite configs whose top-level model_type is "internvl_chat" but whose text-config model_type is something else), so it stays as a bare @register_vision_hook + body predicate; same rationale as gemma4_audio in PR2. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
* test_extractors_test now covers vision hooks too — phi4mm and
hunyuan_vl_mot must not cross-fire for unrelated model_types, and
plain text architectures must produce no vision output.
* adding-a-new-model SKILL.md previously told contributors to edit
_configs.py for new audio/vision fields. After the registry
refactor, the canonical path for audio/vision config is to add a
per-model file under src/mobius/_configs/per_model/ with
@register_vision_hook("type") / @register_audio_hook("type"). The
skill now points there and mentions src/mobius/_configs/_base.py
for non-modality top-level config edits.
* multimodal-models projector-variants.md updates the VisionConfig
location to src/mobius/_configs/_sub_configs.py.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
force-pushed
the
refactor-configs-3-vision-extractor-registry
branch
from
May 20, 2026 14:26
4f30881 to
c6d15c3
Compare
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Refactors vision config extraction to use the same per-model hook registry pattern introduced for audio (PR #310), replacing the large _extract_vision_config switch with discoverable, per-model hook modules under src/mobius/_configs/per_model/. This keeps _configs/_base.py slim and makes adding new vision-capable architectures a “drop in a file + import it” change.
Changes:
- Added vision hook registry dispatch (
extract_vision_config) that assembles aVisionConfigand lifts shared fields (e.g.,image_token_id) to the top-level. - Introduced per-model vision hook modules (
_vision_default,_phi4mm_vision,_internvl_vision,_hunyuan_vl_mot_vision) and wired them intoper_model/__init__.pyfor side-effect registration. - Extended extractor registry unit tests to cover vision hooks and updated internal skills docs to point to the new locations/pattern.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/_configs/per_model/_vision_default.py | Adds the default “HF vision_config + shared fields” extractor hook. |
| src/mobius/_configs/per_model/_phi4mm_vision.py | Adds Phi4MM-specific SigLIP hard-coded vision parameters hook. |
| src/mobius/_configs/per_model/_internvl_vision.py | Adds InternVL defaulting for image_token_id when missing. |
| src/mobius/_configs/per_model/_hunyuan_vl_mot_vision.py | Adds HunYuan VL-MoT flat-config hard-coded vision parameters hook. |
| src/mobius/_configs/per_model/init.py | Imports new per-model vision modules to register hooks at import time. |
| src/mobius/_configs/_extractors.py | Implements extract_vision_config() dispatcher that builds VisionConfig and lifts shared fields. |
| src/mobius/_configs/_extractors_test.py | Adds unit tests ensuring vision hooks don’t fire for unrelated model types. |
| src/mobius/_configs/_base.py | Replaces _extract_vision_config mega-switch with a shim delegating to the registry dispatcher. |
| .agents/skills/multimodal-models/references/projector-variants.md | Updates doc reference to VisionConfig dataclass location. |
| .agents/skills/adding-a-new-model/SKILL.md | Updates “adding a model” guidance to prefer per-model audio/vision hooks. |
justinchuby
added a commit
that referenced
this pull request
May 21, 2026
…el fields Three review comments from copilot-pull-request-reviewer on PR #311: 1. (_vision_default.py) Unconditional fields["image_token_id"] / fields[ "mm_tokens_per_image"] assignments could clobber values that a per-model hook set. Switched to setdefault semantics so the default only fills in a value when no per-model hook has supplied one. 2. (per_model/__init__.py) Default hooks were imported in alphabetical order, putting _vision_default last. That meant _vision_default ran last and overwrote per-model image_token_id (phi4mm / hunyuan_vl_mot). Reordered so _audio_default / _vision_default import (and therefore register and run) first, with fmt/ruff overrides to keep the ordering stable. 3. (_extractors_test.py) Added regression tests asserting that phi4mm and hunyuan_vl_mot image_token_id survive the full vision-hook pipeline, both on the VisionConfig dataclass and as the lifted top-level field. Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
…el fields Three review comments from copilot-pull-request-reviewer on PR #311: 1. (_vision_default.py) Unconditional fields["image_token_id"] / fields[ "mm_tokens_per_image"] assignments could clobber values that a per-model hook set. Switched to setdefault semantics so the default only fills in a value when no per-model hook has supplied one. 2. (per_model/__init__.py) Default hooks were imported in alphabetical order, putting _vision_default last. That meant _vision_default ran last and overwrote per-model image_token_id (phi4mm / hunyuan_vl_mot). Reordered so _audio_default / _vision_default import (and therefore register and run) first, with fmt/ruff overrides to keep the ordering stable. 3. (_extractors_test.py) Added regression tests asserting that phi4mm and hunyuan_vl_mot image_token_id survive the full vision-hook pipeline, both on the VisionConfig dataclass and as the lifted top-level field. Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
refactor-configs-3-vision-extractor-registry
branch
from
May 21, 2026 02:22
b1247cf to
cec7274
Compare
Member
Author
|
Reliance on import order is brittle. We should find a more robust way of doing this. |
Per review on PR #311: the previous fix-up wrapped the per_model package __init__ in 'fmt: off' / 'ruff: noqa: I001' to keep _audio_default and _vision_default at the top of the alphabetical import block — that sequence was the only thing ensuring default hooks ran before per-model overrides and therefore did not clobber fields like image_token_id. That's too fragile (anyone running isort or hand-editing the imports without noticing the magic comment puts us back in the original bug). Replace import-order reliance with an explicit 'priority' kwarg on the @register_*_hook decorators: * DEFAULT_PRIORITY (0) — defaults that fill in HF fields * PER_MODEL_PRIORITY (100) — per-model hooks that may override The dispatcher sorts (priority, insertion_index) before invoking, so: - defaults always run before per-model regardless of import order - equal priority preserves insertion order (stable sort) - import order in per_model/__init__.py is intentionally irrelevant and ruff/isort may freely re-sort the block Updated: * _extractors._make_register, _run: priority + index + sort * per_model/_vision_default, per_model/_audio_default: priority=DEFAULT_PRIORITY * per_model/__init__.py: alphabetical, no more fmt-off / noqa comments * tests: - update tuple shape assertion to (priority, idx, filter, fn) - add test_priority_overrides_registration_order - add test_equal_priority_preserves_registration_order Signed-off-by: justinchuby <11205safetensors048+justinchuby@users.noreply.github.com>
Member
Author
|
You're right — relying on import order (and trying to lock it with fmt: off / ruff: noqa: I001) is too fragile. Anyone re-sorting the block without noticing the magic comment puts us back in the original bug. Replaced with explicit priority in commit 7b34f4b. The dispatcher sorts hooks by (priority, insertion_index) before iterating, so defaults always run before per-model overrides regardless of import order; equal priority preserves registration order via stable sort; per_model/init.py is back to plain alphabetical imports so ruff/isort can freely re-sort it. Added test_priority_overrides_registration_order and test_equal_priority_preserves_registration_order as regression guards. |
…nery
Per review feedback: explicit priority on @register_*_hook was already
better than relying on import order, but still over-engineered for what
the system actually needs (one always-runs first pass + per-model
overrides). Replace with a simpler design that puts the run order
visibly at the call site.
Changes:
* Move 'always-runs default' logic out of the hook registry. The
previous _audio_default.py / _vision_default.py modules become
_configs/_audio_defaults.py and _configs/_vision_defaults.py at the
package root (not under per_model/), exposing
apply_audio_defaults / apply_vision_defaults plain functions.
extract_audio_config and extract_vision_config now call these
explicitly as the pipeline's first step, then dispatch to the
per-model hook registry:
def extract_vision_config(...):
fields = {}
apply_vision_defaults(config, parent_config, model_type, fields)
... # then run per-model hooks
This makes the run order self-evident at the call site, kills any
dependence on import order, and removes the priority + insertion
index machinery entirely.
* per_model/__init__.py drops the _audio_default / _vision_default
imports and goes back to plain alphabetical (no fmt-off /
noqa: I001).
* Drop DEFAULT_PRIORITY / PER_MODEL_PRIORITY constants and the priority
kwarg on _make_register. The registry is back to
list[tuple[frozenset[str] | None, Hook]] and _run iterates in
registration order.
* Why the defaults had to move OUT of per_model/: the dispatcher uses
a lazy 'from mobius._configs.per_model._vision_default import ...'
inside extract_vision_config. That import triggers
per_model/__init__.py side-effects, which during an _isolated_registry
test (where register_*_hook is monkeypatched) permanently registers
the per-model hooks into the monkeypatched empty list. After undo,
the original _AUDIO_HOOKS / _VISION_HOOKS are still empty, and
subsequent 'from mobius._configs import per_model' is a cache
no-op so the per-model hooks never re-register. Loading the defaults
from a sibling module (not from per_model/) avoids triggering the
per-model package init.
* Add test_defaults_run_before_per_model_hooks regression guard. Drop
the priority-specific tests (test_priority_overrides_registration_order
and test_equal_priority_preserves_registration_order) — no longer
applicable.
Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
This comment was marked as resolved.
This comment was marked as resolved.
…sion-extractor-registry
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part 3 of 3 — stacked on #310
Mirrors part 2/3 for the vision side: replaces the 185-line
_extract_vision_configswitch with a per-model hook registry. The vision dispatcher carries slightly more logic than audio because vision also lifts a fixed set of shared fields (image_token_id,mrope_section,spatial_merge_size, …) up to the top-level of the returned dict.New per_model files
_vision_default.pyvision_config+ LoRA / embd_layer / mrope_section_phi4mm_vision.py_hunyuan_vl_mot_vision.py_internvl_vision.pyimage_token_idfor InternVL chainEnd state
After this 3-PR series, every model-type switch in the original
_extract_*_configfunctions has been moved into discoverable per-model files. New audio- or vision-capable models add a single file undersrc/mobius/_configs/per_model/instead of editing_base.py.The follow-on refactor (out of scope here) is to do the same for per-model config subclasses (
Gemma2Config,MllamaConfig, etc.) and to carve upArchitectureConfig.from_transformersitself.Tests
pytest src/ tests/build_graph_test.py tests/cli_test.py -n autopytest tests/arch_validation_test.py -k 'phi4mm or hunyuan_vl_mot or qwen2_vl or qwen3_vl or gemma4'