# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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_cls": mock_transformer_cls,
        "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, BooguImageTurboPipeline

    assert BooguImagePipeline is not None
    assert issubclass(BooguImageTurboPipeline, BooguImagePipeline)


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_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"
    assert source.subfolder == "transformer"
    assert source.prefix == "transformer."
    assert source.fall_back_to_pt is True


def test_constructor_routes_only_transformer_config(mock_dependencies, mocker):
    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
    from vllm_omni.quantization.component_config import ComponentQuantizationConfig

    transformer_config = mocker.MagicMock(spec=QuantizationConfig)
    encoder_config = mocker.MagicMock(spec=QuantizationConfig)
    od_config = OmniDiffusionConfig(
        model="dummy-boogu",
        tf_model_config=TransformerConfig(params={}),
        dtype=torch.bfloat16,
        quantization_config=ComponentQuantizationConfig(
            {"transformer": transformer_config, "mllm": encoder_config, "vae": None}
        ),
    )
    BooguImagePipeline(od_config=od_config)
    kwargs = mock_dependencies["transformer_cls"].call_args.kwargs
    assert kwargs["quant_config"] is transformer_config
    assert kwargs["prefix"] == "transformer"


@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(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()


@pytest.mark.parametrize("cfg_parallel_size", [2, 3])
def test_constructor_accepts_cfg_parallel(mock_dependencies, cfg_parallel_size):
    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=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size),
    )

    pipeline = BooguImagePipeline(od_config=od_config)

    assert pipeline.od_config.parallel_config.cfg_parallel_size == cfg_parallel_size
    assert hasattr(pipeline, "predict_noise_maybe_with_cfg")
    assert hasattr(pipeline, "predict_noise_with_multi_branch_cfg")


# ---------------------------------------------------------------------------
# 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_real, instruction_attention_mask, **kwargs):
        assert all(not freqs.is_complex() for pair in freqs_real for freqs in pair)
        return torch.zeros_like(latents)


class _FakeScheduler:
    def __init__(self):
        self.set_calls = 0
        self.step_calls = 0

    def set_timesteps(self, num_inference_steps, device=None, num_tokens=None):
        self.set_calls += 1
        self.timesteps = torch.linspace(0, 1, num_inference_steps + 1)[:-1]

    def step(self, model_output, t, latents, return_dict=False):
        self.step_calls += 1
        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_cls=None):
    if pipeline_cls is None:
        from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline

        pipeline_cls = BooguImagePipeline
    base_pipeline = _make_encode_pipeline()
    pipeline = object.__new__(pipeline_cls)
    nn.Module.__init__(pipeline)
    pipeline.__dict__.update(base_pipeline.__dict__)
    pipeline.transformer = _FakeTransformer()
    pipeline.scheduler = _FakeScheduler()
    pipeline.vae = _FakeDecodeVAE()
    pipeline.vae_scale_factor = 8
    pipeline.default_sample_size = 128
    return pipeline


def _wrap_request_batch(items):
    """Wrap ``(prompt, sampling_params)`` pairs in a real DiffusionRequestBatch.

    The pipeline's request-batch ``forward`` reads ``sampling_params_list`` and
    ``collate_request_generators``, so the fake single-namespace shim no longer
    satisfies the contract; construct the genuine wrapper instead.
    """
    from vllm_omni.diffusion.request import OmniDiffusionRequest
    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

    requests = [
        OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id=f"req-{i}")
        for i, (prompt, sampling) in enumerate(items)
    ]
    return DiffusionRequestBatch(requests=requests)


def _make_request_batch(prompt, **sampling_overrides):
    return _wrap_request_batch([(prompt, _sampling(**sampling_overrides))])


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)

    outs = pipeline.forward(req)

    # Request-batch forward returns one DiffusionOutput per request.
    assert isinstance(outs, list) and len(outs) == 1
    out = outs[0]
    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
    assert pipeline.scheduler.set_calls == 1
    assert pipeline.scheduler.step_calls == 2


