[diffusion] fix: crop GLM-Image output to requested size (#33859)

Co-authored-by: AuFlow <AuFlow@users.noreply.github.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>

python/sglang/multimodal_gen/configs/sample/glmimage.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -16,10 +16,21 @@ class GlmImageSamplingParams(SamplingParams):
     guidance_scale: float = 1.5
     num_inference_steps: int = 30
 
+    # Preserve the user-facing canvas before width/height are expanded to the
+    # D32 generation grid. These fields intentionally participate in dynamic
+    # batch compatibility because every item in a decoded tensor batch must use
+    # the same crop.
+    requested_width: int | None = field(default=None, init=False)
+    requested_height: int | None = field(default=None, init=False)
+
     def _adjust(self, server_args):
         requested_width = self.width
         requested_height = self.height
         if self.width is not None and self.height is not None:
+            if self.requested_width is None:
+                self.requested_width = requested_width
+            if self.requested_height is None:
+                self.requested_height = requested_height
             self.width, self.height = align_glm_image_resolution(
                 self.width, self.height
             )

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -65,6 +65,39 @@ except RuntimeError:
     pass
 
 
+def _replace_sampling_params_for_prompt(
+    sampling_params_orig: SamplingParams,
+    prompt: str,
+    output_file_name: str | None,
+    image_path: str | list[str] | None,
+) -> SamplingParams:
+    """Clone per-prompt parameters without losing model-internal state."""
+    sampling_params = dataclasses.replace(
+        sampling_params_orig,
+        prompt=prompt,
+        output_file_name=output_file_name,
+        image_path=image_path,
+    )
+
+    # dataclasses.replace() resets fields declared with init=False. Preserve
+    # model-internal output geometry so GLM-Image can crop the aligned canvas
+    # back to the user's requested size.
+    for field_name in ("requested_width", "requested_height"):
+        if hasattr(sampling_params_orig, field_name):
+            setattr(
+                sampling_params,
+                field_name,
+                getattr(sampling_params_orig, field_name),
+            )
+
+    # dataclasses.replace() also drops non-field attributes. Keep the explicit
+    # user fields so InputValidationStage honors values such as width/height.
+    sampling_params._explicit_fields = getattr(
+        sampling_params_orig, "_explicit_fields", set()
+    ) | {"prompt", "output_file_name", "image_path"}
+    return sampling_params
+
+
 class DiffGenerator:
     """
     A unified class for generating images/videos using diffusion models.
@@ -228,18 +261,12 @@ class DiffGenerator:
         )
 
         for i, p in enumerate(prompts):
-            sampling_params = dataclasses.replace(
+            sampling_params = _replace_sampling_params_for_prompt(
                 sampling_params_orig,
                 prompt=p,
                 output_file_name=user_output_file_name,
                 image_path=image_paths_per_prompt[i],
             )
-            # `dataclasses.replace` drops non-field attrs; restore
-            # `_explicit_fields` so InputValidationStage honors user-supplied
-            # width/height, and mark the keys overridden above as explicit.
-            sampling_params._explicit_fields = getattr(
-                sampling_params_orig, "_explicit_fields", set()
-            ) | {"prompt", "output_file_name", "image_path"}
             sampling_params._set_output_file_name()
             req = prepare_request(
                 server_args=self.server_args,

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -257,13 +257,15 @@ def _get_response_resize(
                 width, height = output_image.size
             return f"{width}x{height}"
         except (OSError, ValueError):
-            # Fall back to the aligned sampling canvas if the output cannot be
+            # Fall back to request metadata if the output cannot be inspected
-            # inspected (for example, for a custom output transport).
+            # (for example, for a custom output transport).
             pass
 
-    if sampling_params.width is None or sampling_params.height is None:
+    width = sampling_params.requested_width or sampling_params.width
+    height = sampling_params.requested_height or sampling_params.height
+    if width is None or height is None:
         return None
-    return sampling_params.output_size_str()
+    return f"{width}x{height}"
 
 
 @router.post("/generations", response_model=ImageResponse)

python/sglang/multimodal_gen/runtime/pipelines/glm_image.py

@@ -1,3 +1,4 @@
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
@@ -6,6 +7,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.glm_image import (
     GlmImageAR,
     GlmImageBeforeDenoisingStage,
+    GlmImageDecodingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 
@@ -50,7 +52,14 @@ class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
             ),
         )
 
-        self.add_standard_decoding_stage()
+        self.add_stage_factory(
+            RoleType.DECODER,
+            lambda: GlmImageDecodingStage(
+                vae=self.get_module("vae"),
+                pipeline=self,
+            ),
+            "decoding_stage",
+        )
 
 
 EntryClass = [GlmImagePipeline]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py

@@ -21,11 +21,12 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im
     ComponentUse,
 )
 from sglang.multimodal_gen.runtime.models.dits.glm_image import GlmImageKVCache
-from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
     PipelineStage,
     StageParallelismType,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.precision import (
@@ -142,6 +143,32 @@ def _validate_glm_image_resolution_alignment(width: int, height: int) -> None:
         )
 
 
+def center_crop_glm_image_output(
+    frames: torch.Tensor,
+    target_width: int | None,
+    target_height: int | None,
+) -> torch.Tensor:
+    """Center-crop decoded GLM-Image pixels back to the requested canvas."""
+    if None in (target_width, target_height):
+        return frames
+
+    decoded_height, decoded_width = frames.shape[-2:]
+    if target_width > decoded_width or target_height > decoded_height:
+        raise ValueError(
+            "Cannot crop GLM-Image output to a canvas larger than the decoded "
+            f"image: requested {target_width}x{target_height}, decoded "
+            f"{decoded_width}x{decoded_height}"
+        )
+    if (target_width, target_height) == (decoded_width, decoded_height):
+        return frames
+
+    left = (decoded_width - target_width) // 2
+    top = (decoded_height - target_height) // 2
+    return frames[
+        ..., top : top + target_height, left : left + target_width
+    ].contiguous()
+
+
 def pooled_image_features_to_tensor(image_features) -> torch.Tensor:
     pooler_output = getattr(image_features, "pooler_output", None)
     if pooler_output is not None:
@@ -681,7 +708,7 @@ class GlmImageAR(PipelineStage):
         width = batch.width
         if batch.image_path is not None:
             ar_condition_images = [
-                resize_glm_image_to_alignment(load_image(img_path))
+                load_image(img_path)
                 for img_path in image_path_to_list(batch.image_path)
             ]
         else:
@@ -693,6 +720,11 @@ class GlmImageAR(PipelineStage):
             height = height or ar_condition_images[0].height
             width = width or ar_condition_images[0].width
 
+        if getattr(batch, "requested_width", None) is None:
+            batch.requested_width = width
+        if getattr(batch, "requested_height", None) is None:
+            batch.requested_height = height
+
         requested_width = width
         requested_height = height
         width, height = align_glm_image_resolution(width, height)
@@ -707,6 +739,11 @@ class GlmImageAR(PipelineStage):
                 height,
             )
 
+        if ar_condition_images is not None:
+            ar_condition_images = [
+                resize_glm_image_to_alignment(image) for image in ar_condition_images
+            ]
+
         time_start = time.time()
         num_outputs = _num_outputs_per_prompt(batch)
         seed = getattr(batch, "seed", None)
@@ -786,6 +823,34 @@ class GlmImageAR(PipelineStage):
         return batch
 
 
+class GlmImageDecodingStage(DecodingStage):
+    """Decode on the D32 canvas, then restore the user-requested dimensions."""
+
+    @torch.no_grad()
+    def forward(
+        self,
+        batch: Req,
+        server_args: ServerArgs,
+    ) -> OutputBatch:
+        output_batch = super().forward(batch, server_args)
+        if output_batch.output is not None:
+            output_batch.output = center_crop_glm_image_output(
+                output_batch.output,
+                batch.requested_width,
+                batch.requested_height,
+            )
+        if output_batch.trajectory_decoded is not None:
+            output_batch.trajectory_decoded = [
+                center_crop_glm_image_output(
+                    decoded,
+                    batch.requested_width,
+                    batch.requested_height,
+                )
+                for decoded in output_batch.trajectory_decoded
+            ]
+        return output_batch
+
+
 class GlmImageBeforeDenoisingStage(PipelineStage):
     r"""
     Pipeline for text-to-image generation using GLM-Image.

python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json

@@ -398,7 +398,7 @@
                 "GlmImageAR": 69033.12,
                 "GlmImageBeforeDenoisingStage": 46.04,
                 "DenoisingStage": 18392.5,
-                "DecodingStage": 153.7
+                "GlmImageDecodingStage": 153.7
             },
             "denoise_step_ms": {
                 "0": 479.04,

python/sglang/multimodal_gen/test/unit/test_glm_image_ar.py

@@ -10,8 +10,11 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.image_api import (
     _build_image_response_kwargs,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.glm_image import (
     GlmImageAR,
+    GlmImageDecodingStage,
+    center_crop_glm_image_output,
 )
 from sglang.multimodal_gen.runtime.server_args import set_global_server_args
 
@@ -232,6 +235,9 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
                 stage.forward(batch, self._server_args())
 
                 self.assertEqual((batch.width, batch.height), expected)
+                self.assertEqual(
+                    (batch.requested_width, batch.requested_height), requested
+                )
                 stage.generate_prior_tokens.assert_called_once_with(
                     prompt="A simple product sketch",
                     image=None,
@@ -270,9 +276,77 @@ class TestGlmImageARSrtBackend(unittest.TestCase):
 
         call_kwargs = stage.generate_prior_tokens.call_args.kwargs
         self.assertEqual((batch.width, batch.height), (1280, 736))
+        self.assertEqual((batch.requested_width, batch.requested_height), (1280, 720))
         self.assertEqual(call_kwargs["image"][0].size, (1280, 736))
         self.assertEqual((call_kwargs["width"], call_kwargs["height"]), (1280, 736))
 
+    @patch(
+        "sglang.multimodal_gen.runtime.pipelines_core.stages."
+        "model_specific_stages.glm_image.get_local_torch_device",
+        return_value=torch.device("cpu"),
+    )
+    @patch(
+        "sglang.multimodal_gen.runtime.pipelines_core.stages."
+        "model_specific_stages.glm_image.load_image",
+        return_value=Image.new("RGB", (1280, 720)),
+    )
+    def test_forward_preserves_implicit_edit_image_size(
+        self, _mock_load_image, _mock_device
+    ):
+        stage = GlmImageAR(processor=_FakeProcessor(), vision_language_encoder=None)
+        stage.generate_prior_tokens = MagicMock(
+            return_value=(torch.zeros((1, 1), dtype=torch.long), None, None)
+        )
+        sampling = GlmImageSamplingParams(
+            prompt="Edit this image",
+            image_path="input.png",
+        )
+        sampling.seed = None
+        batch = Req(sampling_params=sampling)
+
+        stage.forward(batch, self._server_args())
+
+        self.assertEqual((batch.width, batch.height), (1280, 736))
+        self.assertEqual((batch.requested_width, batch.requested_height), (1280, 720))
+        call_kwargs = stage.generate_prior_tokens.call_args.kwargs
+        self.assertEqual(call_kwargs["image"][0].size, (1280, 736))
+        self.assertEqual((call_kwargs["width"], call_kwargs["height"]), (1280, 736))
+
+    def test_center_crop_restores_requested_size(self):
+        frames = torch.arange(1024 * 1024).reshape(1, 1, 1024, 1024)
+
+        cropped = center_crop_glm_image_output(frames, 1000, 999)
+
+        self.assertEqual(tuple(cropped.shape), (1, 1, 999, 1000))
+        self.assertEqual(cropped[0, 0, 0, 0], frames[0, 0, 12, 12])
+        self.assertEqual(cropped[0, 0, -1, -1], frames[0, 0, 1010, 1011])
+        self.assertTrue(cropped.is_contiguous())
+
+    @patch.object(DecodingStage, "forward")
+    def test_decoding_stage_crops_outputs_and_trajectory(self, mock_decode):
+        frames = torch.zeros((2, 3, 736, 1280))
+        trajectory = [
+            torch.zeros((2, 3, 1, 736, 1280)),
+            torch.ones((2, 3, 1, 736, 1280)),
+        ]
+        mock_decode.return_value = OutputBatch(
+            output=frames,
+            trajectory_decoded=trajectory,
+        )
+        stage = GlmImageDecodingStage(vae=None)
+        sampling = GlmImageSamplingParams(width=1280, height=736)
+        sampling.requested_width = 1280
+        sampling.requested_height = 720
+        batch = Req(sampling_params=sampling)
+
+        output_batch = stage.forward(batch, self._server_args())
+
+        self.assertEqual(tuple(output_batch.output.shape), (2, 3, 720, 1280))
+        self.assertEqual(len(output_batch.trajectory_decoded), 2)
+        for decoded in output_batch.trajectory_decoded:
+            self.assertEqual(tuple(decoded.shape), (2, 3, 1, 720, 1280))
+        mock_decode.assert_called_once_with(batch, self._server_args())
+
     @patch(
         "sglang.multimodal_gen.runtime.pipelines_core.stages."
         "model_specific_stages.glm_image.get_local_torch_device",

python/sglang/multimodal_gen/test/unit/test_openai_image_api.py

@@ -72,6 +72,14 @@ def test_response_resize_is_only_populated_for_glm_image():
     assert _get_response_resize(SamplingParams(width=1280, height=736)) is None
 
 
+def test_response_resize_prefers_requested_size_over_generation_canvas():
+    glm_sampling = GlmImageSamplingParams(width=1280, height=736)
+    glm_sampling.requested_width = 1280
+    glm_sampling.requested_height = 720
+
+    assert _get_response_resize(glm_sampling) == "1280x720"
+
+
 def test_response_resize_uses_actual_generated_image_size(tmp_path):
     output_path = tmp_path / "output.png"
     Image.new("RGB", (1280, 736)).save(output_path)

python/sglang/multimodal_gen/test/unit/test_sampling_params.py

@@ -158,6 +158,9 @@ class TestSamplingParamsSubclass(unittest.TestCase):
                     params._adjust(server_args)
 
                 self.assertEqual((params.width, params.height), expected)
+                self.assertEqual(
+                    (params.requested_width, params.requested_height), requested
+                )
                 mock_warning.assert_called_once_with(
                     "GLM-Image requires dimensions divisible by %s; adjusted "
                     "requested resolution from %sx%s to %sx%s",
@@ -459,8 +462,10 @@ class TestSamplingParamsCliArgs(unittest.TestCase):
         )
 
     def test_dataclasses_replace_preserves_explicit_fields(self):
-        """`dataclasses.replace` drops `_explicit_fields`; DiffGenerator must restore it."""
+        """Per-prompt clones retain explicit and model-internal fields."""
-        import dataclasses
+        from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import (
+            _replace_sampling_params_for_prompt,
+        )
 
         server_args = MagicMock()
         server_args.backend = "sglang"
@@ -484,18 +489,12 @@ class TestSamplingParamsCliArgs(unittest.TestCase):
         self.assertIn("width", sampling_params_orig._explicit_fields)
         self.assertIn("height", sampling_params_orig._explicit_fields)
 
-        cloned = dataclasses.replace(
+        cloned = _replace_sampling_params_for_prompt(
             sampling_params_orig,
             prompt="new",
             output_file_name=None,
             image_path="/tmp/in2.png",
         )
-        self.assertFalse(hasattr(cloned, "_explicit_fields"))
-
-        # Mirror the restore done in DiffGenerator.generate().
-        cloned._explicit_fields = getattr(
-            sampling_params_orig, "_explicit_fields", set()
-        ) | {"prompt", "output_file_name", "image_path"}
 
         explicit = set(cloned.build_request_extra()["explicit_fields"])
         self.assertIn("width", explicit)
@@ -503,6 +502,29 @@ class TestSamplingParamsCliArgs(unittest.TestCase):
         self.assertIn("prompt", explicit)
         self.assertIn("image_path", explicit)
 
+    def test_per_prompt_clone_preserves_glm_image_crop_size(self):
+        from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import (
+            _replace_sampling_params_for_prompt,
+        )
+
+        sampling_params_orig = GlmImageSamplingParams(
+            prompt="orig",
+            width=1024,
+            height=1024,
+        )
+        sampling_params_orig.requested_width = 1000
+        sampling_params_orig.requested_height = 999
+
+        cloned = _replace_sampling_params_for_prompt(
+            sampling_params_orig,
+            prompt="new",
+            output_file_name=None,
+            image_path=None,
+        )
+
+        self.assertEqual((cloned.width, cloned.height), (1024, 1024))
+        self.assertEqual((cloned.requested_width, cloned.requested_height), (1000, 999))
+
 
 if __name__ == "__main__":
     unittest.main()