diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index 8f594b894..cd593b340 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/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] diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py index c7c94df4d..eaf37534e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/glm_image.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index bbb54b178..8ed772d70 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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 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 c909619a6..fb24cca80 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 @@ -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,13 +235,29 @@ class GlmImageAR(PipelineStage): prior_token_image_ids = None if image is not None: - prior_token_image_embed = self.vision_language_encoder.get_image_features( - inputs["pixel_values"], image_grid_thw[:-1] + source_grids = 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_embed, image_grid_thw[:-1] + prior_token_image_ids_d32 = self.vision_language_encoder.get_image_tokens( + 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. # max_new_tokens must be exactly grid_h * grid_w + 1 (the +1 is for EOS). @@ -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,25 +819,32 @@ 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 set_forward_context(current_timestep=1, attn_metadata=None): - _ = self.transformer( - hidden_states=condition_latent, - encoder_hidden_states=torch.zeros_like(prompt_embeds)[ - :1, :0, ... - ], - prior_token_id=condition_image_prior_token_id, - prior_token_drop=torch.full_like( - condition_image_prior_token_id, False, dtype=torch.bool - ), - timestep=torch.zeros((1,), device=device), - target_size=torch.tensor( - [condition_image.shape[-2:]], device=device - ), - crop_coords=torch.zeros((1, 2), device=device), - attention_kwargs=attention_kwargs, - kv_caches=kv_caches, - kv_caches_mode="write", - ) + 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): + _ = transformer( + hidden_states=condition_latent, + encoder_hidden_states=torch.zeros_like(prompt_embeds)[ + :1, :0, ... + ], + prior_token_id=condition_image_prior_token_id, + prior_token_drop=torch.full_like( + condition_image_prior_token_id, False, dtype=torch.bool + ), + timestep=torch.zeros((1,), device=device), + target_size=torch.tensor( + [condition_image.shape[-2:]], device=device + ), + crop_coords=torch.zeros((1, 2), device=device), + attention_kwargs=attention_kwargs, + kv_caches=kv_caches, + kv_caches_mode="write", + ) # 6. Prepare additional timestep conditions target_size = (height, width)