From 7ce8faae28eeb778c75130aeca5623504a7dc075 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 1 Dec 2025 16:35:34 +0800 Subject: [PATCH] [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> --- .../configs/pipeline_configs/base.py | 50 +++++++- .../configs/pipeline_configs/flux.py | 25 +++- .../configs/pipeline_configs/qwen_image.py | 38 +++++++ .../runtime/loader/component_loader.py | 1 - .../pipelines_core/stages/image_encoding.py | 107 ++---------------- .../stages/timestep_preparation.py | 19 +--- 6 files changed, 120 insertions(+), 120 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 80de61c49..f0c9c85d3 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/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] diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index 825e64d00..402cbf686 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index 18e7d46ad..f6f589a7e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index d0d81c552..3e93e80a5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index aff633b65..1ecbd11b5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/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 ( - isinstance(server_args.pipeline_config, QwenImageEditPipelineConfig) - 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_processor_kwargs = ( + server_args.pipeline_config.prepare_image_processor_kwargs(batch) + ) 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 - - # 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 + batch.image_latent = server_args.pipeline_config.postprocess_image_latent( + latent_condition, batch + ) self.maybe_free_model_hooks() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py index 91bc6e695..ca0e65f61 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py +++ b/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 = ( - 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 - ) + sigmas = server_args.pipeline_config.prepare_sigmas(sigmas, num_inference_steps) # Prepare extra kwargs for set_timesteps extra_set_timesteps_kwargs = {}