def test_forward_cfg_off_when_guidance_one():
    pipeline = _make_forward_pipeline()
    cfg_calls = []
    original_predict_noise_maybe_with_cfg = pipeline.predict_noise_maybe_with_cfg

    def record_cfg_call(**kwargs):
        cfg_calls.append(kwargs)
        return original_predict_noise_maybe_with_cfg(**kwargs)

    pipeline.predict_noise_maybe_with_cfg = record_cfg_call
    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
    assert len(cfg_calls) == 2
    assert all(call["do_true_cfg"] is False for call in cfg_calls)
    assert all(call["negative_kwargs"] is None for call in cfg_calls)


def test_double_guidance_combine_matches_legacy_formula():
    pipeline = _make_forward_pipeline()
    positive_with_reference = torch.randn(2, 4, 8, 8)
    negative_with_reference = torch.randn(2, 4, 8, 8)
    uncond = torch.randn(2, 4, 8, 8)
    text_scale = 5.0
    image_scale = 2.0

    actual = pipeline.combine_multi_branch_cfg_noise(
        [positive_with_reference, negative_with_reference, uncond],
        {"text": text_scale, "image": image_scale},
    )
    legacy = (
        positive_with_reference
        + (text_scale - 1.0) * (positive_with_reference - negative_with_reference)
        + (image_scale - 1.0) * (negative_with_reference - uncond)
    )

    torch.testing.assert_close(actual, legacy, rtol=0, atol=1e-5)


class _RecordingDMDTransformer(_FakeTransformer):
    def __init__(self, velocity=0.0):
        self.timesteps = []
        self.refs = []
        self.freqs_real_calls = []
        self.velocity = velocity

    def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
        self.timesteps.append(float(timestep[0]))
        self.refs.append(kwargs.get("ref_image_hidden_states"))
        self.freqs_real_calls.append(freqs_real)
        return torch.full_like(latents, self.velocity)


def _make_turbo_forward_pipeline():
    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImageTurboPipeline

    pipeline = _make_forward_pipeline(BooguImageTurboPipeline)
    pipeline.transformer = _RecordingDMDTransformer(velocity=0.125)
    return pipeline


def test_turbo_dmd_uses_four_steps_with_real_rope_without_regular_scheduler():
    pipeline = _make_turbo_forward_pipeline()
    renoise_calls = []
    original_renoise = pipeline._renoise_dmd_latents

    def recording_renoise(latents, sigma, generator=None):
        renoise_calls.append(sigma)
        return original_renoise(latents, sigma, generator)

    pipeline._renoise_dmd_latents = recording_renoise
    req = _make_request_batch(
        "a cat",
        height=64,
        width=64,
        output_type="latent",
        generator=torch.Generator(device="cpu").manual_seed(7),
    )

    out = pipeline.forward(req)[0]

    assert out.output.shape[0] == 1
    assert len(pipeline.transformer.timesteps) == 4
    assert len(pipeline.transformer.freqs_real_calls) == 4
    assert all(
        not freq.is_complex()
        for freqs_real in pipeline.transformer.freqs_real_calls
        for pair in freqs_real
        for freq in pair
    )
    expected_sigmas = torch.linspace(0.001, 1.0, 5, dtype=torch.bfloat16)[:-1].float().tolist()
    assert pipeline.transformer.timesteps == pytest.approx(expected_sigmas)
    assert pipeline.scheduler.set_calls == 0
    assert pipeline.scheduler.step_calls == 0
    assert len(renoise_calls) == 3  # The final x0 must not be re-noised.
    assert len(pipeline.processor.calls) == 1  # DMD has no negative/CFG encoding.


def test_turbo_dmd_decode_uses_vae_attention_context():
    pipeline = _make_turbo_forward_pipeline()
    events = []

    class _RecordingContext:
        def __init__(self, device):
            self.device = device

        def __enter__(self):
            events.append(("enter", self.device.type))

        def __exit__(self, exc_type, exc_value, traceback):
            events.append(("exit", self.device.type))

    pipeline._vae_attention_context = lambda device: _RecordingContext(device)
    original_decode = pipeline.vae.decode

    def recording_decode(latents, return_dict=False):
        events.append(("decode", latents.device.type))
        return original_decode(latents, return_dict=return_dict)

    pipeline.vae.decode = recording_decode
    req = _make_request_batch(
        "a cat",
        height=64,
        width=64,
        output_type="pt",
        generator=torch.Generator(device="cpu").manual_seed(7),
    )

    pipeline.forward(req)

    assert events == [("enter", "cpu"), ("decode", "cpu"), ("exit", "cpu")]


