[Bugfix] Preserve NextStep rectangular latent dimensions (#7853)

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

tests/diffusion/models/nextstep_1_1/test_nextstep_rectangular.py

@@ -0,0 +1,52 @@
+# 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)