[diffusion] refactor: remove hard-code of instanceof on PipelineConfig (#14186)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

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

@@ -7,6 +7,7 @@ from dataclasses import asdict, dataclass, field, fields
 from enum import Enum, auto
 from typing import Any
 
+import numpy as np
 import PIL
 import torch
 from einops import rearrange
@@ -183,9 +184,48 @@ class PipelineConfig:
         height, width = get_default_height_width(image, vae_scale_factor, height, width)
         return width, height
 
+    ## For timestep preparation stage
+
+    def prepare_sigmas(self, sigmas, num_inference_steps):
+        return sigmas
+
+    ## For ImageVAEEncodingStage
     def resize_condition_image(self, image, target_width, target_height):
         return image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS)
 
+    def prepare_image_processor_kwargs(self, batch):
+        return {}
+
+    def postprocess_image_latent(self, latent_condition, batch):
+        vae_arch_config = self.vae_config.arch_config
+        spatial_compression_ratio = vae_arch_config.spatial_compression_ratio
+        temporal_compression_ratio = vae_arch_config.temporal_compression_ratio
+        num_frames = batch.num_frames
+        latent_height = batch.height // spatial_compression_ratio
+        latent_width = batch.width // spatial_compression_ratio
+        mask_lat_size = torch.ones(1, 1, num_frames, latent_height, latent_width)
+        mask_lat_size[:, :, 1:] = 0
+        first_frame_mask = mask_lat_size[:, :, 0:1]
+        first_frame_mask = torch.repeat_interleave(
+            first_frame_mask,
+            repeats=temporal_compression_ratio,
+            dim=2,
+        )
+        mask_lat_size = torch.concat(
+            [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2
+        )
+        mask_lat_size = mask_lat_size.view(
+            1,
+            -1,
+            temporal_compression_ratio,
+            latent_height,
+            latent_width,
+        )
+        mask_lat_size = mask_lat_size.transpose(1, 2)
+        mask_lat_size = mask_lat_size.to(latent_condition.device)
+        image_latents = torch.concat([mask_lat_size, latent_condition], dim=1)
+        return image_latents
+
     def slice_noise_pred(self, noise, latents):
         return noise
 
@@ -230,7 +270,7 @@ class PipelineConfig:
         return None
 
     # called after vae encode
-    def post_process_vae_encode(self, image_latents, vae):
+    def postprocess_vae_encode(self, image_latents, vae):
         return image_latents
 
     # called after scale_and_shift, before vae decoding
@@ -553,6 +593,14 @@ class PipelineConfig:
 class ImagePipelineConfig(PipelineConfig):
     """Base config for image generation pipelines with token-like latents [B, S, D]."""
 
+    def _prepare_sigmas(self, sigmas, num_inference_steps):
+        sigmas = (
+            np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+            if sigmas is None
+            else sigmas
+        )
+        return sigmas
+
     def shard_latents_for_sp(self, batch, latents):
         sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank()
         seq_len = latents.shape[1]

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

@@ -29,6 +29,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import (
     clip_preprocess_text,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import _pack_latents
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 
 
 def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
@@ -81,6 +82,9 @@ class FluxPipelineConfig(ImagePipelineConfig):
         ]
     )
 
+    def prepare_sigmas(self, sigmas, num_inference_steps):
+        return self._prepare_sigmas(sigmas, num_inference_steps)
+
     def prepare_latent_shape(self, batch, batch_size, num_frames):
         height = 2 * (
             batch.height // (self.vae_config.arch_config.vae_scale_factor * 2)
@@ -464,6 +468,25 @@ class Flux2PipelineConfig(FluxPipelineConfig):
     def resize_condition_image(self, image, target_width, target_height):
         return image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS)
 
+    def postprocess_image_latent(self, latent_condition, batch):
+        batch_size = batch.batch_size
+        # get image_latent_ids right after scale & shift
+        image_latent_ids = _prepare_image_ids([latent_condition])
+        image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1)
+        image_latent_ids = image_latent_ids.to(get_local_torch_device())
+        batch.condition_image_latent_ids = image_latent_ids
+
+        # latent: (1, 128, 32, 32)
+        packed = self.maybe_pack_latents(
+            latent_condition, None, batch
+        )  # (1, 1024, 128)
+        packed = packed.squeeze(0)  # (1024, 128) - remove batch dim
+
+        # Concatenate all reference tokens along sequence dimension
+        image_latents = packed.unsqueeze(0)  # (1, N*1024, 128)
+        image_latents = image_latents.repeat(batch_size, 1, 1)
+        return image_latents
+
     def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
 
         txt_ids = _prepare_text_ids(prompt_embeds).to(device=device)