def test_turbo_dmd_step_and_renoise_match_upstream_equations():
    from diffusers.utils.torch_utils import randn_tensor

    pipeline = _make_turbo_forward_pipeline()
    latents = torch.full((2, 4, 3, 3), 2.0)
    embeds = torch.zeros(2, _SEQ_LEN, _EMBED_DIM)
    mask = torch.ones(2, _SEQ_LEN, dtype=torch.long)

    x0 = pipeline._predict_dmd_student_step(latents, torch.tensor(0.25), embeds, None, mask)
    assert torch.allclose(x0, latents + 0.75 * 0.125)

    expected_generator = torch.Generator(device="cpu").manual_seed(13)
    expected_noise = randn_tensor(x0.shape, generator=expected_generator, device=x0.device, dtype=x0.dtype)
    expected = 0.4 * expected_noise + 0.6 * x0

    actual_generator = torch.Generator(device="cpu").manual_seed(13)
    actual = pipeline._renoise_dmd_latents(x0, torch.tensor(0.6), actual_generator)
    assert torch.allclose(actual, expected, rtol=1e-6, atol=1e-7)


def test_turbo_dmd_seeded_renoise_is_deterministic():
    def run_once():
        pipeline = _make_turbo_forward_pipeline()
        req = _make_request_batch(
            "a cat",
            height=64,
            width=64,
            output_type="latent",
            generator=torch.Generator(device="cpu").manual_seed(11),
        )
        return pipeline.forward(req)[0].output

    assert torch.equal(run_once(), run_once())


def test_turbo_dmd_custom_timesteps_and_validation():
    pipeline = _make_turbo_forward_pipeline()
    sigmas = pipeline._build_dmd_student_sigmas(99, torch.device("cpu"), torch.float32, 0.001, [1000, 500, 0])
    assert torch.equal(sigmas, torch.tensor([1.0, 0.5, 0.0]))

    with pytest.raises(ValueError, match="non-empty 1D"):
        pipeline._build_dmd_student_sigmas(4, torch.device("cpu"), torch.float32, 0.001, [])
    with pytest.raises(ValueError, match="non-empty 1D"):
        pipeline._build_dmd_student_sigmas(4, torch.device("cpu"), torch.float32, 0.001, [[0.1]])
    with pytest.raises(ValueError, match="finite values"):
        pipeline._build_dmd_student_sigmas(4, torch.device("cpu"), torch.float32, 0.001, [0.0, 1001.0])


def test_turbo_dmd_conditioning_sigma_override_reaches_forward_loop():
    pipeline = _make_turbo_forward_pipeline()
    req = _make_request_batch(
        "a cat",
        height=64,
        width=64,
        output_type="latent",
        extra_args={"dmd_conditioning_sigma": 0.125},
    )

    pipeline.forward(req)

    assert pipeline.transformer.timesteps[0] == pytest.approx(0.125)


def test_turbo_dmd_rejects_timesteps_and_sigmas_together():
    pipeline = _make_turbo_forward_pipeline()
    req = _make_request_batch(
        "a cat",
        height=64,
        width=64,
        timesteps=torch.tensor([0.0, 0.5]),
        sigmas=[0.0, 0.5],
    )
    with pytest.raises(ValueError, match="only one of timesteps or sigmas"):
        pipeline.forward(req)


def test_turbo_dmd_rejects_zero_steps():
    pipeline = _make_turbo_forward_pipeline()
    req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=0)
    with pytest.raises(ValueError, match="num_inference_steps must be >= 1"):
        pipeline.forward(req)


