[Bugfix] Preserve default diffusion sampling params in default stage#2780
Conversation
Signed-off-by: David Chen <530634352@qq.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
| default_sampling_params = None | ||
| if not isinstance(default_sampling_params, dict): | ||
| default_sampling_params = None | ||
| stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} |
There was a problem hiding this comment.
So here the logic is, for multiple stages model, we take stage-0 as defaulte params for all stages. Right?
There was a problem hiding this comment.
Not exactly.
This code path is only used when we synthesize the fallback single-stage diffusion config in _create_default_diffusion_stage_cfg(). In that case, there is only one generated stage, and its stage_id is always 0, so reading default_sampling_params["0"] is intentional.
For multi-stage models, we do not go through this helper. We load the resolved stage configs directly, and each stage keeps its own default_sampling_params, which are later read from that stage’s config during metadata extraction. So this change does not apply stage-0 defaults to all stages; it only preserves stage-0 defaults for the synthetic single-stage diffusion path.
lishunyang12
left a comment
lishunyang12
left a comment
There was a problem hiding this comment.
LGTM. The change is correct and well-scoped.
What it does: When _create_default_diffusion_stage_cfg builds a single-stage diffusion config from CLI kwargs, it now parses default_sampling_params (JSON string keyed by stage index), extracts the stage-0 entry, and includes it in the generated stage config dict. This aligns with how stage_config.py / parse_stage_deploy_config reads default_sampling_params from stage data.
Correctness notes:
- The hardcoded key
"0"is correct here since this method always creates a single stage withstage_id: 0. - JSON parse error handling is reasonable (logs a warning, falls back to empty dict).
- The guard
if not isinstance(default_sampling_params, dict)handles the case where the parsed JSON is a non-dict type (e.g., a list or number). - The
"default_sampling_params"key is placed at the correct level in the stage config dict (sibling ofengine_args,runtime, etc.), matching whatparse_stage_deploy_configexpects viastage_data.get("default_sampling_params").
Tests: Unit test for the builder method and extended integration assertions for generator_device propagation both look good.
Minor optional suggestion (non-blocking): the variable name stage_default_sampling_params could be just stage_sampling_defaults for brevity, but this is purely stylistic.
…llm-project#2780) Signed-off-by: David Chen <530634352@qq.com>
…llm-project#2780) Signed-off-by: David Chen <530634352@qq.com>
…llm-project#2780) Signed-off-by: David Chen <530634352@qq.com>
3 participants
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
Purpose
Fix the default diffusion-stage bootstrap path so
generator_deviceand other diffusion sampling fields can be provided through--default-sampling-params.Before this change,
generator_devicewas already supported inOmniDiffusionSamplingParamsand honored by the diffusion runner, but the auto-generated single-stage diffusion config did not preservedefault_sampling_paramsfrom the CLI. As a result, defaults passed through--default-sampling-paramscould be dropped in the default stage-config path.This change keeps the existing CLI surface unchanged and aligns startup defaults with the per-request API behavior, instead of introducing a separate dedicated CLI flag.
Test Plan
AsyncOmniEngine._create_default_diffusion_stage_cfg()preserves parsed stage-0default_sampling_params.generator_devicecoming from default sampling params for both single-stage and AsyncOmni paths.python -m py_compile vllm_omni/engine/async_omni_engine.py tests/entrypoints/test_async_omni_diffusion_config.py tests/entrypoints/openai_api/test_image_server.pyTest Result
py_compilepassed for all modified files.pytestexecution was not completed in the local environment because test startup was blocked by missing dependencies (cv2,aenum).Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model. Please runmkdocs serveto sync the documentation editions to./docs.BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)