fix: NextStep pipeline class in config fallback path

[Joshna-Medisetty]

Fix: NextStep pipeline class in fallback path (omni_diffusion.py)

Summary

Fix NextStep-1.1 failing with ValueError: Unknown model type: nextstep, architectures: ['LlamaForCausalLM'] when the fallback path is used (no model_index.json). The fix is to always set pipeline_class = "NextStep11Pipeline" when model_type == "nextstep", and remove the incorrect if od_config.model_class_name is None: guard around it.

Why the original line was added

In PR #612 ([Model] Add new nextstep_1 (Diffusion) model (only T2I)), the NextStep branch was added with:

elif model_type == "nextstep":
    if od_config.model_class_name is None:
        od_config.model_class_name = "NextStep11Pipeline"

The intent was to respect caller override: the CLI (text_to_image.py) explicitly sets omni_kwargs["model_class_name"] = "NextStep11Pipeline" for NextStep. The author wanted “only set the pipeline class when the caller hasn’t already set it,” so they guarded the assignment with if od_config.model_class_name is None.

Why that line is wrong in the current code

The fallback path was later refactored to a single pattern:

1. Each branch (bagel, nextstep, glm-image, …) sets only a local variable pipeline_class.
2. Once at the end, we do: if od_config.model_class_name is None: od_config.model_class_name = pipeline_class.

So “don’t overwrite when the caller already set it” is already enforced in that one place. In the refactored flow, the nextstep branch must always set pipeline_class = "NextStep11Pipeline" so that (a) we don’t raise “Unknown model type” when pipeline_class is None, and (b) the end block can assign it to od_config.model_class_name when the caller did not set it.

elif model_type == "nextstep":
    if od_config.model_class_name is None: 
        pipeline_class = "NextStep11Pipeline"

That is wrong here because:

• When the caller does set model_class_name (e.g. from the CLI), we do not set pipeline_class, so pipeline_class stays None, and we hit if pipeline_class is None: raise ValueError(...).
• The “don’t overwrite” behavior is already handled at the end; the branch should not tie setting pipeline_class to whether the caller set model_class_name.

Why we don’t need the guard anymore

• Before the refactor: The nextstep branch could correctly do “if None, then set”: if od_config.model_class_name is None: od_config.model_class_name = "NextStep11Pipeline".
• After the refactor: The branch should only set the local variable: pipeline_class = "NextStep11Pipeline". The shared block at the end (if od_config.model_class_name is None: od_config.model_class_name = pipeline_class) already implements “only set when the caller didn’t set it” for all model types, including NextStep. So we do not need the if od_config.model_class_name is None: guard in the nextstep branch; we only need to set pipeline_class unconditionally.

Change

• Before: if od_config.model_class_name is None: pipeline_class = "NextStep11Pipeline"
• After: pipeline_class = "NextStep11Pipeline"

Behavior after the fix:

• Caller sets model_class_name → we don’t overwrite (end block skips assignment).
• Caller does not set it → we set it from pipeline_class (end block assigns).

[Joshna-Medisetty]

cc @xuechendi

[hsliuustc0106]

any ci test for this bug?

[xuechendi]

@Joshna-Medisetty , may you help to add a new test file in https://github.com/vllm-project/vllm-omni/tree/main/tests/e2e/offline_inference ?

[Joshna-Medisetty]

@hsliuustc0106
Test added: NextStep has no model_index.json, so loading uses the config.json fallback. When model_class_name="NextStep11Pipeline" is set up front (same as the offline text_to_image CLI), the buggy path only assigned pipeline_class if model_class_name was None, leaving pipeline_class unset and raising Unknown model type: nextstep. The fix always sets pipeline_class for nextstep; this test fails on old code and passes after the fix. XPU: tensor_parallel_size=2 only to avoid OOM.

[fhfuih]

Duplicate #1947 and already (indirectly) fixed by #1908

For the test, all tests of the example code snippets in the documentation will be introduced in #1910 , which include the NextStep use example in Offline Text to Image page. Because according to my own test, seems this is related to that specific text-to-image.py script. If we create Omni ourselves without setting model_class_name, it runs fine.

So overall this PR is duplicate. Thanks for the contribution though. If you are interested in, we are also calling for model tests (similar to the scripts your have written here, but your version does not follow the official guideline) you can drop your name at #1832 and offer relevant test to NextStep in another PR