1766 lines
65 KiB
Python
1766 lines
65 KiB
Python
# 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.
|
|
"""
|
|
|
|
import os
|
|
import zlib
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
from transformers import Qwen3VLConfig
|
|
|
|
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,
|
|
)
|
|
mllm_loader = mocker.patch(
|
|
f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
|
|
return_value=mllm_wrapper,
|
|
)
|
|
mllm_config_loader = mocker.patch(
|
|
f"{_MODULE}.Qwen3VLConfig.from_pretrained",
|
|
return_value=Qwen3VLConfig(),
|
|
)
|
|
mllm_builder = mocker.patch(
|
|
f"{_MODULE}.create_transformers_model_with_vllm_linears",
|
|
return_value=nn.Linear(1, 1),
|
|
)
|
|
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 only the dummy model id as local. Other filesystem checks (for
|
|
# example lazy imports in the quantization registry) must remain real.
|
|
path_exists = os.path.exists
|
|
mocker.patch("os.path.exists", side_effect=lambda path: str(path).startswith("dummy-boogu") or path_exists(path))
|
|
|
|
return {
|
|
"inner_encoder": inner_encoder,
|
|
"mllm_wrapper": mllm_wrapper,
|
|
"mllm_loader": mllm_loader,
|
|
"mllm_config_loader": mllm_config_loader,
|
|
"mllm_builder": mllm_builder,
|
|
"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")
|
|
assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
|
|
mock_dependencies["mllm_config_loader"].assert_not_called()
|
|
mock_dependencies["mllm_builder"].assert_not_called()
|
|
|
|
|
|
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_rejects_unsupported_mllm_quantization(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}
|
|
),
|
|
)
|
|
with pytest.raises(ValueError, match="Boogu MLLM only supports FP8 quantization"):
|
|
BooguImagePipeline(od_config=od_config)
|
|
mock_dependencies["mllm_loader"].assert_not_called()
|
|
mock_dependencies["mllm_config_loader"].assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("quantization_config", "quantize_mllm"),
|
|
[
|
|
pytest.param("fp8", True, id="global-fp8"),
|
|
pytest.param({"mllm": {"method": "fp8"}, "transformer": {"method": "fp8"}}, True, id="mllm-and-dit-fp8"),
|
|
pytest.param({"mllm": None, "transformer": "fp8"}, False, id="dit-only-fp8"),
|
|
],
|
|
)
|
|
def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm):
|
|
from transformers import AutoModel
|
|
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym
|
|
|
|
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
|
|
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
|
|
|
od_config = OmniDiffusionConfig(
|
|
model="dummy-boogu",
|
|
revision="test-revision",
|
|
tf_model_config=TransformerConfig(params={}),
|
|
dtype=torch.bfloat16,
|
|
quantization_config=quantization_config,
|
|
)
|
|
pipeline = BooguImagePipeline(od_config=od_config)
|
|
|
|
if quantize_mllm:
|
|
mock_dependencies["mllm_loader"].assert_not_called()
|
|
mock_dependencies["mllm_config_loader"].assert_called_once_with(
|
|
"dummy-boogu", subfolder="mllm", local_files_only=True, revision="test-revision"
|
|
)
|
|
builder = mock_dependencies["mllm_builder"]
|
|
builder.assert_called_once()
|
|
auto_cls, hf_config, online_quant_config = builder.call_args.args
|
|
assert auto_cls is AutoModel
|
|
assert hf_config is mock_dependencies["mllm_config_loader"].return_value
|
|
assert isinstance(online_quant_config, OnlineQuantizationConfig)
|
|
assert online_quant_config.args.linear.weight == kFp8Static128BlockSym
|
|
assert builder.call_args.kwargs == {
|
|
"dtype": torch.bfloat16,
|
|
"device": pipeline._execution_device,
|
|
"prefix": "mllm",
|
|
"skip_modules": ("mllm.visual",),
|
|
}
|
|
assert pipeline.mllm is builder.return_value
|
|
assert not pipeline.mllm.training
|
|
assert all(not parameter.requires_grad for parameter in pipeline.mllm.parameters())
|
|
transformer_source, mllm_source = pipeline.weights_sources
|
|
assert transformer_source.subfolder == "transformer"
|
|
assert mllm_source.model_or_path == "dummy-boogu"
|
|
assert mllm_source.subfolder == "mllm"
|
|
assert mllm_source.prefix == "mllm."
|
|
assert mllm_source.revision == "test-revision"
|
|
else:
|
|
assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
|
|
assert pipeline.mllm is mock_dependencies["inner_encoder"]
|
|
mock_dependencies["mllm_config_loader"].assert_not_called()
|
|
mock_dependencies["mllm_builder"].assert_not_called()
|
|
assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
|
|
assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], DiffusionFp8Config)
|
|
|
|
|
|
def test_constructor_preserves_mllm_ignored_layers(mock_dependencies):
|
|
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
|
|
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
|
|
|
quant_config = DiffusionFp8Config(
|
|
ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
|
|
)
|
|
od_config = OmniDiffusionConfig(
|
|
model="dummy-boogu",
|
|
tf_model_config=TransformerConfig(params={}),
|
|
dtype=torch.bfloat16,
|
|
quantization_config=quant_config,
|
|
)
|
|
BooguImagePipeline(od_config=od_config)
|
|
|
|
online_quant_config = mock_dependencies["mllm_builder"].call_args.args[2]
|
|
assert online_quant_config.ignored_layers == quant_config.ignored_layers
|
|
|
|
|
|
def test_constructor_preserves_serialized_mllm_quantization(mock_dependencies):
|
|
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
|
|
|
|
mock_dependencies["mllm_config_loader"].return_value = Qwen3VLConfig(
|
|
quantization_config={"quant_method": "fp8", "modules_to_not_convert": ["model.visual"]}
|
|
)
|
|
od_config = OmniDiffusionConfig(
|
|
model="dummy-boogu-fp8",
|
|
tf_model_config=TransformerConfig(params={}),
|
|
dtype=torch.bfloat16,
|
|
quantization_config="fp8",
|
|
)
|
|
pipeline = BooguImagePipeline(od_config=od_config)
|
|
|
|
# No override: HF reads the serialized checkpoint's scales and skip list.
|
|
assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
|
|
assert pipeline.mllm is mock_dependencies["inner_encoder"]
|
|
mock_dependencies["mllm_builder"].assert_not_called()
|
|
assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
|
|
|
|
|
|
def test_load_weights_maps_mllm_and_preserves_transformer():
|
|
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
|
|
|
|
pipeline = object.__new__(BooguImagePipeline)
|
|
nn.Module.__init__(pipeline)
|
|
pipeline.mllm = nn.ModuleDict(
|
|
{
|
|
"language_model": nn.Linear(2, 2, bias=False),
|
|
"visual": nn.Linear(2, 2, bias=False),
|
|
}
|
|
)
|
|
pipeline.transformer = nn.Linear(2, 2, bias=False)
|
|
language_weight = torch.full((2, 2), 1.0)
|
|
visual_weight = torch.full((2, 2), 2.0)
|
|
transformer_weight = torch.full((2, 2), 3.0)
|
|
|
|
loaded = pipeline.load_weights(
|
|
[
|
|
("mllm.lm_head.weight", torch.full((2, 2), 9.0)),
|
|
("mllm.model.language_model.weight", language_weight),
|
|
("mllm.model.visual.weight", visual_weight),
|
|
("transformer.weight", transformer_weight),
|
|
]
|
|
)
|
|
|
|
assert loaded == {"mllm.language_model.weight", "mllm.visual.weight", "transformer.weight"}
|
|
torch.testing.assert_close(pipeline.mllm["language_model"].weight, language_weight)
|
|
torch.testing.assert_close(pipeline.mllm["visual"].weight, visual_weight)
|
|
torch.testing.assert_close(pipeline.transformer.weight, transformer_weight)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("parallel_config", "cache_backend", "message"),
|
|
[
|
|
(DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"),
|
|
(DiffusionParallelConfig(ring_degree=2), "none", "Ulysses only"),
|
|
(
|
|
DiffusionParallelConfig(ulysses_degree=2, cfg_parallel_size=2),
|
|
"none",
|
|
"CFG parallelism is not validated",
|
|
),
|
|
(
|
|
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")
|
|
|
|
|
|
def test_constructor_accepts_ulysses_uaa(
|
|
mock_dependencies,
|
|
monkeypatch,
|
|
):
|
|
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
|
BooguImagePipeline,
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
f"{_MODULE}.current_omni_platform.is_cuda",
|
|
lambda: True,
|
|
)
|
|
od_config = OmniDiffusionConfig(
|
|
model="dummy-boogu",
|
|
tf_model_config=TransformerConfig(
|
|
params={
|
|
"num_attention_heads": 28,
|
|
"num_kv_heads": 7,
|
|
}
|
|
),
|
|
dtype=torch.float32,
|
|
parallel_config=DiffusionParallelConfig(
|
|
ulysses_degree=2,
|
|
ulysses_mode="advanced_uaa",
|
|
),
|
|
)
|
|
|
|
pipeline = BooguImagePipeline(od_config=od_config)
|
|
assert pipeline.transformer is mock_dependencies["transformer"]
|
|
|
|
|
|
def test_constructor_requires_uaa_for_boogu_gqa(
|
|
mock_dependencies,
|
|
monkeypatch,
|
|
):
|
|
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
|
BooguImagePipeline,
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
f"{_MODULE}.current_omni_platform.is_cuda",
|
|
lambda: True,
|
|
)
|
|
od_config = OmniDiffusionConfig(
|
|
model="dummy-boogu",
|
|
tf_model_config=TransformerConfig(
|
|
params={
|
|
"num_attention_heads": 28,
|
|
"num_kv_heads": 7,
|
|
}
|
|
),
|
|
dtype=torch.float32,
|
|
parallel_config=DiffusionParallelConfig(
|
|
ulysses_degree=2,
|
|
ulysses_mode="strict",
|
|
),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="advanced_uaa"):
|
|
BooguImagePipeline(od_config=od_config)
|
|
|
|
|
|
def test_constructor_rejects_non_cuda_sequence_parallelism(
|
|
mock_dependencies,
|
|
monkeypatch,
|
|
):
|
|
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
|
|
BooguImagePipeline,
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
f"{_MODULE}.current_omni_platform.is_cuda",
|
|
lambda: False,
|
|
)
|
|
od_config = OmniDiffusionConfig(
|
|
model="dummy-boogu",
|
|
tf_model_config=TransformerConfig(params={}),
|
|
dtype=torch.float32,
|
|
parallel_config=DiffusionParallelConfig(
|
|
ulysses_degree=2,
|
|
ulysses_mode="advanced_uaa",
|
|
),
|
|
)
|
|
|
|
with pytest.raises(NotImplementedError, match="requires CUDA"):
|
|
BooguImagePipeline(od_config=od_config)
|
|
|
|
mock_dependencies["mllm_wrapper"].model.to.assert_not_called()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Prompt-encoding tests (deterministic fakes, no constructor)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _content_id(text: str) -> int:
|
|
"""Bounded text IDs, stable across processes unlike builtin ``hash()``.
|
|
|
|
Seed/reference checks use separate transformers so BF16 rounding cannot
|
|
hide their smaller signals beneath the content term.
|
|
"""
|
|
return zlib.crc32(text.encode()) % 9973
|
|
|
|
|
|
class _RecordingProcessor:
|
|
"""Fake Qwen3VLProcessor with content-derived token IDs."""
|
|
|
|
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] = _content_id(system_text)
|
|
input_ids[i, 1] = _content_id(user_text)
|
|
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:
|
|
"""Record reference-image presence and encode text with stable IDs.
|
|
|
|
Find text by type because TI2I puts the image first.
|
|
"""
|
|
|
|
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.zeros(batch, _SEQ_LEN, dtype=torch.long)
|
|
for i, messages in enumerate(prompts):
|
|
system_text = messages[0]["content"][0]["text"]
|
|
user_text = next(c["text"] for c in messages[1]["content"] if c.get("type") == "text")
|
|
input_ids[i, 0] = _content_id(system_text)
|
|
input_ids[i, 1] = _content_id(user_text)
|
|
input_ids[i, 2:] = torch.arange(2, _SEQ_LEN) + i
|
|
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: generator routing, output split
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
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_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
|