@@ -510,7 +533,7 @@ class Flux2PipelineConfig(FluxPipelineConfig):
     def maybe_prepare_latent_ids(self, latents):
         return _prepare_latent_ids(latents)
 
-    def post_process_vae_encode(self, image_latents, vae):
+    def postprocess_vae_encode(self, image_latents, vae):
         # patchify
         image_latents = _patchify_latents(image_latents)
         return image_latents

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

@@ -104,6 +104,17 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
         ]
     )
 
+    def prepare_sigmas(self, sigmas, num_inference_steps):
+        return self._prepare_sigmas(sigmas, num_inference_steps)
+
+    def prepare_image_processor_kwargs(self, batch):
+        if batch.prompt:
+            prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
+            txt = prompt_template_encode.format(batch.prompt)
+            return dict(text=[txt], padding=True)
+        else:
+            return {}
+
     def get_vae_scale_factor(self):
         return self.vae_config.arch_config.vae_scale_factor
 
@@ -265,6 +276,33 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
     def resize_condition_image(self, image, target_width, target_height):
         return resize(image, target_height, target_width, resize_mode="default")
 
+    def postprocess_image_latent(self, latent_condition, batch):
+        batch_size = batch.batch_size
+        if batch_size > latent_condition.shape[0]:
+            if batch_size % latent_condition.shape[0] == 0:
+                # expand init_latents for batch_size
+                additional_image_per_prompt = batch_size // latent_condition.shape[0]
+                image_latents = latent_condition.repeat(
+                    additional_image_per_prompt, 1, 1, 1
+                )
+            else:
+                raise ValueError(
+                    f"Cannot duplicate `image` of batch size {latent_condition.shape[0]} to {batch_size} text prompts."
+                )
+        else:
+            image_latents = latent_condition
+        image_latent_height, image_latent_width = image_latents.shape[3:]
+        num_channels_latents = self.dit_config.arch_config.in_channels // 4
+        image_latents = _pack_latents(
+            image_latents,
+            batch_size,
+            num_channels_latents,
+            image_latent_height,
+            image_latent_width,
+        )
+
+        return image_latents
+
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
         return self._prepare_edit_cond_kwargs(
             batch, batch.prompt_embeds, rotary_emb, device, dtype

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -562,7 +562,6 @@ class VAELoader(ComponentLoader):
 
         server_args.model_paths["vae"] = component_model_path
 
-        # TODO: abstract these logics
         logger.info("HF model config: %s", config)
         vae_config = server_args.pipeline_config.vae_config
         vae_config.update_model_arch(config)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -12,14 +12,7 @@ import torch
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 
 from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
-from sglang.multimodal_gen.configs.pipeline_configs.flux import (
-    Flux2PipelineConfig,
-    _prepare_image_ids,
-)
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
-    QwenImageEditPipelineConfig,
-    QwenImagePipelineConfig,
-    _pack_latents,
     qwen_image_postprocess_text,
 )
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
@@ -111,15 +104,9 @@ class ImageEncodingStage(PipelineStage):
 
         image = batch.condition_image
 
-        if batch.prompt and (
+        image_processor_kwargs = (
-            isinstance(server_args.pipeline_config, QwenImageEditPipelineConfig)
+            server_args.pipeline_config.prepare_image_processor_kwargs(batch)
-            or isinstance(server_args.pipeline_config, QwenImagePipelineConfig)
+        )
-        ):
-            prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
-            txt = prompt_template_encode.format(batch.prompt)
-            image_processor_kwargs = dict(text=[txt], padding=True)
-        else:
-            image_processor_kwargs = {}
 
         image_inputs = self.image_processor(
             images=image, return_tensors="pt", **image_processor_kwargs
@@ -295,12 +282,12 @@ class ImageVAEEncodingStage(PipelineStage):
             encoder_output, generator, sample_mode=sample_mode
         )
 
