From 2c2cffd2aacda119e1708a6302163d9b9684f071 Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Mon, 21 Sep 2026 23:26:14 +0800
Subject: [PATCH] [Bugfix] Preserve NextStep rectangular latent dimensions
 (#7853)

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 .../nextstep_1_1/test_nextstep_rectangular.py | 52 +++++++++++++++++++
 .../nextstep_1_1/pipeline_nextstep_1_1.py     |  4 +-
 2 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 tests/diffusion/models/nextstep_1_1/test_nextstep_rectangular.py

diff --git a/tests/diffusion/models/nextstep_1_1/test_nextstep_rectangular.py b/tests/diffusion/models/nextstep_1_1/test_nextstep_rectangular.py
new file mode 100644
index 000000000..279feb445
--- /dev/null
+++ b/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)
diff --git a/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py
index 392e42868..9a3358d0d 100644
--- a/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py
+++ b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Adapted from NextStep-1.1 (https://huggingface.co/stepfun-ai/NextStep-1.1)
 
 import os
@@ -704,7 +704,7 @@ class NextStep11Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelin
         )
 
         # Unpatchify
-        latents = self.model.unpatchify(tokens)
+        latents = self.model.unpatchify(tokens, h=height // self.down_factor, w=width // self.down_factor)
         latents = (latents / self.scaling_factor) + self.shift_factor
 
         # Decode latents