[Diffusion] Add Boogu Edit-Turbo support docs and tests (#6701)

Signed-off-by: xRay2016 <1150722393@qq.com>

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """L1 unit tests for the native Boogu-Image pipeline.
 
@@ -133,6 +133,51 @@ def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
     assert boogu_pipeline.mllm is mock_dependencies["inner_encoder"]
 
 
+def test_constructor_forwards_revision_to_all_component_loaders(mock_dependencies, mocker):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    revision = "hotfix-1k-20260708"
+    prefetch = mocker.patch(f"{_MODULE}.prefetch_subfolders")
+    scheduler_loader = mocker.patch(
+        f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
+        return_value=mock_dependencies["scheduler"],
+    )
+    mllm_loader = mocker.patch(
+        f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
+        return_value=mock_dependencies["mllm_wrapper"],
+    )
+    processor_loader = mocker.patch(
+        f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
+        return_value=mock_dependencies["processor"],
+    )
+    vae_loader = mocker.patch(
+        f"{_MODULE}.AutoencoderKL.from_pretrained",
+        return_value=mock_dependencies["vae"],
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        revision=revision,
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        num_gpus=1,
+    )
+
+    pipeline = BooguImagePipeline(od_config=od_config)
+
+    (source,) = pipeline.weights_sources
+    assert source.revision == revision
+    prefetch.assert_called_once_with(
+        "dummy-boogu",
+        ["scheduler", "vae", "mllm", "processor"],
+        local_files_only=True,
+        revision=revision,
+    )
+    for loader in (scheduler_loader, mllm_loader, processor_loader, vae_loader):
+        assert loader.call_args.kwargs["revision"] == revision
+
+
 def test_constructor_weights_sources(boogu_pipeline):
     (source,) = boogu_pipeline.weights_sources
     assert source.model_or_path == "dummy-boogu"

tests/e2e/online_serving/test_boogu_image_edit.py

@@ -1,15 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
-Online serving tests for Boogu-Image-0.1-Edit (image-to-image via chat completions).
+Online serving tests for Boogu-Image-0.1-Edit and Edit-Turbo.
 
-The Edit checkpoint shares the ``BooguImagePipeline`` class with the Base
+Both Edit checkpoints share the ``BooguImagePipeline`` class with the Base
 text-to-image checkpoint; the editing (TI2I) path activates when the request
-carries a reference image. Text guidance rides on ``guidance_scale`` (upstream
-default 4.0), and image guidance rides on ``guidance_scale_2`` (upstream default
-1.0 = off; a value > 1 enables the double-guidance path). Only a single
-reference image is supported.
+carries a reference image. The regular Edit checkpoint uses text guidance via
+``guidance_scale`` and optional image guidance via ``guidance_scale_2``.
+Edit-Turbo is guidance-distilled and uses four steps with guidance scale 1.0.
+Only a single reference image is supported.
 
 - ``test_single_image_to_image_001``: one reference image, text guidance only.
 - ``test_double_guidance_001``: one reference image, text + image guidance.
@@ -40,6 +40,10 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 EDIT_PROMPT = "Change the style to a colored pencil drawing."
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
 
+EDIT_MODEL = "Boogu/Boogu-Image-0.1-Edit"
+EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo"
+EDIT_TURBO_REVISION = "hotfix-1k-20260708"
+
 
 def _get_diffusion_feature_cases(model: str):
     """Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``)."""
@@ -54,13 +58,44 @@ def _get_diffusion_feature_cases(model: str):
     ]
 
 
+def _get_single_edit_cases():
+    """Return the regular and distilled Edit server/request configurations."""
+    return [
+        pytest.param(
+            OmniServerParams(model=EDIT_MODEL),
+            {
+                "num_inference_steps": 2,
+                "guidance_scale": 5.0,
+            },
+            id="edit",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=EDIT_TURBO_MODEL,
+                server_args=["--revision", EDIT_TURBO_REVISION],
+            ),
+            {
+                "num_inference_steps": 4,
+                "guidance_scale": 1.0,
+            },
+            id="edit-turbo-hotfix-1k",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
 @pytest.mark.parametrize(
-    "omni_server",
-    _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
-    indirect=True,
+    ("omni_server", "sampling_params"),
+    _get_single_edit_cases(),
+    indirect=["omni_server"],
 )
-def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Single-reference text-guided edit smoke for ``Boogu/Boogu-Image-0.1-Edit``."""
+def test_single_image_to_image_001(
+    omni_server: OmniServer,
+    sampling_params: dict,
+    openai_client: OpenAIClientHandler,
+):
+    """Single-reference text-guided edit smoke for regular and Turbo Edit."""
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
     messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
@@ -69,9 +104,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
         "model": omni_server.model,
         "messages": messages,
         "extra_body": {
-            "num_inference_steps": 2,
-            # Boogu text guidance (upstream default 4.0); > 1 enables CFG.
-            "guidance_scale": 5.0,
+            **sampling_params,
             "seed": 42,
         },
     }
@@ -81,7 +114,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
+    _get_diffusion_feature_cases(EDIT_MODEL),
     indirect=True,
 )
 def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
@@ -105,15 +138,20 @@ def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClien
     openai_client.send_diffusion_request(request_config)
 
 
