53 lines
2.6 KiB
Python
53 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from types import MethodType, SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.nextstep_1_1 import pipeline_nextstep_1_1 as pipeline_module
|
|
from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep import NextStepModel
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
@pytest.mark.parametrize("height,width", [(32, 32), (32, 48), (48, 32)])
|
|
@pytest.mark.parametrize("batch", [1, 2])
|
|
def test_forward_preserves_latent_grid(monkeypatch, height, width, batch):
|
|
pipeline = object.__new__(pipeline_module.NextStep11Pipeline)
|
|
torch.nn.Module.__init__(pipeline)
|
|
pipeline._NextStep11Pipeline__device = torch.device("cpu")
|
|
pipeline.config = SimpleNamespace(use_gen_pos_embed=False)
|
|
pipeline.down_factor = 16
|
|
pipeline.scaling_factor = 2.0
|
|
pipeline.shift_factor = 0.25
|
|
grid = SimpleNamespace(config=SimpleNamespace(latent_channels=2, latent_patch_size=2))
|
|
expected = torch.arange(batch * 2 * (height // 8) * (width // 8), dtype=torch.float32)
|
|
expected = expected.reshape(batch, 2, height // 8, width // 8)
|
|
tokens = NextStepModel.patchify(grid, expected)
|
|
pipeline.model = Mock(
|
|
return_value=SimpleNamespace(last_hidden_state=torch.zeros(batch, 1, 4), past_key_values=None)
|
|
)
|
|
pipeline.model.forward_model.return_value = pipeline.model.return_value
|
|
pipeline.model.unpatchify = MethodType(NextStepModel.unpatchify, grid)
|
|
pipeline.tokenizer = Mock(
|
|
bos_token=None,
|
|
return_value=SimpleNamespace(
|
|
input_ids=torch.ones(batch, 2, dtype=torch.long), attention_mask=torch.ones(batch, 2, dtype=torch.long)
|
|
),
|
|
)
|
|
pipeline._check_input = Mock(return_value=(["teapot"], None))
|
|
pipeline._build_captions = Mock(return_value=(["teapot"] * batch, None, 1, 1.0))
|
|
pipeline._add_prefix_ids = lambda _hw, ids, mask: (ids, mask)
|
|
pipeline.decoding = Mock(return_value=tokens)
|
|
pipeline.vae = SimpleNamespace(dtype=torch.float32, decode=lambda latent: SimpleNamespace(sample=latent))
|
|
monkeypatch.setattr(pipeline_module, "StaticCache", lambda **_kwargs: None)
|
|
params = OmniDiffusionSamplingParams(
|
|
height=height, width=width, guidance_scale=1.0, num_outputs_per_prompt=batch, seed=None
|
|
)
|
|
output = pipeline.forward(SimpleNamespace(prompts=["teapot"], sampling_params=params))
|
|
torch.testing.assert_close(output.output, expected / 2 + 0.25, rtol=0, atol=0)
|