def test_turbo_dmd_normalizes_internal_dummy_guidance_only():
    pipeline = _make_turbo_forward_pipeline()
    req = _make_request_batch(
        "dummy run",
        height=64,
        width=64,
        num_inference_steps=1,
        guidance_scale=0.0,
        guidance_scale_2=0.0,
        output_type="latent",
    )
    req.is_dummy_run = lambda: True

    pipeline.forward(req)

    assert len(pipeline.transformer.timesteps) == 1
    assert len(pipeline.processor.calls) == 1


@pytest.mark.parametrize(
    "sampling_overrides",
    [
        {"guidance_scale": 0.0},
        {"guidance_scale": 2.0},
        {"guidance_scale_2": 2.0},
        {"extra_args": {"empty_instruction_guidance_scale": 1.0}},
    ],
)
def test_turbo_dmd_rejects_non_unit_guidance(sampling_overrides):
    pipeline = _make_turbo_forward_pipeline()
    req = _make_request_batch("a cat", height=64, width=64, **sampling_overrides)
    with pytest.raises(ValueError, match="requires guidance_scale=1.0"):
        pipeline.forward(req)


# ---------------------------------------------------------------------------
# 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


def test_vae_attention_context_is_noop_on_cpu():
    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
        BooguImagePipeline,
    )

    with BooguImagePipeline._vae_attention_context(torch.device("cpu")):
        pass


def test_vae_attention_context_prioritizes_efficient_cuda(monkeypatch):
    from vllm_omni.diffusion.models.boogu_image import pipeline_boogu_image

    recorded = {}

    def fake_sdpa_kernel(backends, *, set_priority):
        recorded["backends"] = backends
        recorded["set_priority"] = set_priority
        return pipeline_boogu_image.nullcontext()

    monkeypatch.setattr(pipeline_boogu_image, "sdpa_kernel", fake_sdpa_kernel)

    with pipeline_boogu_image.BooguImagePipeline._vae_attention_context(torch.device("cuda")):
        pass

    assert recorded == {
        "backends": [
            pipeline_boogu_image.SDPBackend.EFFICIENT_ATTENTION,
            pipeline_boogu_image.SDPBackend.MATH,
        ],
        "set_priority": True,
    }


# ---------------------------------------------------------------------------
# Editing / TI2I: forward CFG branch selection
# ---------------------------------------------------------------------------