-def _edits_request_config(model: str, *, guidance_scale_2: float | None = None) -> dict:
+def _edits_request_config(
+    model: str,
+    *,
+    num_inference_steps: int = 2,
+    guidance_scale: float = 5.0,
+    guidance_scale_2: float | None = None,
+) -> dict:
     """Multipart ``/v1/images/edits`` config: one synthetic reference + smoke-depth params."""
     image_bytes = base64.b64decode(generate_synthetic_image(512, 512)["base64"])
     data = {
         "model": model,
         "prompt": EDIT_PROMPT,
-        "num_inference_steps": 2,
-        # Boogu text guidance (upstream default 4.0); > 1 enables CFG.
-        "guidance_scale": 5.0,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
         "seed": 42,
     }
     if guidance_scale_2 is not None:
@@ -133,19 +171,25 @@ def _assert_edits_response_has_image(responses) -> None:
 
 
 @pytest.mark.parametrize(
-    "omni_server",
-    _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
-    indirect=True,
+    ("omni_server", "sampling_params"),
+    _get_single_edit_cases(),
+    indirect=["omni_server"],
 )
-def test_images_edits_endpoint_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Single-reference text-guided edit via the OpenAI-compatible ``/v1/images/edits`` endpoint."""
-    responses = openai_client.send_images_edits_http_request(_edits_request_config(omni_server.model))
+def test_images_edits_endpoint_001(
+    omni_server: OmniServer,
+    sampling_params: dict,
+    openai_client: OpenAIClientHandler,
+):
+    """Single-reference regular and Turbo Edit via ``/v1/images/edits``."""
+    responses = openai_client.send_images_edits_http_request(
+        _edits_request_config(omni_server.model, **sampling_params)
+    )
     _assert_edits_response_has_image(responses)
 
 
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
+    _get_diffusion_feature_cases(EDIT_MODEL),
     indirect=True,
 )
 def test_images_edits_endpoint_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):

tests/e2e/online_serving/test_boogu_image_edit_expansion.py

@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
-L4 expansion coverage for ``Boogu/Boogu-Image-0.1-Edit`` (image editing / TI2I).
+L4 expansion coverage for Boogu-Image Edit and Edit-Turbo (image editing / TI2I).
 
 The nightly smoke module (``test_boogu_image_edit.py``) already covers single-image
 text-guided editing and the text+image double-guidance path. This expansion
@@ -16,11 +16,12 @@ module adds the two edit paths smoke does not exercise:
   ``height`` / ``width``: asserts the output resolution follows the reference
   (upstream ``align_res``, on by default for a single reference).
 
-The Edit checkpoint shares ``BooguImagePipeline`` with the Base checkpoint; the
-TI2I path activates when the request carries a reference image. Only a single
-reference image is supported. Cases stay smoke-depth (``num_inference_steps=2``).
-CPU offload, Cache-DiT, and multi-GPU parallelism are not yet supported for this
-model and are intentionally omitted.
+Both Edit checkpoints share ``BooguImagePipeline`` with the Base checkpoint;
+the TI2I path activates when the request carries a reference image. Only a
+single reference image is supported. The regular Edit cases stay smoke-depth
+(``num_inference_steps=2``), while Edit-Turbo uses its tuned four-step
+configuration. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet
+supported for these models and are intentionally omitted.
 
 From ``tests/``::
 
@@ -41,6 +42,8 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 
 MODEL = "Boogu/Boogu-Image-0.1-Edit"
+EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo"
+EDIT_TURBO_REVISION = "hotfix-1k-20260708"
 EDIT_PROMPT = "Change the style to a colored pencil drawing."
 
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
@@ -57,6 +60,33 @@ def _get_diffusion_feature_cases(model: str):
     ]
 
 
+def _get_align_res_cases():
+    """Return regular and Turbo Edit configurations for align-res coverage."""
+    return [
+        pytest.param(
+            OmniServerParams(model=MODEL),
+            {
+                "num_inference_steps": 2,
+                "guidance_scale": 5.0,
+            },
+            id="edit",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=EDIT_TURBO_MODEL,
+                server_args=["--revision", EDIT_TURBO_REVISION],
+            ),
+            {
+                "num_inference_steps": 4,
+                "guidance_scale": 1.0,
+            },
+            id="edit-turbo-hotfix-1k",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
 @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
 def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
     """Image-only guidance: text CFG off (1.0), image CFG on (2.0) -> image-only CFG branch."""
@@ -76,8 +106,16 @@ def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClien
     openai_client.send_diffusion_request(request_config)
 
 
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+@pytest.mark.parametrize(
+    ("omni_server", "sampling_params"),
+    _get_align_res_cases(),
+    indirect=["omni_server"],
+)
+def test_align_res_output_size(
+    omni_server: OmniServer,
+    sampling_params: dict,
+    openai_client: OpenAIClientHandler,
+) -> None:
     """Output resolution follows the reference image (upstream ``align_res``).
 
     No ``height`` / ``width`` is requested. The reference is 512x384 — both
@@ -92,8 +130,7 @@ def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAICli
         "model": omni_server.model,
         "messages": messages,
         "extra_body": {
-            "num_inference_steps": 2,
-            "guidance_scale": 5.0,
+            **sampling_params,
             "seed": 42,
         },
     }

tests/entrypoints/test_utils.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Unit tests for vllm_omni.entrypoints.utils module."""
 
 import os
@@ -420,6 +423,7 @@ class TestLoadAndResolveStageConfigs:
         """Ensure that kwargs survive default stage creation."""
         engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
         kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
+        kwargs["revision"] = "pinned-revision"
         mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
         mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
 
@@ -433,6 +437,7 @@ class TestLoadAndResolveStageConfigs:
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
         assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
+        assert stage_configs[0]["engine_args"]["revision"] == "pinned-revision"
 
     def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
         deploy_path = tmp_path / "qwen3_multi.yaml"