From 8f096b853ad91b32c4f7eef2abf50db7d3ac06de Mon Sep 17 00:00:00 2001 From: AuFlow <73925903+AuFlow@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:50:47 +0800 Subject: [PATCH] [diffusion] fix: crop GLM-Image output to requested size (#33859) Co-authored-by: AuFlow Co-authored-by: ronnie_zheng --- .../multimodal_gen/configs/sample/glmimage.py | 13 +++- .../entrypoints/diffusion_generator.py | 41 ++++++++-- .../runtime/entrypoints/openai/image_api.py | 10 ++- .../runtime/pipelines/glm_image.py | 11 ++- .../stages/model_specific_stages/glm_image.py | 69 ++++++++++++++++- .../server/ascend/perf_baselines_npu.json | 2 +- .../test/unit/test_glm_image_ar.py | 74 +++++++++++++++++++ .../test/unit/test_openai_image_api.py | 8 ++ .../test/unit/test_sampling_params.py | 40 +++++++--- 9 files changed, 243 insertions(+), 25 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/sample/glmimage.py b/python/sglang/multimodal_gen/configs/sample/glmimage.py index 7824d8858..1dad65548 100644 --- a/python/sglang/multimodal_gen/configs/sample/glmimage.py +++ b/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 ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index e026ae895..0523f5650 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index ee0881813..f4db1d78c 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/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 - # inspected (for example, for a custom output transport). + # Fall back to request metadata if the output cannot be inspected + # (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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py b/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py index 90e32015b..4b72b5a84 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py +++ b/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] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py index 72645883b..3faaa54da 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py +++ b/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. diff --git a/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json b/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json index d9fd7146e..56074ea12 100644 --- a/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json +++ b/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, diff --git a/python/sglang/multimodal_gen/test/unit/test_glm_image_ar.py b/python/sglang/multimodal_gen/test/unit/test_glm_image_ar.py index 2cf9dc980..9f11c8e78 100644 --- a/python/sglang/multimodal_gen/test/unit/test_glm_image_ar.py +++ b/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", diff --git a/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py b/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py index 112fe2b67..e4498fbeb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py b/python/sglang/multimodal_gen/test/unit/test_sampling_params.py index 58f3cf19e..9b322079e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sampling_params.py +++ b/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.""" - import dataclasses + """Per-prompt clones retain explicit and model-internal fields.""" + 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()