class _RecordingRefTransformer(_FakeTransformer):
    """Counts predictions per step and records the reference-latent argument."""

    def __init__(self):
        self.calls = []
        self.instruction_embed_calls = []

    def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
        self.calls.append(kwargs.get("ref_image_hidden_states"))
        self.instruction_embed_calls.append(instruction_embeds)
        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 _wrap_request_batch([(prompt, 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)
    for i in range(0, len(pipeline.transformer.instruction_embed_calls), 2):
        positive, negative = pipeline.transformer.instruction_embed_calls[i : i + 2]
        assert positive is not negative


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
    embed_calls = pipeline.transformer.instruction_embed_calls
    for i in range(0, len(embed_calls), 3):
        positive, negative_with_reference, uncond = embed_calls[i : i + 3]
        assert positive is not negative_with_reference
        assert negative_with_reference is uncond


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
    embed_calls = pipeline.transformer.instruction_embed_calls
    for i in range(0, len(embed_calls), 2):
        positive_with_reference, positive_without_reference = embed_calls[i : i + 2]
        assert positive_with_reference is positive_without_reference


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 = _wrap_request_batch([({"prompt": "a cat"}, 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)


# ---------------------------------------------------------------------------
# Request-batch: compatibility key, generator routing, output split
# ---------------------------------------------------------------------------


def test_boogu_batch_compatibility_key_t2i_stable_ti2i_unique():
    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import _boogu_batch_compatibility_key

    # t2i: request_id does not enter the key -> t2i requests batch together.
    assert _boogu_batch_compatibility_key(False, "req-a") == _boogu_batch_compatibility_key(False, "req-b")
    assert _boogu_batch_compatibility_key(False, "req-a")[1] == "t2i"

    # ti2i: request_id in the key -> each edit gets a unique key, never co-batched.
    assert _boogu_batch_compatibility_key(True, "req-a") != _boogu_batch_compatibility_key(True, "req-b")
    assert _boogu_batch_compatibility_key(True, "req-a")[1] == "ti2i"

    # t2i and ti2i never share a key.
    assert _boogu_batch_compatibility_key(False, "req-a") != _boogu_batch_compatibility_key(True, "req-a")


def test_pre_process_key_wiring_t2i_batches_ti2i_isolated(tmp_path):
    # End-to-end: real pre-process sets request.batch_compatibility_key, and the
    # scheduler's key builder reads it into condition_key. Two t2i requests share
    # a key (co-batchable); two edit requests get distinct keys (batch=1).
    import PIL.Image

    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import get_boogu_image_pre_process_func
    from vllm_omni.diffusion.request import OmniDiffusionRequest
    from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams

    pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))

    def condition_key(prompt, rid):
        req = OmniDiffusionRequest(
            prompt=prompt, sampling_params=OmniDiffusionSamplingParams(height=512, width=512), request_id=rid
        )
        pre(req)
        return build_request_batch_sampling_params_key(req).condition_key

    t2i_a = condition_key({"prompt": "a cat"}, "t-a")
    t2i_b = condition_key({"prompt": "a dog"}, "t-b")
    assert t2i_a == t2i_b and t2i_a[1] == "t2i"

    img = PIL.Image.new("RGB", (64, 64))
    ti2i_a = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-a")
    ti2i_b = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-b")
    assert ti2i_a != ti2i_b and ti2i_a[1] == "ti2i"


class _GeneratorRecordingVAE:
    """Fake VAE that records the generator passed to each latent ``sample()``."""

    dtype = torch.float32

    def __init__(self):
        self.config = SimpleNamespace(scaling_factor=1.0, shift_factor=0.0)
        self.seen_generators = []

    def encode(self, img):
        def sample(generator=None):
            self.seen_generators.append(generator)
            return torch.zeros(1, 4, 2, 2)

        return SimpleNamespace(latent_dist=SimpleNamespace(sample=sample))


def test_build_ref_latents_uses_per_request_generators():
    pipeline = _make_encode_pipeline()
    pipeline.vae = _GeneratorRecordingVAE()

    g0 = torch.Generator().manual_seed(0)
    g1 = torch.Generator().manual_seed(1)
    preprocessed = [torch.zeros(1, 3, 16, 16), None, torch.zeros(1, 3, 16, 16)]

    pipeline._build_ref_latents(
        preprocessed,
        num_images_per_prompt=1,
        device=torch.device("cpu"),
        generators=[g0, None, g1],
    )

    # Each present reference row samples with its own request's generator; the
    # None reference row consumes none.
    assert pipeline.vae.seen_generators == [g0, g1]


def test_forward_request_batch_collates_generators_and_splits_outputs():
    from vllm_omni.diffusion.data import DiffusionOutput

    pipeline = _make_forward_pipeline()
    recorded = {}

    def fake_encode_prompt(prompt, **kwargs):
        n = len(prompt)
        embeds = torch.zeros(n, _SEQ_LEN, _EMBED_DIM)
        mask = torch.ones(n, _SEQ_LEN, dtype=torch.long)
        return embeds, mask, None, None

    def fake_prepare_latents(batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
        recorded["generator"] = generator
        # A distinct constant per batch row so the split can be verified.
        rows = torch.arange(batch_size, dtype=torch.float32).view(batch_size, 1, 1, 1)
        return rows.repeat(1, num_channels_latents, 2, 2)

    pipeline.encode_prompt = fake_encode_prompt
    pipeline.prepare_latents = fake_prepare_latents

    # T2I, CFG off, 1 step, latent output -> avoids VAE decode and CFG branches.
    s0 = _sampling(height=64, width=64, num_inference_steps=1, guidance_scale=1.0, output_type="latent")
    s1 = _sampling(height=64, width=64, num_inference_steps=1, guidance_scale=1.0, output_type="latent")
    req = _wrap_request_batch([("a cat", s0), ("a dog", s1)])
    g0 = torch.Generator().manual_seed(0)
    g1 = torch.Generator().manual_seed(1)
    req.requests[0].sampling_params.generator = g0
    req.requests[1].sampling_params.generator = g1

    outs = pipeline.forward(req)

    # forward() collates per-request generators rather than reusing req0's.
    assert recorded["generator"] == [g0, g1]
    # One DiffusionOutput per request, in scheduler order.
    assert len(outs) == 2
    assert all(isinstance(o, DiffusionOutput) for o in outs)
    assert float(outs[0].output[0, 0, 0, 0]) == 0.0
    assert float(outs[1].output[0, 0, 0, 0]) == 1.0


def test_reshape_mask_is_request_major_with_distinct_masks():
    # B=2, N=2 with per-request-distinct masks: reshaped rows must be
    # request-major [p0, p0, p1, p1] (repeat_interleave) to match the embed
    # layout, not tiled [p0, p1, p0, p1] (a plain repeat).
    pipeline = _make_encode_pipeline()
    embeds = torch.zeros(2, _SEQ_LEN, _EMBED_DIM)
    mask = torch.tensor([[1, 1, 1, 1, 1, 0], [1, 1, 1, 0, 0, 0]], dtype=torch.long)
    _, _, _, reshaped_mask = pipeline._reshape_embeds_and_mask(embeds, mask, 2)
    assert reshaped_mask.shape == (4, _SEQ_LEN)
    assert torch.equal(reshaped_mask[0], mask[0])
    assert torch.equal(reshaped_mask[1], mask[0])
    assert torch.equal(reshaped_mask[2], mask[1])
    assert torch.equal(reshaped_mask[3], mask[1])


def test_forward_request_batch_num_outputs_slices_and_generators():
    pipeline = _make_forward_pipeline()
    recorded = {}

    def fake_encode_prompt(prompt, **kwargs):
        n = len(prompt)
        return torch.zeros(n, _SEQ_LEN, _EMBED_DIM), torch.ones(n, _SEQ_LEN, dtype=torch.long), None, None

    def fake_prepare_latents(batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
        recorded["generator"] = generator
        rows = torch.arange(batch_size, dtype=torch.float32).view(batch_size, 1, 1, 1)
        return rows.repeat(1, num_channels_latents, 2, 2)

    pipeline.encode_prompt = fake_encode_prompt
    pipeline.prepare_latents = fake_prepare_latents

    kw = dict(
        height=64, width=64, num_inference_steps=1, guidance_scale=1.0, output_type="latent", num_outputs_per_prompt=2
    )
    req = _wrap_request_batch([("a cat", _sampling(**kw)), ("a dog", _sampling(**kw))])
    g0 = torch.Generator().manual_seed(0)
    g1 = torch.Generator().manual_seed(1)
    req.requests[0].sampling_params.generator = g0
    req.requests[1].sampling_params.generator = g1

    outs = pipeline.forward(req)

    # collated generator is request-major, output-minor.
    assert recorded["generator"] == [g0, g0, g1, g1]
    # each request gets its own 2-output slice: req0 rows [0,1], req1 rows [2,3].
    assert len(outs) == 2
    assert outs[0].output.shape[0] == 2 and outs[1].output.shape[0] == 2
    assert float(outs[0].output[0, 0, 0, 0]) == 0.0 and float(outs[0].output[1, 0, 0, 0]) == 1.0
    assert float(outs[1].output[0, 0, 0, 0]) == 2.0 and float(outs[1].output[1, 0, 0, 0]) == 3.0


def test_forward_batch_isolation_partner_content_and_seed():
    """CFG-on, B=2: request A's output must not change when only the
    co-batched partner's prompt content, negative prompt, or seed changes.

    ``_FakeTransformer``/``_FakeScheduler`` are content-blind (always-zero
    velocity, latents passed through unchanged), so they cannot catch a
    cross-request value leak. This test swaps in a transformer whose output
    depends on both ``instruction_embeds`` (content, positive or negative
    depending on which CFG branch called it) and ``latents`` (seed), and a
    scheduler that actually applies the predicted velocity, so a batching bug
    that mixes rows in either the cond or uncond prediction would change A's
    result. A negative-prompt-only perturbation is required to cover the
    uncond branch: varying only the positive prompt never touches
    ``negative_instruction_embeds``, so an earlier version of this test
    passed even with a synthetic row-mixing bug injected into the uncond
    predict() call (verified via a RED-arm check before this fix).
    """

    class _ContentAwareTransformer(_FakeTransformer):
        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
            content = instruction_embeds.mean(dim=(1, 2)).view(-1, 1, 1, 1)
            return latents + content

    class _ApplyingScheduler(_FakeScheduler):
        def step(self, model_output, t, latents, return_dict=False):
            return (model_output,)

    def run(prompt_a, seed_a, neg_a, prompt_b, seed_b, neg_b):
        pipeline = _make_forward_pipeline()
        pipeline.transformer = _ContentAwareTransformer()
        pipeline.scheduler = _ApplyingScheduler()
        kw = dict(height=64, width=64, num_inference_steps=2, guidance_scale=4.0, output_type="latent")
        req = _wrap_request_batch(
            [
                (
                    {"prompt": prompt_a, "negative_prompt": neg_a},
                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)),
                ),
                (
                    {"prompt": prompt_b, "negative_prompt": neg_b},
                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)),
                ),
            ]
        )
        return pipeline.forward(req)[0].output

    baseline = run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "blurry")
    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a totally different scene", 2, "blurry"))
    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 999, "blurry"))
    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "watermark"))


