anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Bugfix] Fix DataType Handling in Default Diffusion Config (#2530)

Browse Source 
Signed-off-by: Alex Brooks <albrooks@redhat.com>
This commit is contained in:
Alex Brooks authored and GitHub committed 2026-04-07 11:24:19 +08:00
1 parent 7a72f34ce4
commit 08e2e1fc8e
2 files changed
+54 -31
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/entrypoints/test_utils.py
+18
View File
@@ -5,14 +5,17 @@ from collections import Counter
from dataclasses import dataclass
import pytest
import torch
from pytest_mock import MockerFixture
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.entrypoints.utils import (
_convert_dataclasses_to_dict,
_filter_dict_like_object,
filter_dataclass_kwargs,
load_and_resolve_stage_configs,
resolve_model_config_path,
)
@@ -304,3 +307,18 @@ class TestResolveModelConfigPath:
assert result is not None
assert "glm_image.yaml" in result
class TestLoadAndResolveStageConfigs:
def test_load_and_resolve_with_kwargs(self):
"""Ensure that dtype survives default stage creation."""
kwargs = {"dtype": torch.float32}
config_path, stage_configs = load_and_resolve_stage_configs(
model="black-forest-labs/FLUX.2-klein-4B",
stage_configs_path=None,
kwargs=kwargs,
default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
)
assert config_path is None
assert len(stage_configs) == 1
assert "dtype" in stage_configs[0]["engine_args"]
vllm_omni/engine/async_omni_engine.py
+36 -31
View File
@@ -920,6 +920,41 @@ class AsyncOmniEngine:
num_devices = max(1, int(parallel_config.world_size))
devices = ",".join(str(i) for i in range(num_devices))
stage_engine_args = {
"max_num_seqs": 1,
"parallel_config": parallel_config,
"model_class_name": kwargs.get("model_class_name", None),
"step_execution": kwargs.get("step_execution", False),
"vae_use_slicing": kwargs.get("vae_use_slicing", False),
"vae_use_tiling": kwargs.get("vae_use_tiling", False),
"cache_backend": cache_backend,
"cache_config": cache_config,
"enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
"enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
"enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
"enforce_eager": kwargs.get("enforce_eager", False),
"diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
"custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
"worker_extension_cls": kwargs.get("worker_extension_cls", None),
"enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
"enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
"num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
"quantization": kwargs.get("quantization", None),
"enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
**(
{
"profiler_config": asdict(kwargs["profiler_config"])
if hasattr(kwargs["profiler_config"], "__dataclass_fields__")
else kwargs["profiler_config"]
}
if kwargs.get("profiler_config") is not None
else {}
),
}
# Only set dtype if it was already explicitly passed and normalized
if "dtype" in normalized_kwargs:
stage_engine_args["dtype"] = normalized_kwargs["dtype"]
default_stage_cfg = [
{
"stage_id": 0,
@@ -928,37 +963,7 @@ class AsyncOmniEngine:
"process": True,
"devices": devices,
},
"engine_args": {
"max_num_seqs": 1,
"parallel_config": parallel_config,
"model_class_name": kwargs.get("model_class_name", None),
"step_execution": kwargs.get("step_execution", False),
"vae_use_slicing": kwargs.get("vae_use_slicing", False),
"vae_use_tiling": kwargs.get("vae_use_tiling", False),
"cache_backend": cache_backend,
"cache_config": cache_config,
"enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
"enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
"enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
"enforce_eager": kwargs.get("enforce_eager", False),
"diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
"custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
"worker_extension_cls": kwargs.get("worker_extension_cls", None),
"enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
"enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
"num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
"quantization": kwargs.get("quantization", None),
"enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
**(
{
"profiler_config": asdict(kwargs["profiler_config"])
if hasattr(kwargs["profiler_config"], "__dataclass_fields__")
else kwargs["profiler_config"]
}
if kwargs.get("profiler_config") is not None
else {}
),
},
"engine_args": stage_engine_args,
"final_output": True,
"final_output_type": "image",
}