[diffusion] fix: fix GLM-Image /v1/images/edits support (#25697)

python/sglang/jit_kernel/diffusion/triton/scale_shift.py

@@ -343,6 +343,8 @@ def fuse_scale_shift_kernel(
 
     B, L, C = x.shape
     output = torch.empty_like(x)
+    if x.numel() == 0:
+        return output
 
     if scale.dim() == 4:
         # scale/shift: [B, F, 1, C]

python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py

@@ -22,7 +22,7 @@ class GlmImagePipelineConfig(SpatialImagePipelineConfig):
     vae_precision: str = "bf16"
 
     should_use_guidance: bool = False
-    task_type: ModelTaskType = ModelTaskType.T2I
+    task_type: ModelTaskType = ModelTaskType.TI2I
 
     vae_tiling: bool = False
 

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -445,6 +445,9 @@ class _ScaleResidualNormScaleShift(CustomOp):
         shift: torch.Tensor,
         scale: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual.numel() == 0 or x.numel() == 0:
+            return self.forward_native(residual, x, gate, shift, scale)
+
         if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192:
             import warnings
 

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

@@ -12,6 +12,9 @@ from diffusers.utils.torch_utils import randn_tensor
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
+    ComponentUse,
+)
 from sglang.multimodal_gen.runtime.models.dits.glm_image import GlmImageKVCache
 from sglang.multimodal_gen.runtime.models.vision_utils import load_image
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
@@ -106,6 +109,19 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
+def image_path_to_list(image_path: Union[str, List[str]]) -> List[str]:
+    return image_path if isinstance(image_path, list) else [image_path]
+
+
+def pooled_image_features_to_tensor(image_features) -> torch.Tensor:
+    pooler_output = getattr(image_features, "pooler_output", None)
+    if pooler_output is not None:
+        image_features = pooler_output
+    if isinstance(image_features, torch.Tensor):
+        return image_features
+    return torch.cat(tuple(image_features), dim=0)
+
+
 class GlmImageAR(PipelineStage):
     r"""
     Pipeline for text-to-image generation using GLM-Image.
@@ -219,12 +235,28 @@ class GlmImageAR(PipelineStage):
 
         prior_token_image_ids = None
         if image is not None:
-            prior_token_image_embed = self.vision_language_encoder.get_image_features(
+            source_grids = image_grid_thw[:-1]
-                inputs["pixel_values"], image_grid_thw[:-1]
+            prior_token_image_embed = pooled_image_features_to_tensor(
+                self.vision_language_encoder.get_image_features(
+                    inputs["pixel_values"], source_grids
                 )
-            prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0)
+            )
-            prior_token_image_ids = self.vision_language_encoder.get_image_tokens(
+            prior_token_image_ids_d32 = self.vision_language_encoder.get_image_tokens(
-                prior_token_image_embed, image_grid_thw[:-1]
+                prior_token_image_embed, source_grids
+            )
+            prior_token_image_ids = []
+            prior_ids_per_source = torch.split(
+                prior_token_image_ids_d32,
+                source_grids.prod(dim=-1).tolist(),
+            )
+            for prior_ids, source_grid in zip(prior_ids_per_source, source_grids):
+                _, source_h, source_w = source_grid.tolist()
+                prior_token_image_ids.append(
+                    self._upsample_token_ids(
+                        prior_ids,
+                        int(source_h),
+                        int(source_w),
+                    ).squeeze(0)
                 )
 
         # For GLM-Image, greedy decoding is not allowed; it may cause repetitive outputs.
@@ -257,12 +289,24 @@ class GlmImageAR(PipelineStage):
         prompt = batch.prompt
         height = batch.height
         width = batch.width
+        if batch.image_path is not None:
+            ar_condition_images = [
+                load_image(img_path)
+                for img_path in image_path_to_list(batch.image_path)
+            ]
+        else:
+            ar_condition_images = None
 
         device = get_local_torch_device()
 
+        if ar_condition_images is not None:
+            height = height or ar_condition_images[0].height
+            width = width or ar_condition_images[0].width
+
         time_start = time.time()
         prior_token_id, prior_token_image_ids = self.generate_prior_tokens(
             prompt=prompt,
+            image=ar_condition_images,
             height=height,
             width=width,
         )
@@ -272,6 +316,8 @@ class GlmImageAR(PipelineStage):
 
         batch.prior_token_id = prior_token_id
         batch.prior_token_image_ids = prior_token_image_ids
+        batch.height = height
+        batch.width = width
 
         return batch
 
@@ -341,6 +387,22 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
             else 128
         )
 
+    def component_uses(
+        self, server_args: ServerArgs, stage_name: str | None = None
+    ) -> list[ComponentUse]:
+        stage_name = self._component_stage_name(stage_name)
+        uses: list[ComponentUse] = []
+        if self.transformer is not None:
+            uses.append(
+                ComponentUse(
+                    stage_name=stage_name,
+                    component_name="transformer",
+                    phase="reference_image",
+                    memory_intensive=True,
+                )
+            )
+        return uses
+
     def _parse_and_expand_shape_info(
         self, prompt: str
     ) -> Tuple[str, int, int, int, int]:
@@ -656,7 +718,8 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
         num_inference_steps = batch.num_inference_steps
         if batch.image_path is not None:
             ar_condition_images = [
-                load_image(img_path) for img_path in batch.image_path
+                load_image(img_path)
+                for img_path in image_path_to_list(batch.image_path)
             ]
         else:
             ar_condition_images = None
@@ -676,8 +739,6 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
         self._current_timestep = None
         self._interrupt = False
 
-        batch_size = 1
-
         device = get_local_torch_device()
 
         if ar_condition_images is not None:
@@ -758,8 +819,15 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
                 # Do not remove.
                 # It would be use to run the reference image through a
                 # forward pass at timestep 0 and keep the KV cache.
+                with self.use_declared_component(
+                    component_name="transformer",
+                    module=self.transformer,
+                    phase="reference_image",
+                ) as transformer:
+                    assert transformer is not None
+                    self.transformer = transformer
                     with set_forward_context(current_timestep=1, attn_metadata=None):
-                    _ = self.transformer(
+                        _ = transformer(
                             hidden_states=condition_latent,
                             encoder_hidden_states=torch.zeros_like(prompt_embeds)[
                                 :1, :0, ...