[Feat] Support Pipeline Resolution Based on HF Config & Make Pipeline Registry Eager

[alex-jw-brooks]

Purpose

FIX #3757

While most models are consistent about their usage of stages within a pipeline, some models may only use a subset of the stages, and will run into issues if we try to initialize all of the stages. Such an example is Qwen/Qwen3-Omni-30B-A3B-Instruct, which should be thinker only, but breaks when trying to initialize the talker/code2wav.

To this end, we need a mechanism for resolving the correct pipeline configuration given the HF config, e.g., for the case of Qwen3Omni, selecting the thinker only configuration if enable_audio_output is False, or all stages if enable_audio_output is True.

Changes in this PR:

• More general / architectural changes:

  • Removes lazy initialization from the Diffusion registry. This was written to follow the patterns in vLLM, but since we are just mapping to pipeline defs, we can eagerly declare without heavy imports. This is simpler and more explicit with typing, which is important now that we may map to pipeline configs or resolver funcs.
  • Standardizes empty .__init__s across diffusion pipelines to avoid unintentional heavy imports, since things declared here don't seem to be used anyway
  • Support mapping pipeline names -> modules or resolver funcs, where if provided, the resolver func takes a single arg (HF config) and returns the PipelineConfig if applicable, or None.

  Since resolvers should be 1:1 with HF Transformer config types, they should be super simple and wrapped in a decorator that declares the type they are attached to; the decorator just ensures that None is returned (i.e., no pipeline config) if the type is not the same. Example:

# This resolver will return `None` if we pass it something that's not a Qwen3OmniMoeConfig
@pipeline_cfg_resolver(config_type=Qwen3OmniMoeConfig)
def resolve_qwen3_omni_pipeline(
    hf_config: Qwen3OmniMoeConfig,
) -> PipelineConfig:
    # If we have a config and it explicitly disabled audio input, load thinker only
    if not hf_config.enable_audio_output:
        return QWEN3_OMNI_THINKER_ONLY_PIPELINE
    return QWEN3_OMNI_PIPELINE

Minor / model specific changes:

• Disables async_chunk for single stage models; currently it defaults to True and breaks the single stage model, but this doesn't make sense for single stage scenarios
• Removes unused spatial merge / rope deltas from the talker, which were throwing in the wrongly initialized talker in the linked issue and don't appear to be used
• Some refactoring, e.g., moving the stage config factory to its own file; this is mostly needed for avoiding circular imports, but is also a little cleaner since the stage config file is messy

[hsliuustc0106]

BLOCKING:

• Test Coverage — This PR adds new pipeline resolution logic but doesn't include any tests:
  1. Test for the new function (both paths: enable_audio_output=True/False)
  2. Test for the method with both callable and PipelineConfig objects
  3. Test for the behavior on single-stage pipelines
  4. Test that Qwen3-Omni-30B-A3B-Captioner loads successfully as thinker-only

Please add tests to verify the new resolver logic works correctly.

[hsliuustc0106]

BLOCKING: Missing tests for new pipeline resolution logic. See inline comments.

[Gaohan123]

@alex-jw-brooks Hello, any updates?

[alex-jw-brooks]

Hey @Gaohan123, I have been a bit focused on other threads, sorry!

This PR is now pretty large since it's also cleaning up the pipeline registry and removing a lot of dead code. I will rebase this PR, update the description to be more accurate, and move it out of draft tomorrow

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[alex-jw-brooks]

Hi @hsliuustc0106 @Gaohan123! This PR ended up being larger than I originally expected since it took a bit of refactoring to make the registry eager, but I think it is finally ready for a review when you have time 🙏 I've updated the body with the main changes

[hsliuustc0106]

lgtm