Merge remote-tracking branch 'origin/main' into qwen-image-2.1

# Conflicts:
#	vllm_omni/diffusion/executor/multiproc_executor.py
#	vllm_omni/diffusion/offloader/sequential_backend.py

tests/diffusion/test_diffusion_config_propagation.py

@@ -11,12 +11,14 @@ import pytest
 import torch
 
 from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.omni_config import extract_diffusion_stage_config_kwargs
 from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.model_metadata import (
+    FLUX2_KLEIN_MAX_INPUT_IMAGES,
     HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
     QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
 )
@@ -24,6 +26,12 @@ from vllm_omni.diffusion.model_metadata import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def _local_model_paths(monkeypatch):
+    # These tests exercise config transport, not model repository resolution.
+    monkeypatch.setattr("vllm_omni.diffusion.data.get_model_path", lambda model, revision: model)
+
+
 def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
     """Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
 
@@ -33,7 +41,8 @@ def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
     """
     stages = StageConfigFactory.create_default_diffusion(kwargs)
     engine_args = dict(stages[0]["engine_args"])
-    return OmniDiffusionConfig.from_kwargs(**engine_args)
+    diffusion_kwargs = extract_diffusion_stage_config_kwargs(engine_args, stage_id=0)
+    return OmniDiffusionConfig(**{name: value for name, value in diffusion_kwargs.items() if value is not None})
 
 
 def test_direct_config_extras_are_independent_dicts():
@@ -65,9 +74,7 @@ class TestParallelConfigPropagation:
         stages = StageConfigFactory.create_default_diffusion({"parallel_config": pc, "model": "x"})
         assert stages[0]["runtime"]["devices"] == "0,1,2,3"
 
-        # Let __post_init__ reconstruct from dict (real code path)
-        ea = dict(stages[0]["engine_args"])
-        od = OmniDiffusionConfig.from_kwargs(**ea)
+        od = _roundtrip_diffusion_config(parallel_config=pc, model="x")
         assert od.parallel_config.tensor_parallel_size == 4
         assert od.parallel_config.world_size == 4
 
@@ -201,6 +208,17 @@ class TestCreateDefaultDiffusion:
         assert od.max_num_batched_tokens == 2048
         assert od.max_model_len == 4096
 
+    @pytest.mark.parametrize(
+        ("field_name", "value"),
+        [
+            ("enable_sleep_mod", None),
+            ("enable_lora", True),
+        ],
+    )
+    def test_unowned_raw_field_is_rejected(self, field_name, value):
+        with pytest.raises(ValueError, match=field_name):
+            StageConfigFactory.create_default_diffusion({"model": "x", field_name: value})
+
 
 def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
@@ -211,6 +229,18 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 
 
+def test_flux2_klein_sets_generic_multimodal_limit():
+    od_config = OmniDiffusionConfig(
+        model="black-forest-labs/FLUX.2-klein-9B",
+        model_class_name="Flux2KleinPipeline",
+    )
+
+    od_config.update_multimodal_support()
+
+    assert od_config.supports_multimodal_inputs is True
+    assert od_config.max_multimodal_image_inputs == FLUX2_KLEIN_MAX_INPUT_IMAGES
+
+
 def test_task_type_roundtrip():
     od = _roundtrip_diffusion_config(model="x", task_type="model-defined-task")
     assert od.task_type == "model-defined-task"