-        latent_condition = self.server_args.pipeline_config.post_process_vae_encode(
+        latent_condition = server_args.pipeline_config.postprocess_vae_encode(
             latent_condition, self.vae
         )
 
         scaling_factor, shift_factor = (
-            self.server_args.pipeline_config.get_decode_scale_and_shift(
+            server_args.pipeline_config.get_decode_scale_and_shift(
                 device=latent_condition.device,
                 dtype=latent_condition.dtype,
                 vae=self.vae,
@@ -317,87 +304,9 @@ class ImageVAEEncodingStage(PipelineStage):
         latent_condition -= shift_factor
         latent_condition = latent_condition * scaling_factor
 
-        batch_size = batch.batch_size
+        batch.image_latent = server_args.pipeline_config.postprocess_image_latent(
-
+            latent_condition, batch
-        # TODO: abstract this
+        )
-        if isinstance(server_args.pipeline_config, QwenImageEditPipelineConfig):
-            if (
-                batch_size > latent_condition.shape[0]
-                and batch_size % latent_condition.shape[0] == 0
-            ):
-                # expand init_latents for batch_size
-                additional_image_per_prompt = batch_size // latent_condition.shape[0]
-                image_latents = torch.cat(
-                    [latent_condition] * additional_image_per_prompt, dim=0
-                )
-            elif (
-                batch_size > latent_condition.shape[0]
-                and batch_size % latent_condition.shape[0] != 0
-            ):
-                raise ValueError(
-                    f"Cannot duplicate `image` of batch size {latent_condition.shape[0]} to {batch_size} text prompts."
-                )
-            else:
-                image_latents = torch.cat([latent_condition], dim=0)
-            image_latent_height, image_latent_width = image_latents.shape[3:]
-            num_channels_latents = (
-                self.server_args.pipeline_config.dit_config.arch_config.in_channels // 4
-            )
-            image_latents = _pack_latents(
-                image_latents,
-                batch_size,
-                num_channels_latents,
-                image_latent_height,
-                image_latent_width,
-            )
-        elif isinstance(server_args.pipeline_config, Flux2PipelineConfig):
-            # Pack each latent and concatenate
-            image_latents = [latent_condition]
-            # get image_latent_ids right after scale & shift
-            image_latent_ids = _prepare_image_ids(image_latents)
-            image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1)
-            image_latent_ids = image_latent_ids.to(get_local_torch_device())
-            batch.condition_image_latent_ids = image_latent_ids
-
-            packed_latents = []
-            for latent in image_latents:
-                # latent: (1, 128, 32, 32)
-                packed = server_args.pipeline_config.maybe_pack_latents(
-                    latent, None, None
-                )  # (1, 1024, 128)
-                packed = packed.squeeze(0)  # (1024, 128) - remove batch dim
-                packed_latents.append(packed)
-
-            # Concatenate all reference tokens along sequence dimension
-            image_latents = torch.cat(packed_latents, dim=0)  # (N*1024, 128)
-            image_latents = image_latents.unsqueeze(0)  # (1, N*1024, 128)
-            image_latents = image_latents.repeat(batch_size, 1, 1)
-        else:
-            latent_height = batch.height // self.vae.spatial_compression_ratio
-            latent_width = batch.width // self.vae.spatial_compression_ratio
-            mask_lat_size = torch.ones(1, 1, num_frames, latent_height, latent_width)
-            mask_lat_size[:, :, list(range(1, num_frames))] = 0
-            first_frame_mask = mask_lat_size[:, :, 0:1]
-            first_frame_mask = torch.repeat_interleave(
-                first_frame_mask,
-                repeats=self.vae.temporal_compression_ratio,
-                dim=2,
-            )
-            mask_lat_size = torch.concat(
-                [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2
-            )
-            mask_lat_size = mask_lat_size.view(
-                1,
-                -1,
-                self.vae.temporal_compression_ratio,
-                latent_height,
-                latent_width,
-            )
-            mask_lat_size = mask_lat_size.transpose(1, 2)
-            mask_lat_size = mask_lat_size.to(latent_condition.device)
-            image_latents = torch.concat([mask_lat_size, latent_condition], dim=1)
-
-        batch.image_latent = image_latents
 
         self.maybe_free_model_hooks()
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py

@@ -10,13 +10,6 @@ This module contains implementations of timestep preparation stages for diffusio
 import inspect
 from typing import Any, Callable, Tuple
 
-import numpy as np
-
-from sglang.multimodal_gen.configs.pipeline_configs import FluxPipelineConfig
-from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
-    QwenImageEditPipelineConfig,
-    QwenImagePipelineConfig,
-)
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
@@ -79,17 +72,7 @@ class TimestepPreparationStage(PipelineStage):
         sigmas = batch.sigmas
         n_tokens = batch.n_tokens
 
-        is_flux = (
+        sigmas = server_args.pipeline_config.prepare_sigmas(sigmas, num_inference_steps)
-            isinstance(server_args.pipeline_config, FluxPipelineConfig)
-            or isinstance(server_args.pipeline_config, QwenImagePipelineConfig)
-            or isinstance(server_args.pipeline_config, QwenImageEditPipelineConfig)
-        )
-        if is_flux:
-            sigmas = (
-                np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
-                if sigmas is None
-                else sigmas
-            )
 
         # Prepare extra kwargs for set_timesteps
         extra_set_timesteps_kwargs = {}