[Model] Support Boogu/Boogu-Image-0.1-Base and Boogu/Boogu-Image-0.1-Edit (#4995)

Signed-off-by: zzehli <jaeli_ottawa@outlook.com>
Signed-off-by: jl <jaeli_ottawa@outlook.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/diffusion/models/boogu_image/test_boogu_image_scheduler.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+"""L1 unit tests for the Boogu time-shifting flow-match Euler scheduler.
+
+Golden values were generated from the upstream boogu scheduler
+(scripts/test_ported_scheduler_parity.py) with the released checkpoint's
+scheduler_config.json, so CI does not need the `boogu` package installed.
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.boogu_image import FlowMatchEulerDiscreteScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Ground truth from the released checkpoint's scheduler/scheduler_config.json.
+CHECKPOINT_CONFIG = dict(
+    num_train_timesteps=1000,
+    do_shift=True,
+    dynamic_time_shift=False,
+    time_shift_version="v1",
+    seq_len=4096,
+    base_shift=0.5,
+    max_shift=1.15,
+)
+
+NUM_TOKENS = 4096 * 4  # 1024x1024 resolution
+
+# Generated from the upstream boogu scheduler with CHECKPOINT_CONFIG.
+GOLDEN_TIMESTEPS = {
+    4: [0.00000000, 0.09546924, 0.24048907, 0.48715585],
+    10: [
+        0.00000000,
+        0.03398621,
+        0.07335263,
+        0.11948693,
+        0.17429835,
+        0.24048907,
+        0.32201332,
+        0.42489707,
+        0.55880034,
+        0.74024153,
+    ],
+}
+
+
+def make_scheduler(**overrides) -> FlowMatchEulerDiscreteScheduler:
+    return FlowMatchEulerDiscreteScheduler(**{**CHECKPOINT_CONFIG, **overrides})
+
+
+@pytest.mark.parametrize("n_steps", sorted(GOLDEN_TIMESTEPS))
+def test_golden_timesteps_checkpoint_config(n_steps):
+    scheduler = make_scheduler()
+    scheduler.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
+    np.testing.assert_allclose(
+        scheduler.timesteps.numpy(),
+        np.asarray(GOLDEN_TIMESTEPS[n_steps], dtype=np.float32),
+        atol=1e-7,
+    )
+
+
+@pytest.mark.parametrize("n_steps", [4, 10, 30])
+def test_timestep_grid_structure(n_steps):
+    scheduler = make_scheduler()
+    scheduler.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
+
+    assert scheduler.timesteps.shape == (n_steps,)
+    assert scheduler._timesteps.shape == (n_steps + 1,)
+    # Ascending 0 -> 1 convention with the terminal 1.0 appended.
+    assert scheduler.timesteps[0].item() == 0.0
+    assert torch.all(scheduler._timesteps[1:] > scheduler._timesteps[:-1])
+    assert scheduler._timesteps[-1].item() == 1.0
+    assert torch.equal(scheduler._timesteps[:-1], scheduler.timesteps)
+
+
+def test_no_shift_is_uniform_grid():
+    n_steps = 8
+    scheduler = make_scheduler(do_shift=False)
+    scheduler.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
+    expected = np.linspace(0, 1, n_steps + 1, dtype=np.float32)[:-1]
+    np.testing.assert_allclose(scheduler.timesteps.numpy(), expected, atol=1e-7)
+
+
+def test_static_shift_ignores_num_tokens():
+    # dynamic_time_shift=False: the shift comes from configured seq_len, so
+    # per-call num_tokens must not change the schedule.
+    scheduler_a = make_scheduler()
+    scheduler_a.set_timesteps(num_inference_steps=10, num_tokens=NUM_TOKENS)
+    scheduler_b = make_scheduler()
+    scheduler_b.set_timesteps(num_inference_steps=10, num_tokens=320 * 320)
+    assert torch.equal(scheduler_a.timesteps, scheduler_b.timesteps)
+
+
+def test_step_arithmetic_and_index_advancement():
+    n_steps = 4
+    scheduler = make_scheduler()
+    scheduler.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
+
+    sample = torch.zeros(1, 4, 8, 8)
+    model_output = torch.ones(1, 4, 8, 8)
+
+    assert scheduler.step_index is None
+    out = scheduler.step(model_output, scheduler.timesteps[0], sample)
+    assert scheduler.step_index == 1
+
+    # prev_sample = sample + (t_next - t) * model_output
+    expected_dt = scheduler._timesteps[1] - scheduler._timesteps[0]
+    assert torch.allclose(out.prev_sample, expected_dt * model_output)
+
+    # Full loop reaches the end of the grid.
+    sample = out.prev_sample
+    for t in scheduler.timesteps[1:]:
+        sample = scheduler.step(model_output, t, sample, return_dict=False)[0]
+    assert scheduler.step_index == n_steps
+    # Integrating a constant velocity of 1 over t in [0, 1] gives exactly 1.
+    assert torch.allclose(sample, torch.ones_like(sample), atol=1e-6)
+
+
+def test_step_rejects_integer_timestep():
+    scheduler = make_scheduler()
+    scheduler.set_timesteps(num_inference_steps=4, num_tokens=NUM_TOKENS)
+    with pytest.raises(ValueError, match="integer indices"):
+        scheduler.step(torch.zeros(1), 0, torch.zeros(1))
+
+
+def test_set_begin_index():
+    scheduler = make_scheduler()
+    scheduler.set_timesteps(num_inference_steps=4, num_tokens=NUM_TOKENS)
+    scheduler.set_begin_index(2)
+    assert scheduler.begin_index == 2
+
+    scheduler.step(torch.zeros(1), scheduler.timesteps[2], torch.zeros(1))
+    assert scheduler.step_index == 3
+
+
+def test_len_and_config_roundtrip():
+    scheduler = make_scheduler()
+    assert len(scheduler) == CHECKPOINT_CONFIG["num_train_timesteps"]
+    # ConfigMixin must retain the checkpoint config (needed for from_pretrained).
+    assert scheduler.config.do_shift is True
+    assert scheduler.config.dynamic_time_shift is False
+    assert scheduler.config.time_shift_version == "v1"
+    assert scheduler.config.seq_len == 4096
+
+
+def test_matches_upstream_boogu_scheduler():
+    """Differential check against the upstream boogu package, when installed."""
+    boogu_schedulers = pytest.importorskip("boogu.schedulers.scheduling_flow_match_euler_discrete_time_shifting")
+    UpstreamScheduler = boogu_schedulers.FlowMatchEulerDiscreteScheduler
+
+    for dynamic in (False, True):
+        for version in ("v1", "v2"):
+            for n_steps in (4, 10, 30):
+                config = {**CHECKPOINT_CONFIG, "dynamic_time_shift": dynamic, "time_shift_version": version}
+                upstream = UpstreamScheduler(**config)
+                ported = make_scheduler(**config)
+                upstream.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
+                ported.set_timesteps(num_inference_steps=n_steps, num_tokens=NUM_TOKENS)
+                assert torch.equal(upstream.timesteps, ported.timesteps), (dynamic, version, n_steps)
+                assert torch.equal(upstream._timesteps, ported._timesteps), (dynamic, version, n_steps)

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -0,0 +1,433 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+# Small mock dimensions (head_dim = HIDDEN_SIZE / NUM_HEADS = 16 must equal
+# sum(AXES_DIM_ROPE)).
+HIDDEN_SIZE = 64
+NUM_HEADS = 4
+NUM_KV_HEADS = 2
+HEAD_DIM = HIDDEN_SIZE // NUM_HEADS
+AXES_DIM_ROPE = (8, 4, 4)
+AXES_LENS = (32, 16, 16)
+MULTIPLE_OF = 32
+NORM_EPS = 1e-5
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Initialize the minimal single-rank distributed environment required by
+    the vLLM parallel linear layers (tensor-parallel group must exist)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29512")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for tests using CPU tensors.
+
+    vLLM's dispatch_unquantized_gemm() selects the backend by platform, not by
+    tensor device; CPU test tensors can crash on non-default backends."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+def _randomize_parameters(module: torch.nn.Module) -> None:
+    """Fill parameters with small random values.
+
+    vLLM parallel linears allocate weights with `torch.empty` (real weights
+    arrive via `load_weights`), so uninitialized memory must be overwritten
+    before a forward pass."""
+    with torch.no_grad():
+        for param in module.parameters():
+            param.uniform_(-0.02, 0.02)
+
+
+def _identity_rotary_emb(batch_size: int, seq_len: int) -> torch.Tensor:
+    """Complex rotary frequencies encoding a zero rotation."""
+    return torch.polar(
+        torch.ones(batch_size, seq_len, HEAD_DIM // 2),
+        torch.zeros(batch_size, seq_len, HEAD_DIM // 2),
+    )
+
+
+def _tiny_tf_model_config(**overrides):
+    config = {
+        "patch_size": 2,
+        "in_channels": 4,
+        "hidden_size": HIDDEN_SIZE,
+        "num_layers": 4,
+        "num_double_stream_layers": 2,
+        "num_refiner_layers": 2,
+        "num_attention_heads": NUM_HEADS,
+        "num_kv_heads": NUM_KV_HEADS,
+        "multiple_of": MULTIPLE_OF,
+        "norm_eps": NORM_EPS,
+        "axes_dim_rope": list(AXES_DIM_ROPE),
+        "axes_lens": list(AXES_LENS),
+        "instruction_feature_configs": {
+            "instruction_feat_dim": 32,
+            "reduce_type": "mean",
+            "num_instruction_feature_layers": 1,
+        },
+        "prompt_tuning_configs": {"use_prompt_tuning": False},
+        "timestep_scale": 1.0,
+    }
+    config.update(overrides)
+    return config
+
+
+def _tiny_od_config(**overrides):
+    from vllm_omni.diffusion.data import TransformerConfig
+
+    return SimpleNamespace(
+        tf_model_config=TransformerConfig.from_dict(_tiny_tf_model_config(**overrides)),
+        dtype=torch.float32,
+    )
+
+
+def test_boogu_image_transformer_import():
+    from vllm_omni.diffusion.models.boogu_image import BooguImageTransformer2DModel
+
+    assert BooguImageTransformer2DModel is not None
+
+
+def test_single_stream_block_shape():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformerBlock,
+    )
+
+    block = BooguImageTransformerBlock(
+        dim=HIDDEN_SIZE,
+        num_attention_heads=NUM_HEADS,
+        num_kv_heads=NUM_KV_HEADS,
+        multiple_of=MULTIPLE_OF,
+        ffn_dim_multiplier=None,
+        norm_eps=NORM_EPS,
+        modulation=True,
+    )
+    _randomize_parameters(block)
+
+    batch_size, seq_len = 1, 16
+    hidden_states = torch.randn(batch_size, seq_len, HIDDEN_SIZE)
+    attention_mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
+    rotary_emb = _identity_rotary_emb(batch_size, seq_len)
+    temb = torch.randn(batch_size, min(HIDDEN_SIZE, 1024))
+
+    out = block(hidden_states, attention_mask, rotary_emb, temb)
+    assert out.shape == hidden_states.shape
+    assert torch.isfinite(out).all()
+
+    with pytest.raises(ValueError, match="temb"):
+        block(hidden_states, attention_mask, rotary_emb, None)
+
+
+def test_single_stream_block_no_modulation():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformerBlock,
+    )
+
+    block = BooguImageTransformerBlock(
+        dim=HIDDEN_SIZE,
+        num_attention_heads=NUM_HEADS,
+        num_kv_heads=NUM_KV_HEADS,
+        multiple_of=MULTIPLE_OF,
+        ffn_dim_multiplier=None,
+        norm_eps=NORM_EPS,
+        modulation=False,
+    )
+    _randomize_parameters(block)
+
+    batch_size, seq_len = 2, 8
+    hidden_states = torch.randn(batch_size, seq_len, HIDDEN_SIZE)
+    out = block(hidden_states, None, _identity_rotary_emb(batch_size, seq_len))
+    assert out.shape == hidden_states.shape
+    assert torch.isfinite(out).all()
+
+
+def test_double_stream_block_shape():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageDoubleStreamTransformerBlock,
+    )
+
+    block = BooguImageDoubleStreamTransformerBlock(
+        dim=HIDDEN_SIZE,
+        num_attention_heads=NUM_HEADS,
+        num_kv_heads=NUM_KV_HEADS,
+        multiple_of=MULTIPLE_OF,
+        ffn_dim_multiplier=None,
+        norm_eps=NORM_EPS,
+        modulation=True,
+    )
+    _randomize_parameters(block)
+
+    batch_size = 1
+    img_len, instruct_len = 16, 8
+    total_len = img_len + instruct_len
+
+    img_hidden_states = torch.randn(batch_size, img_len, HIDDEN_SIZE)
+    instruct_hidden_states = torch.randn(batch_size, instruct_len, HIDDEN_SIZE)
+    img_attention_mask = torch.ones(batch_size, img_len, dtype=torch.bool)
+    joint_attention_mask = torch.ones(batch_size, total_len, dtype=torch.bool)
+    image_rotary_emb = _identity_rotary_emb(batch_size, img_len)
+    rotary_emb = _identity_rotary_emb(batch_size, total_len)
+    temb = torch.randn(batch_size, min(HIDDEN_SIZE, 1024))
+
+    img_out, instruct_out = block(
+        img_hidden_states,
+        instruct_hidden_states,
+        img_attention_mask,
+        joint_attention_mask,
+        image_rotary_emb,
+        rotary_emb,
+        temb=temb,
+        encoder_seq_lengths=[instruct_len] * batch_size,
+        seq_lengths=[total_len] * batch_size,
+    )
+
+    assert img_out.shape == img_hidden_states.shape
+    assert instruct_out.shape == instruct_hidden_states.shape
+    assert torch.isfinite(img_out).all()
+    assert torch.isfinite(instruct_out).all()
+
+
+def test_transformer_instantiates():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+
+    assert BooguImageTransformer2DModel._repeated_blocks == [
+        "BooguImageTransformerBlock",
+        "BooguImageContextRefinerTransformerBlock",
+        "BooguImageSingleStreamTransformerBlock",
+    ]
+    assert BooguImageTransformer2DModel._layerwise_offload_blocks_attrs == [
+        "single_stream_layers",
+        "double_stream_layers",
+    ]
+
+    assert len(model.noise_refiner) == 2
+    assert len(model.ref_image_refiner) == 2
+    assert len(model.context_refiner) == 2
+    assert len(model.double_stream_layers) == 2
+    assert len(model.single_stream_layers) == 2  # num_layers - num_double_stream_layers
+    assert model.image_index_embedding.shape == (5, HIDDEN_SIZE)
+    # patch_size^2 * in_channels -> hidden_size
+    assert model.x_embedder.in_features == 2 * 2 * 4
+    assert model.x_embedder.out_features == HIDDEN_SIZE
+
+
+def test_transformer_preprocesses_multiple_instruction_feature_layers():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    instruction_feature_configs = {
+        "instruction_feat_dim": 32,
+        "reduce_type": "concat",
+        "num_instruction_feature_layers": 2,
+    }
+    model = BooguImageTransformer2DModel(
+        od_config=_tiny_od_config(instruction_feature_configs=instruction_feature_configs)
+    )
+    hidden_states = [torch.randn(1, 8, 32), torch.randn(1, 8, 32)]
+
+    processed = model.preprocess_instruction_hidden_states(hidden_states)
+
+    assert model.preprocessed_instruction_feat_dim == 64
+    assert processed.shape == (1, 8, 64)
+    assert torch.equal(processed, torch.cat(hidden_states, dim=-1))
+
+
+def test_transformer_validates_rope_dims():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    with pytest.raises(ValueError, match="axes_dim_rope"):
+        BooguImageTransformer2DModel(od_config=_tiny_od_config(axes_dim_rope=[8, 8, 8]))
+
+
+def test_transformer_rejects_prompt_tuning():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    with pytest.raises(NotImplementedError, match="[Pp]rompt tuning"):
+        BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True}))
+
+
+def _native_to_checkpoint_name(name: str) -> str:
+    """Inverse of ``load_weights`` remapping: native param name -> diffusers name.
+
+    - ``.to_out.<suffix>`` -> ``.to_out.0.<suffix>`` (diffusers ModuleList wrap).
+    - promoted joint-attention projections move back under ``.processor.``.
+    """
+    if ".to_out." in name:
+        name = name.replace(".to_out.", ".to_out.0.")
+    for proj in (
+        "img_to_q",
+        "img_to_k",
+        "img_to_v",
+        "instruct_to_q",
+        "instruct_to_k",
+        "instruct_to_v",
+        "instruct_out",
+        "img_out",
+    ):
+        token = f".img_instruct_attn.{proj}."
+        if token in name:
+            name = name.replace(token, f".img_instruct_attn.processor.{proj}.")
+            break
+    return name
+
+
+def test_transformer_load_weights_round_trip():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+    native_params = dict(model.named_parameters())
+
+    # Build synthetic diffusers-named weights (one per native parameter).
+    checkpoint_weights = {}
+    for native_name, param in native_params.items():
+        checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param)
+
+    # The remapping must be a bijection over the parameter set.
+    assert len(checkpoint_weights) == len(native_params)
+
+    loaded = model.load_weights(list(checkpoint_weights.items()))
+
+    # No missing / unexpected parameters.
+    assert loaded == set(native_params.keys())
+
+    # Values landed on the right parameters (TP=1: weight_loader copies verbatim).
+    reloaded = dict(model.named_parameters())
+    for native_name in native_params:
+        expected = checkpoint_weights[_native_to_checkpoint_name(native_name)]
+        assert torch.allclose(reloaded[native_name], expected)
+
+
+def test_transformer_load_weights_warns_for_unexpected_and_unloaded(caplog):
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+    native_params = dict(model.named_parameters())
+    loaded_name = next(iter(native_params))
+    checkpoint_name = _native_to_checkpoint_name(loaded_name)
+
+    loaded = model.load_weights(
+        [
+            (checkpoint_name, torch.randn_like(native_params[loaded_name])),
+            ("unexpected.weight", torch.ones(1)),
+        ]
+    )
+
+    assert loaded == {loaded_name}
+    assert "Skipping unexpected checkpoint weight unexpected.weight" in caplog.text
+    assert "Model parameters not loaded from checkpoint" in caplog.text
+    assert next(name for name in native_params if name != loaded_name) in caplog.text
+
+
+def test_transformer_forward_t2i_shape():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageDoubleStreamRotaryPosEmbed,
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+    _randomize_parameters(model)
+    model.eval()
+
+    batch_size = 1
+    in_channels = 4
+    latent_h = latent_w = 8  # multiples of patch_size (2)
+    instruct_len = 8
+    instruction_feat_dim = 32  # matches _tiny_tf_model_config
+
+    latents = torch.randn(batch_size, in_channels, latent_h, latent_w)
+    timestep = torch.full((batch_size,), 0.5)
+    instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
+    instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
+    freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
+
+    with torch.no_grad():
+        out = model(latents, timestep, instruction_hidden_states, freqs_cis, instruction_attention_mask)
+
+    assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
+    assert torch.isfinite(out).all()
+
+
+def test_transformer_forward_ti2i_shape():
+    """Editing path: a non-empty ``ref_image_hidden_states`` exercises the
+    reference-image patch embedder + refiner and must not change the output
+    shape (the output tracks the noise-latent dimensions)."""
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageDoubleStreamRotaryPosEmbed,
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+    _randomize_parameters(model)
+    model.eval()
+
+    batch_size = 1
+    in_channels = 4
+    latent_h = latent_w = 8
+    ref_h, ref_w = 6, 10  # a differently-sized reference latent
+    instruct_len = 8
+    instruction_feat_dim = 32
+
+    latents = torch.randn(batch_size, in_channels, latent_h, latent_w)
+    timestep = torch.full((batch_size,), 0.5)
+    instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
+    instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
+    freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
+
+    # One sample, one reference image (Boogu editing supports a single ref).
+    ref_image_hidden_states = [[torch.randn(in_channels, ref_h, ref_w)]]
+
+    with torch.no_grad():
+        out = model(
+            latents,
+            timestep,
+            instruction_hidden_states,
+            freqs_cis,
+            instruction_attention_mask,
+            ref_image_hidden_states=ref_image_hidden_states,
+        )
+
+    assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
+    assert torch.isfinite(out).all()

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -0,0 +1,900 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""L1 unit tests for the native Boogu-Image pipeline.
+
+Two groups:
+
+1. Constructor tests: all ``from_pretrained`` calls are mocked (Ovis pattern)
+   so the pipeline ``__init__`` wiring is exercised without downloading
+   weights or building the real transformer.
+2. Prompt-encoding tests: a pipeline shell (``object.__new__``) is wired with
+   small deterministic fakes so the ported chat templating, processor kwargs,
+   CFG handling, and reshape logic can be verified numerically on CPU.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig, TransformerConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_MODULE = "vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image"
+
+_EMBED_DIM = 8
+_SEQ_LEN = 6
+
+
+# ---------------------------------------------------------------------------
+# Constructor tests (mocked from_pretrained)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def mock_dependencies(mocker, monkeypatch):
+    """Mock external components so ``__init__`` runs without real weights."""
+    # The full VLM wrapper has an ``lm_head``; the pipeline must strip it and
+    # keep the inner ``.model`` as the encoder.
+    inner_encoder = mocker.MagicMock(name="inner_qwen3vl_model")
+    inner_encoder.dtype = torch.float32
+    mllm_wrapper = mocker.MagicMock(name="qwen3vl_wrapper")
+    mllm_wrapper.model.to.return_value = inner_encoder
+
+    mock_processor = mocker.MagicMock(name="processor")
+
+    mock_vae = mocker.MagicMock(name="vae")
+    mock_vae.config.block_out_channels = [128, 256, 512, 512]  # scale factor 8
+    mock_vae.to.return_value = mock_vae
+
+    mock_scheduler = mocker.MagicMock(name="scheduler")
+
+    monkeypatch.setattr(
+        f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
+        lambda *a, **k: mock_scheduler,
+    )
+    monkeypatch.setattr(
+        f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
+        lambda *a, **k: mllm_wrapper,
+    )
+    monkeypatch.setattr(
+        f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
+        lambda *a, **k: mock_processor,
+    )
+    monkeypatch.setattr(
+        f"{_MODULE}.AutoencoderKL.from_pretrained",
+        lambda *a, **k: mock_vae,
+    )
+
+    mock_transformer_cls = mocker.MagicMock(name="transformer_cls")
+    mock_transformer_instance = mocker.MagicMock(name="transformer")
+    mock_transformer_cls.return_value = mock_transformer_instance
+    monkeypatch.setattr(f"{_MODULE}.BooguImageTransformer2DModel", mock_transformer_cls)
+
+    # Treat the dummy model id as a local path: skips hub prefetch.
+    mocker.patch("os.path.exists", return_value=True)
+
+    return {
+        "inner_encoder": inner_encoder,
+        "mllm_wrapper": mllm_wrapper,
+        "processor": mock_processor,
+        "vae": mock_vae,
+        "scheduler": mock_scheduler,
+        "transformer": mock_transformer_instance,
+    }
+
+
+@pytest.fixture
+def boogu_pipeline(mock_dependencies):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        num_gpus=1,
+    )
+    return BooguImagePipeline(od_config=od_config)
+
+
+def test_boogu_image_pipeline_import():
+    from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline
+
+    assert BooguImagePipeline is not None
+
+
+def test_component_discovery_declarations():
+    from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline
+
+    # CPU offload / HSDP discovery must find ``mllm`` (there is no
+    # ``text_encoder`` attribute on this pipeline).
+    assert BooguImagePipeline._dit_modules == ["transformer"]
+    assert BooguImagePipeline._encoder_modules == ["mllm"]
+    assert BooguImagePipeline._vae_modules == ["vae"]
+
+
+def test_constructor_wires_components(boogu_pipeline, mock_dependencies):
+    assert boogu_pipeline.scheduler is mock_dependencies["scheduler"]
+    assert boogu_pipeline.processor is mock_dependencies["processor"]
+    assert boogu_pipeline.vae is mock_dependencies["vae"]
+    assert boogu_pipeline.transformer is mock_dependencies["transformer"]
+    assert boogu_pipeline.vae_scale_factor == 8
+    assert boogu_pipeline.default_sample_size == 128
+    assert hasattr(boogu_pipeline, "load_weights")
+
+
+def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
+    # Upstream encodes with the inner Qwen3VLModel, not the generation wrapper.
+    assert boogu_pipeline.mllm is mock_dependencies["inner_encoder"]
+
+
+def test_constructor_weights_sources(boogu_pipeline):
+    (source,) = boogu_pipeline.weights_sources
+    assert source.model_or_path == "dummy-boogu"
+    assert source.subfolder == "transformer"
+    assert source.prefix == "transformer."
+    assert source.fall_back_to_pt is True
+
+
+@pytest.mark.parametrize(
+    ("parallel_config", "cache_backend", "message"),
+    [
+        (DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"),
+        (DiffusionParallelConfig(ulysses_degree=2), "none", "Sequence parallelism"),
+        (DiffusionParallelConfig(ring_degree=2), "none", "Sequence parallelism"),
+        (DiffusionParallelConfig(cfg_parallel_size=2), "none", "CFG parallelism"),
+        (
+            DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=2),
+            "none",
+            "HSDP",
+        ),
+        (DiffusionParallelConfig(), "cache_dit", "Cache backend 'cache_dit'"),
+        (DiffusionParallelConfig(), "tea_cache", "Cache backend 'tea_cache'"),
+    ],
+)
+def test_constructor_rejects_unsupported_execution_modes(
+    mock_dependencies,
+    parallel_config,
+    cache_backend,
+    message,
+):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        parallel_config=parallel_config,
+        cache_backend=cache_backend,
+    )
+
+    with pytest.raises(NotImplementedError, match=message):
+        BooguImagePipeline(od_config=od_config)
+
+    # Validation happens before any checkpoint component is constructed.
+    mock_dependencies["mllm_wrapper"].model.to.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# Prompt-encoding tests (deterministic fakes, no constructor)
+# ---------------------------------------------------------------------------
+
+
+class _RecordingProcessor:
+    """Fake Qwen3VLProcessor: deterministic token ids derived from the text."""
+
+    def __init__(self):
+        self.calls = []
+
+    def apply_chat_template(self, prompts, **kwargs):
+        self.calls.append({"prompts": prompts, "kwargs": kwargs})
+        batch = len(prompts)
+        input_ids = torch.zeros(batch, _SEQ_LEN, dtype=torch.long)
+        for i, messages in enumerate(prompts):
+            system_text = messages[0]["content"][0]["text"]
+            user_text = messages[1]["content"][0]["text"]
+            input_ids[i, 0] = len(system_text) % 997
+            input_ids[i, 1] = len(user_text) % 997
+            input_ids[i, 2:] = torch.arange(2, _SEQ_LEN) + i
+        attention_mask = torch.ones(batch, _SEQ_LEN, dtype=torch.long)
+        attention_mask[:, -1] = 0  # fake right-padding
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+
+class _FakeMLLM:
+    """Fake Qwen3VLModel: hidden states are a deterministic function of ids."""
+
+    dtype = torch.bfloat16
+
+    def __init__(self):
+        self.calls = []
+
+    def __call__(self, input_ids=None, attention_mask=None, output_hidden_states=False, **kwargs):
+        self.calls.append({"output_hidden_states": output_hidden_states, **kwargs})
+        base = input_ids.to(torch.float32).unsqueeze(-1).repeat(1, 1, _EMBED_DIM)
+        hidden = base + torch.arange(_EMBED_DIM, dtype=torch.float32)
+        return SimpleNamespace(last_hidden_state=hidden, hidden_states=(hidden - 1, hidden))
+
+
+def _make_encode_pipeline(num_instruction_feature_layers: int = 1):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        SYSTEM_PROMPT_4_T2I_UNIFIED,
+        SYSTEM_PROMPT_4_TI2I_UNIFIED,
+        BooguImagePipeline,
+    )
+
+    pipeline = object.__new__(BooguImagePipeline)
+    nn.Module.__init__(pipeline)
+    pipeline._execution_device = torch.device("cpu")
+    pipeline.processor = _RecordingProcessor()
+    pipeline.mllm = _FakeMLLM()
+    pipeline.transformer = SimpleNamespace(
+        instruction_feature_configs={
+            "instruction_feat_dim": _EMBED_DIM,
+            "num_instruction_feature_layers": num_instruction_feature_layers,
+            "reduce_type": "mean",
+        },
+        dtype=torch.float32,
+    )
+    pipeline.SYSTEM_PROMPT_4_T2I = SYSTEM_PROMPT_4_T2I_UNIFIED
+    pipeline.SYSTEM_PROMPT_DROP = SYSTEM_PROMPT_4_TI2I_UNIFIED
+    return pipeline
+
+
+def test_apply_chat_template_system_prompt_selection():
+    pipeline = _make_encode_pipeline()
+
+    messages = pipeline._apply_chat_template("a cat on a mat")
+    assert messages[0]["role"] == "system"
+    assert messages[0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_4_T2I
+    assert messages[1]["role"] == "user"
+    assert messages[1]["content"][0]["text"] == "a cat on a mat"
+
+    # Empty and whitespace-only instructions select the DROP prompt (this is
+    # the path the default negative prompt "" takes).
+    for empty in ("", "   ", None):
+        messages = pipeline._apply_chat_template(empty)
+        assert messages[0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_DROP
+
+
+def test_processor_called_with_upstream_kwargs():
+    pipeline = _make_encode_pipeline()
+    pipeline.encode_prompt("a dog", do_classifier_free_guidance=False)
+
+    (call,) = pipeline.processor.calls
+    kwargs = call["kwargs"]
+    assert kwargs["padding"] == "longest"
+    assert kwargs["padding_side"] == "right"
+    assert kwargs["truncation"] is False
+    assert kwargs["max_length"] == 1280  # upstream __call__ default
+    assert kwargs["tokenize"] is True
+    assert kwargs["return_dict"] is True
+    assert kwargs["return_tensors"] == "pt"
+
+
+def test_encode_prompt_shapes_and_dtype():
+    pipeline = _make_encode_pipeline()
+    embeds, mask, neg_embeds, neg_mask = pipeline.encode_prompt(["a dog", "a cat"])
+
+    assert embeds.shape == (2, _SEQ_LEN, _EMBED_DIM)
+    assert mask.shape == (2, _SEQ_LEN)
+    assert neg_embeds.shape == (2, _SEQ_LEN, _EMBED_DIM)
+    assert neg_mask.shape == (2, _SEQ_LEN)
+    # Cast to the MLLM dtype, mask passed through from the processor.
+    assert embeds.dtype == torch.bfloat16
+    assert neg_embeds.dtype == torch.bfloat16
+    assert torch.equal(mask[:, -1], torch.zeros(2, dtype=torch.long))
+
+
+def test_single_layer_encoding_requests_hidden_states_once():
+    pipeline = _make_encode_pipeline()
+
+    pipeline._get_instruction_feature_embeds("a dog")
+
+    assert pipeline.mllm.calls == [{"output_hidden_states": True, "return_dict": True}]
+
+
+def test_single_layer_encoding_propagates_mllm_failure():
+    pipeline = _make_encode_pipeline()
+    failure = RuntimeError("mllm failed")
+
+    class _FailingMLLM:
+        dtype = torch.bfloat16
+
+        def __init__(self):
+            self.calls = 0
+
+        def __call__(self, **kwargs):
+            self.calls += 1
+            raise failure
+
+    pipeline.mllm = _FailingMLLM()
+
+    with pytest.raises(RuntimeError, match="mllm failed") as exc_info:
+        pipeline._get_instruction_feature_embeds("a dog")
+
+    assert exc_info.value is failure
+    assert pipeline.mllm.calls == 1
+
+
+def test_encode_prompt_cfg_negative_default_is_empty_string():
+    pipeline = _make_encode_pipeline()
+    pipeline.encode_prompt("a dog")
+
+    assert len(pipeline.processor.calls) == 2
+    negative_messages = pipeline.processor.calls[1]["prompts"]
+    assert len(negative_messages) == 1
+    # The default "" negative prompt goes through the DROP system prompt.
+    assert negative_messages[0][0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_DROP
+    assert negative_messages[0][1]["content"][0]["text"] == ""
+
+
+def test_encode_prompt_without_cfg_skips_negative():
+    pipeline = _make_encode_pipeline()
+    embeds, mask, neg_embeds, neg_mask = pipeline.encode_prompt("a dog", do_classifier_free_guidance=False)
+
+    assert len(pipeline.processor.calls) == 1
+    assert embeds.shape == (1, _SEQ_LEN, _EMBED_DIM)
+    assert neg_embeds is None
+    assert neg_mask is None
+
+
+def test_encode_prompt_explicit_negative_prompt():
+    pipeline = _make_encode_pipeline()
+    pipeline.encode_prompt("a dog", negative_prompt="blurry, low quality")
+
+    negative_messages = pipeline.processor.calls[1]["prompts"]
+    assert negative_messages[0][0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_4_T2I
+    assert negative_messages[0][1]["content"][0]["text"] == "blurry, low quality"
+
+
+def test_encode_prompt_negative_batch_mismatch_raises():
+    pipeline = _make_encode_pipeline()
+    with pytest.raises(ValueError, match="batch size"):
+        pipeline.encode_prompt(["a dog", "a cat"], negative_prompt=["only one"])
+
+
+def test_encode_prompt_num_images_per_prompt_repeats_embeds():
+    pipeline = _make_encode_pipeline()
+    embeds, mask, neg_embeds, neg_mask = pipeline.encode_prompt("a dog", num_images_per_prompt=3)
+
+    assert embeds.shape == (3, _SEQ_LEN, _EMBED_DIM)
+    assert mask.shape == (3, _SEQ_LEN)
+    assert torch.equal(embeds[0], embeds[1])
+    assert torch.equal(embeds[0], embeds[2])
+    assert torch.equal(mask[0], mask[1])
+    assert neg_embeds.shape == (3, _SEQ_LEN, _EMBED_DIM)
+
+
+def test_encode_prompt_precomputed_embeds_bypass_encoder():
+    pipeline = _make_encode_pipeline()
+    precomputed = torch.randn(1, _SEQ_LEN, _EMBED_DIM)
+    precomputed_mask = torch.ones(1, _SEQ_LEN, dtype=torch.long)
+    neg_precomputed = torch.randn(1, _SEQ_LEN, _EMBED_DIM)
+
+    embeds, mask, neg_embeds, _ = pipeline.encode_prompt(
+        "ignored",
+        prompt_embeds=precomputed,
+        prompt_attention_mask=precomputed_mask,
+        negative_prompt_embeds=neg_precomputed,
+    )
+
+    assert len(pipeline.processor.calls) == 0
+    assert torch.equal(embeds, precomputed)
+    assert torch.equal(neg_embeds, neg_precomputed)
+
+
+def test_reshape_embeds_and_mask_list_branch():
+    # Multi-layer configs return a list of per-layer tensors; the reshape
+    # helper must handle both forms.
+    pipeline = _make_encode_pipeline()
+    layer = torch.arange(2 * _SEQ_LEN * _EMBED_DIM, dtype=torch.float32).view(2, _SEQ_LEN, _EMBED_DIM)
+    mask = torch.ones(2, _SEQ_LEN, dtype=torch.long)
+
+    batch_size, seq_len, reshaped, reshaped_mask = pipeline._reshape_embeds_and_mask([layer, layer + 1], mask, 2)
+
+    assert batch_size == 2
+    assert seq_len == _SEQ_LEN
+    assert isinstance(reshaped, list) and len(reshaped) == 2
+    assert reshaped[0].shape == (4, _SEQ_LEN, _EMBED_DIM)
+    # repeat(1, n, 1).view(b*n, ...) interleaves per-sample: rows [s0, s0, s1, s1]
+    assert torch.equal(reshaped[0][0], reshaped[0][1])
+    assert torch.equal(reshaped[0][2], reshaped[0][3])
+
+
+class _FakeTransformer:
+    """Fake denoiser: velocity 0 (identity flow) with the config the loop reads."""
+
+    in_channels = 4
+    axes_dim_rope = (8, 4, 4)
+    axes_lens = (32, 16, 16)
+    dtype = torch.float32
+    instruction_feature_configs = {
+        "instruction_feat_dim": _EMBED_DIM,
+        "num_instruction_feature_layers": 1,
+        "reduce_type": "mean",
+    }
+
+    def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
+        return torch.zeros_like(latents)
+
+
+class _FakeScheduler:
+    def set_timesteps(self, num_inference_steps, device=None, num_tokens=None):
+        self.timesteps = torch.linspace(0, 1, num_inference_steps + 1)[:-1]
+
+    def step(self, model_output, t, latents, return_dict=False):
+        return (latents,)
+
+
+class _FakeDecodeVAE:
+    dtype = torch.float32
+
+    def __init__(self):
+        self.config = SimpleNamespace(scaling_factor=1.0, shift_factor=0.0, block_out_channels=[128, 256, 512, 512])
+
+    def decode(self, latents, return_dict=False):
+        batch = latents.shape[0]
+        return (torch.zeros(batch, 3, 16, 16),)
+
+
+def _make_forward_pipeline():
+    pipeline = _make_encode_pipeline()
+    pipeline.transformer = _FakeTransformer()
+    pipeline.scheduler = _FakeScheduler()
+    pipeline.vae = _FakeDecodeVAE()
+    pipeline.vae_scale_factor = 8
+    pipeline.default_sample_size = 128
+    return pipeline
+
+
+def _make_request_batch(prompt, **sampling_overrides):
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    sampling = OmniDiffusionSamplingParams(**sampling_overrides)
+    return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
+
+
+def test_forward_returns_diffusion_output():
+    from vllm_omni.diffusion.data import DiffusionOutput
+
+    pipeline = _make_forward_pipeline()
+    req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2)
+
+    out = pipeline.forward(req)
+
+    assert isinstance(out, DiffusionOutput)
+    assert isinstance(out.output, torch.Tensor)
+    assert out.output.shape[0] == 1
+    assert torch.isfinite(out.output).all()
+    # CFG default is on (text guidance 4.0), so the encoder ran twice (pos + neg).
+    assert len(pipeline.processor.calls) == 2
+
+
+def test_forward_cfg_off_when_guidance_one():
+    pipeline = _make_forward_pipeline()
+    req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2, guidance_scale=1.0)
+    # guidance_scale=1.0 is falsy-adjacent but explicitly disables CFG; the
+    # request layer would set guidance_scale_provided, so emulate that here.
+    req.sampling_params.guidance_scale_provided = True
+
+    pipeline.forward(req)
+
+    # Only the positive prompt is encoded when CFG is off.
+    assert len(pipeline.processor.calls) == 1
+
+
+# ---------------------------------------------------------------------------
+# Editing / TI2I: pre-process function
+# ---------------------------------------------------------------------------
+
+
+def _make_edit_od_config(tmp_path, block_out_channels=(128, 256, 512, 512)):
+    import json
+
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir(parents=True, exist_ok=True)
+    (vae_dir / "config.json").write_text(json.dumps({"block_out_channels": list(block_out_channels)}))
+    return OmniDiffusionConfig(
+        model=str(tmp_path),
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        num_gpus=1,
+    )
+
+
+@pytest.mark.parametrize(
+    "factory_name",
+    ["get_boogu_image_pre_process_func", "get_boogu_image_post_process_func"],
+)
+@pytest.mark.parametrize("config_contents", [None, "{not valid json"])
+def test_process_factory_reports_invalid_vae_config(tmp_path, factory_name, config_contents):
+    import importlib
+
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    if config_contents is not None:
+        (vae_dir / "config.json").write_text(config_contents)
+
+    od_config = OmniDiffusionConfig(
+        model=str(tmp_path),
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        num_gpus=1,
+    )
+    factory = getattr(importlib.import_module(_MODULE), factory_name)
+
+    with pytest.raises(RuntimeError, match=r"Failed to load Boogu VAE config from .*vae/config\.json"):
+        factory(od_config)
+
+
+@pytest.mark.parametrize(("height", "width"), [(0, 16), (16, 0), (0, 0)])
+def test_image_processor_rejects_non_positive_dimensions(height, width):
+    from vllm_omni.diffusion.models.boogu_image.image_processor import BooguImageProcessor
+
+    image_processor = BooguImageProcessor()
+    image = torch.zeros(1, 3, 16, 16)
+
+    with pytest.raises(ValueError, match=rf"height={height}, width={width}"):
+        image_processor.get_new_height_width(image, height=height, width=width)
+
+
+def _make_diffusion_request(prompt, **sampling_overrides):
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    sampling = OmniDiffusionSamplingParams(**sampling_overrides)
+    return OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="req-0")
+
+
+def test_pre_process_no_image_is_noop(tmp_path):
+    import PIL.Image  # noqa: F401  (import guard: PIL must be available)
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        get_boogu_image_pre_process_func,
+    )
+
+    pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
+
+    # Text-to-image request: no multimodal image -> returned unchanged.
+    req = _make_diffusion_request({"prompt": "a cat"}, height=123, width=456)
+    out = pre(req)
+    assert "additional_information" not in out.prompt
+    assert out.sampling_params.height == 123
+    assert out.sampling_params.width == 456
+
+    # A plain-string prompt cannot carry an image either.
+    str_req = _make_diffusion_request("a cat")
+    assert pre(str_req).prompt == "a cat"
+
+
+def test_pre_process_populates_reference_and_align_res(tmp_path):
+    import PIL.Image
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        get_boogu_image_pre_process_func,
+    )
+
+    pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
+
+    image = PIL.Image.new("RGB", (1000, 500))  # (width, height)
+    req = _make_diffusion_request({"prompt": "make it winter", "multi_modal_data": {"image": image}})
+    out = pre(req)
+
+    ai = out.prompt["additional_information"]
+    assert "prompt_image" in ai and "preprocessed_image" in ai
+
+    # VLM copy is a PIL image, downscaled, never upscaled, aligned to 16.
+    prompt_image = ai["prompt_image"]
+    assert isinstance(prompt_image, PIL.Image.Image)
+    assert prompt_image.width % 16 == 0 and prompt_image.height % 16 == 0
+    assert prompt_image.width <= 1000 and prompt_image.height <= 500
+    assert max(prompt_image.width, prompt_image.height) <= 768
+
+    # VAE copy is a normalized [1, C, H, W] tensor aligned to 16.
+    vae = ai["preprocessed_image"]
+    assert isinstance(vae, torch.Tensor) and vae.ndim == 4 and vae.shape[0] == 1
+    assert vae.shape[-1] % 16 == 0 and vae.shape[-2] % 16 == 0
+    assert -1.0 <= float(vae.min()) and float(vae.max()) <= 1.0
+
+    # align_res: the request resolution follows the VAE-encoded reference dims.
+    assert out.sampling_params.height == vae.shape[-2]
+    assert out.sampling_params.width == vae.shape[-1]
+
+
+def test_pre_process_rejects_multiple_images(tmp_path):
+    import PIL.Image
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        get_boogu_image_pre_process_func,
+    )
+
+    pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
+    imgs = [PIL.Image.new("RGB", (64, 64)), PIL.Image.new("RGB", (64, 64))]
+    req = _make_diffusion_request({"prompt": "combine", "multi_modal_data": {"image": imgs}})
+
+    with pytest.raises(ValueError, match="single reference image"):
+        pre(req)
+
+
+def test_pre_process_single_image_in_list_is_accepted(tmp_path):
+    import PIL.Image
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        get_boogu_image_pre_process_func,
+    )
+
+    pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
+    req = _make_diffusion_request({"prompt": "edit", "multi_modal_data": {"image": [PIL.Image.new("RGB", (128, 128))]}})
+    out = pre(req)
+    assert "preprocessed_image" in out.prompt["additional_information"]
+
+
+# ---------------------------------------------------------------------------
+# Editing / TI2I: chat template + image-aware encoding
+# ---------------------------------------------------------------------------
+
+
+def _make_edit_encode_pipeline():
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        SYSTEM_PROMPT_4_TI2I_UNIFIED,
+    )
+
+    pipeline = _make_encode_pipeline()
+    pipeline.SYSTEM_PROMPT_4_TI2I = SYSTEM_PROMPT_4_TI2I_UNIFIED
+    pipeline.SYSTEM_PROMPT_4_I2I = SYSTEM_PROMPT_4_TI2I_UNIFIED
+    return pipeline
+
+
+def test_apply_chat_template_ti2i_places_image_before_text():
+    import PIL.Image
+
+    pipeline = _make_edit_encode_pipeline()
+    image = PIL.Image.new("RGB", (16, 16))
+
+    messages = pipeline._apply_chat_template("turn day into night", [image])
+    assert messages[0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_4_TI2I
+    user_content = messages[1]["content"]
+    # Image content comes first, then the instruction text.
+    assert user_content[0]["type"] == "image"
+    assert user_content[0]["image"] is image
+    assert user_content[-1] == {"type": "text", "text": "turn day into night"}
+
+    # Empty instruction with an image selects the I2I system prompt.
+    empty_messages = pipeline._apply_chat_template("", [image])
+    assert empty_messages[0]["content"][0]["text"] == pipeline.SYSTEM_PROMPT_4_I2I
+
+
+class _ImageAwareRecordingProcessor:
+    """Records whether reference images reached the processor."""
+
+    def __init__(self):
+        self.calls = []
+
+    def apply_chat_template(self, prompts, **kwargs):
+        has_image = []
+        for messages in prompts:
+            user_content = messages[1]["content"]
+            has_image.append(any(c.get("type") == "image" for c in user_content))
+        self.calls.append({"prompts": prompts, "kwargs": kwargs, "has_image": has_image})
+        batch = len(prompts)
+        input_ids = torch.arange(batch * _SEQ_LEN, dtype=torch.long).view(batch, _SEQ_LEN)
+        attention_mask = torch.ones(batch, _SEQ_LEN, dtype=torch.long)
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+
+def test_encode_prompt_attaches_images_to_positive_only():
+    import PIL.Image
+
+    pipeline = _make_edit_encode_pipeline()
+    pipeline.processor = _ImageAwareRecordingProcessor()
+    image = PIL.Image.new("RGB", (16, 16))
+
+    pipeline.encode_prompt(
+        "add a rainbow",
+        do_classifier_free_guidance=True,
+        input_images=[[image]],
+    )
+
+    # Positive call carries the image; the negative (CFG) call is image-free.
+    assert pipeline.processor.calls[0]["has_image"] == [True]
+    assert pipeline.processor.calls[1]["has_image"] == [False]
+
+
+# ---------------------------------------------------------------------------
+# Editing / TI2I: reference-latent VAE encode
+# ---------------------------------------------------------------------------
+
+
+class _FakeRefVAE:
+    """Fake VAE whose ``encode`` yields a known latent for shape/scaling checks."""
+
+    dtype = torch.float32
+
+    def __init__(self, scaling_factor=2.0, shift_factor=0.5):
+        self.config = SimpleNamespace(scaling_factor=scaling_factor, shift_factor=shift_factor)
+        self._latent = torch.ones(1, 4, 3, 5)
+
+    def encode(self, img):
+        dist = SimpleNamespace(sample=lambda generator=None: self._latent.clone())
+        return SimpleNamespace(latent_dist=dist)
+
+
+def test_build_ref_latents_shape_and_scaling():
+    pipeline = _make_encode_pipeline()
+    pipeline.vae = _FakeRefVAE(scaling_factor=2.0, shift_factor=0.5)
+
+    preprocessed = torch.zeros(1, 3, 48, 80)  # normalized image tensor
+    ref_latents = pipeline._build_ref_latents([preprocessed], num_images_per_prompt=1, device=torch.device("cpu"))
+
+    assert len(ref_latents) == 1
+    (sample_latents,) = ref_latents
+    assert isinstance(sample_latents, list) and len(sample_latents) == 1
+    latent = sample_latents[0]
+    # squeeze(0) -> [C, H, W]; (1 - shift) * scaling = (1 - 0.5) * 2 = 1.0
+    assert latent.shape == (4, 3, 5)
+    assert torch.allclose(latent, torch.ones(4, 3, 5))
+
+
+def test_build_ref_latents_expands_per_output_and_handles_none():
+    pipeline = _make_encode_pipeline()
+    pipeline.vae = _FakeRefVAE()
+
+    preprocessed = torch.zeros(1, 3, 48, 80)
+    ref_latents = pipeline._build_ref_latents([preprocessed, None], num_images_per_prompt=2, device=torch.device("cpu"))
+
+    # Two samples x 2 outputs each = 4 entries; the None sample stays None.
+    assert len(ref_latents) == 4
+    assert ref_latents[0] is ref_latents[1]  # same sample repeated
+    assert ref_latents[2] is None and ref_latents[3] is None
+
+
+# ---------------------------------------------------------------------------
+# Editing / TI2I: forward CFG branch selection
+# ---------------------------------------------------------------------------
+
+
+class _RecordingRefTransformer(_FakeTransformer):
+    """Counts predictions per step and records the reference-latent argument."""
+
+    def __init__(self):
+        self.calls = []
+
+    def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
+        self.calls.append(kwargs.get("ref_image_hidden_states"))
+        return torch.zeros_like(latents)
+
+
+class _EditForwardVAE(_FakeDecodeVAE):
+    """Adds a fake ``encode`` so the editing forward path can build ref latents."""
+
+    def encode(self, img):
+        dist = SimpleNamespace(sample=lambda generator=None: torch.zeros(1, 4, 8, 8))
+        return SimpleNamespace(latent_dist=dist)
+
+
+def _make_edit_forward_pipeline():
+    pipeline = _make_edit_encode_pipeline()
+    # Image-aware processor: reference images appear as ``{"type": "image"}``
+    # content entries, which the default text-first fake cannot parse.
+    pipeline.processor = _ImageAwareRecordingProcessor()
+    pipeline.transformer = _RecordingRefTransformer()
+    pipeline.scheduler = _FakeScheduler()
+    pipeline.vae = _EditForwardVAE()
+    pipeline.vae_scale_factor = 8
+    pipeline.default_sample_size = 128
+    return pipeline
+
+
+def _make_edit_request(**sampling_overrides):
+    import PIL.Image
+
+    image = PIL.Image.new("RGB", (64, 64))
+    prompt = {
+        "prompt": "make it winter",
+        "additional_information": {
+            "prompt_image": image,
+            "preprocessed_image": torch.zeros(1, 3, 64, 64),
+        },
+    }
+    sampling = _sampling(**sampling_overrides)
+    return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
+
+
+def _sampling(**overrides):
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    return OmniDiffusionSamplingParams(**overrides)
+
+
+def _count_per_step(calls, num_steps):
+    assert len(calls) % num_steps == 0
+    return len(calls) // num_steps
+
+
+def test_forward_ti2i_text_only_two_predictions_with_ref():
+    pipeline = _make_edit_forward_pipeline()
+    num_steps = 2
+    req = _make_edit_request(height=64, width=64, num_inference_steps=num_steps, guidance_scale=5.0)
+    req.sampling_params.guidance_scale_provided = True
+
+    pipeline.forward(req)
+
+    # Text-only ti2i: cond+ref and neg+ref -> 2 predictions/step, both carry ref.
+    assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
+    assert all(ref is not None for ref in pipeline.transformer.calls)
+
+
+def test_forward_ti2i_double_guidance_three_predictions():
+    pipeline = _make_edit_forward_pipeline()
+    num_steps = 2
+    req = _make_edit_request(
+        height=64, width=64, num_inference_steps=num_steps, guidance_scale=5.0, guidance_scale_2=2.0
+    )
+    req.sampling_params.guidance_scale_provided = True
+    req.sampling_params.guidance_scale_2_provided = True
+
+    pipeline.forward(req)
+
+    # Double guidance: cond+ref, neg+ref, neg+no-ref -> 3 predictions/step.
+    assert _count_per_step(pipeline.transformer.calls, num_steps) == 3
+    # Exactly one of the three per step drops the reference (neg+no-ref).
+    per_step = [pipeline.transformer.calls[i : i + 3] for i in range(0, len(pipeline.transformer.calls), 3)]
+    for step_calls in per_step:
+        assert sum(ref is None for ref in step_calls) == 1
+
+
+def test_forward_ti2i_image_only_two_predictions_drop_ref():
+    pipeline = _make_edit_forward_pipeline()
+    num_steps = 2
+    # text guidance 1.0 (off, provided) + image guidance 2.0 (provided).
+    req = _make_edit_request(
+        height=64, width=64, num_inference_steps=num_steps, guidance_scale=1.0, guidance_scale_2=2.0
+    )
+    req.sampling_params.guidance_scale_provided = True
+    req.sampling_params.guidance_scale_2_provided = True
+
+    pipeline.forward(req)
+
+    # Image-only ti2i: cond+ref and cond+no-ref -> 2 predictions/step.
+    assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
+    per_step = [pipeline.transformer.calls[i : i + 2] for i in range(0, len(pipeline.transformer.calls), 2)]
+    for step_calls in per_step:
+        assert sum(ref is None for ref in step_calls) == 1
+
+
+def test_forward_ti2i_no_guidance_single_prediction_with_ref():
+    pipeline = _make_edit_forward_pipeline()
+    num_steps = 2
+    # Both guidances off/unprovided -> image guidance forced to 1.0, no CFG.
+    req = _make_edit_request(height=64, width=64, num_inference_steps=num_steps, guidance_scale=1.0)
+    req.sampling_params.guidance_scale_provided = True
+
+    pipeline.forward(req)
+
+    assert _count_per_step(pipeline.transformer.calls, num_steps) == 1
+    assert all(ref is not None for ref in pipeline.transformer.calls)
+
+
+def test_forward_image_guidance_ignored_without_reference():
+    # guidance_scale_2 is set but there is no reference image (t2i request);
+    # image guidance must be forced off so this stays plain t2i CFG.
+    pipeline = _make_edit_forward_pipeline()
+    num_steps = 2
+    sampling = _sampling(height=64, width=64, num_inference_steps=num_steps, guidance_scale=5.0, guidance_scale_2=2.0)
+    sampling.guidance_scale_provided = True
+    sampling.guidance_scale_2_provided = True
+    req = SimpleNamespace(prompts=[{"prompt": "a cat"}], sampling_params=sampling)
+
+    pipeline.forward(req)
+
+    # t2i text CFG: cond + uncond -> 2 predictions/step, no ref anywhere.
+    assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
+    assert all(ref is None for ref in pipeline.transformer.calls)