def test_forward_batched_ti2i_fails_closed():
    # A batched ti2i must fail closed (it is gated to batch=1).
    pipeline = _make_forward_pipeline()

    def edit_prompt():
        return {
            "prompt": "make it winter",
            "additional_information": {"preprocessed_image": torch.zeros(1, 3, 64, 64), "prompt_image": None},
        }

    req = _wrap_request_batch(
        [
            (edit_prompt(), _sampling(num_inference_steps=1, guidance_scale=1.0)),
            (edit_prompt(), _sampling(num_inference_steps=1, guidance_scale=1.0)),
        ]
    )
    with pytest.raises(RuntimeError, match="gated to batch=1"):
        pipeline.forward(req)


def test_supports_request_batch_enabled():
    from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline

    assert BooguImagePipeline.supports_request_batch is True


def test_turbo_request_batching_disabled_and_fails_closed():
    from vllm_omni.diffusion.models.boogu_image import BooguImageTurboPipeline

    assert BooguImageTurboPipeline.supports_request_batch is False

    pipeline = _make_turbo_forward_pipeline()
    sampling = dict(height=64, width=64, num_inference_steps=4, guidance_scale=1.0, output_type="latent")
    req = _wrap_request_batch(
        [
            ("a cat", _sampling(**sampling)),
            ("a dog", _sampling(**sampling)),
        ]
    )

    with pytest.raises(RuntimeError, match="does not support request batching"):
        pipeline.forward(req)


def _make_turbo_edit_forward_pipeline():
    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImageTurboPipeline

    base_pipeline = _make_edit_forward_pipeline()
    pipeline = object.__new__(BooguImageTurboPipeline)
    nn.Module.__init__(pipeline)
    pipeline.__dict__.update(base_pipeline.__dict__)
    pipeline.transformer = _RecordingDMDTransformer()
    return pipeline


def test_turbo_dmd_ti2i_keeps_reference_latents_and_uses_zero_conditioning_sigma():
    pipeline = _make_turbo_edit_forward_pipeline()
    req = _make_edit_request(
        height=64,
        width=64,
        num_inference_steps=4,
        guidance_scale=1.0,
        output_type="latent",
        generator=torch.Generator(device="cpu").manual_seed(3),
    )
    req.sampling_params.guidance_scale_provided = True

    pipeline.forward(req)

    assert len(pipeline.transformer.timesteps) == 4
    assert pipeline.transformer.timesteps[0] == pytest.approx(0.0)
    assert all(ref is not None for ref in pipeline.transformer.refs)
    assert pipeline.scheduler.set_calls == 0
    assert pipeline.scheduler.step_calls == 0