diff --git a/python/sglang/multimodal_gen/configs/models/vaes/base.py b/python/sglang/multimodal_gen/configs/models/vaes/base.py index e7a078b6e..18f12003c 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/base.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/base.py @@ -9,7 +9,6 @@ from typing import Any import torch from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig -from sglang.multimodal_gen.runtime.models.vision_utils import get_default_height_width from sglang.multimodal_gen.utils import StoreBoolean @@ -51,13 +50,6 @@ class VAEConfig(ModelConfig): def post_init(self): pass - # returns width, height - def calculate_dimensions( - self, image, vae_scale_factor, width, height - ) -> tuple[int, int]: - height, width = get_default_height_width(image, vae_scale_factor, height, width) - return width, height - @staticmethod def add_cli_args(parser: Any, prefix: str = "vae-config") -> Any: """Add CLI arguments for VAEConfig fields""" @@ -148,6 +140,9 @@ class VAEConfig(ModelConfig): return parser + def get_vae_scale_factor(self): + return 2 ** (len(self.arch_config.block_out_channels) - 1) + @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig": kwargs = {} diff --git a/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py b/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py index 1ba1a2098..d08ab5da9 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py @@ -4,7 +4,6 @@ from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig -from sglang.multimodal_gen.utils import calculate_dimensions @dataclass @@ -28,9 +27,6 @@ class QwenImageVAEArchConfig(VAEArchConfig): scale_factor_spatial: int = 8 clip_output: bool = True - def __post_init__(self): - self.vae_scale_factor = 2 ** len(self.temperal_downsample) - @dataclass class QwenImageVAEConfig(VAEConfig): @@ -42,11 +38,8 @@ class QwenImageVAEConfig(VAEConfig): use_temporal_tiling: bool = False use_parallel_tiling: bool = False - def calculate_dimensions(self, image, vae_scale_factor, width, height): - width = image.size[0] - height = image.size[1] - width, height, _ = calculate_dimensions(1024 * 1024, width / height) - return width, height + def get_vae_scale_factor(self): + return 2 ** len(self.arch_config.temperal_downsample) def __post_init__(self): self.blend_num_frames = ( diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index e7c2a8be3..80de61c49 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -7,8 +7,8 @@ from dataclasses import asdict, dataclass, field, fields from enum import Enum, auto from typing import Any +import PIL import torch -from diffusers.image_processor import VaeImageProcessor from einops import rearrange from sglang.multimodal_gen.configs.models import ( @@ -24,6 +24,7 @@ from sglang.multimodal_gen.runtime.distributed import ( get_sp_world_size, sequence_model_parallel_all_gather, ) +from sglang.multimodal_gen.runtime.models.vision_utils import get_default_height_width from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import ( FlexibleArgumentParser, @@ -174,15 +175,20 @@ class PipelineConfig: # Compilation # enable_torch_compile: bool = False + # calculate the adjust size for condition image + # width: original condition image width + # height: original condition image height + def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]: + vae_scale_factor = self.vae_config.arch_config.spatial_compression_ratio + height, width = get_default_height_width(image, vae_scale_factor, height, width) + return width, height + + def resize_condition_image(self, image, target_width, target_height): + return image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS) + def slice_noise_pred(self, noise, latents): return noise - def maybe_resize_condition_image(self, width, height, image): - """ - image: input image - """ - return image, width, height - def adjust_num_frames(self, num_frames): return num_frames @@ -190,10 +196,6 @@ class PipelineConfig: def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict: return tokenizer(prompt, **tok_kwargs) - # called in ImageEncodingStage, preprocess the image - def preprocess_image(self, image, image_processor: VaeImageProcessor): - return image - def prepare_latent_shape(self, batch, batch_size, num_frames): height = batch.height // self.vae_config.arch_config.spatial_compression_ratio width = batch.width // self.vae_config.arch_config.spatial_compression_ratio diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index 94e5d7196..b8a2601c3 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -2,7 +2,7 @@ import math from dataclasses import dataclass, field from typing import Callable, List, Optional -import PIL.Image +import PIL import torch from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig @@ -448,18 +448,21 @@ class Flux2PipelineConfig(FluxPipelineConfig): def get_neg_prompt_embeds(self, batch): return batch.negative_prompt_embeds[0] - def maybe_resize_condition_image(self, width, height, image): + def calculate_condition_image_size( + self, image, width, height + ) -> Optional[tuple[int, int]]: target_area: int = 1024 * 1024 - if width is not None and height is not None: if width * height > target_area: scale = math.sqrt(target_area / (width * height)) width = int(width * scale) height = int(height * scale) - image = image.resize((width, height), PIL.Image.Resampling.LANCZOS) - width, height = image.size + return width, height - return image, width, height + return None + + def resize_condition_image(self, image, target_width, target_height): + return image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS) def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch): 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 fb39712a9..18e7d46ad 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -14,6 +14,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, shard_rotary_emb_for_sp, ) +from sglang.multimodal_gen.runtime.models.vision_utils import resize from sglang.multimodal_gen.utils import calculate_dimensions @@ -215,8 +216,7 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig): assert batch_size == 1 height = batch.height width = batch.width - image = batch.condition_image - image_size = image[0].size if isinstance(image, list) else image.size + image_size = batch.original_condition_image_size edit_width, edit_height, _ = calculate_dimensions( 1024 * 1024, image_size[0] / image_size[1] ) @@ -262,6 +262,9 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig): "freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)), } + def resize_condition_image(self, image, target_width, target_height): + return resize(image, target_height, target_width, resize_mode="default") + 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 @@ -272,26 +275,11 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig): batch, batch.negative_prompt_embeds, rotary_emb, device, dtype ) - def preprocess_image(self, image, image_processor): - image_size = image[0].size if isinstance(image, list) else image.size + def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]: calculated_width, calculated_height, _ = calculate_dimensions( - 1024 * 1024, image_size[0] / image_size[1] + 1024 * 1024, width / height ) - image = image_processor.resize(image, calculated_height, calculated_width) - return image - - def maybe_resize_condition_image(self, width, height, image): - image_size = image[0].size if isinstance(image, list) else image.size - calculated_width, calculated_height, _ = calculate_dimensions( - 1024 * 1024, image_size[0] / image_size[1] - ) - height = height or calculated_height - width = width or calculated_width - - multiple_of = self.get_vae_scale_factor() * 2 - width = width // multiple_of * multiple_of - height = height // multiple_of * multiple_of - return width, height + return calculated_width, calculated_height def slice_noise_pred(self, noise, latents): # remove noise over input image diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index ab408e95e..20ed203a7 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -288,7 +288,7 @@ class DiffGenerator: # TODO: send batch when supported for request_idx, req in enumerate(requests): logger.info( - "Processing prompt: %d/%d: %s", + "Processing prompt %d/%d: %s", request_idx + 1, len(requests), req.prompt[:100], diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 80d90c114..1f650676e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -55,6 +55,8 @@ class Req: image_path: str | None = None # Image encoder hidden states image_embeds: list[torch.Tensor] = field(default_factory=list) + + original_condition_image_size: tuple[int, int] = None condition_image: torch.Tensor | PIL.Image.Image | None = None pixel_values: torch.Tensor | PIL.Image.Image | None = None preprocessed_image: torch.Tensor | None = None @@ -216,14 +218,6 @@ class Req: self.guidance_scale_2 = self.guidance_scale def adjust_size(self, server_args: ServerArgs): - if self.height is None or self.width is None: - _image, width, height = ( - server_args.pipeline_config.maybe_resize_condition_image( - self.width, self.height, self.condition_image - ) - ) - self.width = width - self.height = height if self.height is None or self.width is None: self.width = 1280 self.height = 720 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 6fef33510..e79ea187b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py @@ -92,6 +92,9 @@ class PipelineStage(ABC): # Default implementation - no verification return VerificationResult() + def maybe_free_model_hooks(self): + pass + # execute on all ranks by default @property def parallelism_type(self) -> StageParallelismType: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index 9a9e5fe0d..2d9cf7f07 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -199,8 +199,7 @@ class DecodingStage(PipelineStage): ) # Offload models if needed - if hasattr(self, "maybe_free_model_hooks"): - self.maybe_free_model_hooks() + self.maybe_free_model_hooks() if server_args.vae_cpu_offload: self.vae.to("cpu") diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py index 91fe9c43d..6fe00884c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/encoding.py @@ -97,8 +97,7 @@ class EncodingStage(PipelineStage): batch.latents = latents # Offload models if needed - if hasattr(self, "maybe_free_model_hooks"): - self.maybe_free_model_hooks() + self.maybe_free_model_hooks() if server_args.vae_cpu_offload: self.vae.to("cpu") 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 11bd59813..aff633b65 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 @@ -29,7 +29,6 @@ from sglang.multimodal_gen.runtime.models.vision_utils import ( normalize, numpy_to_pt, pil_to_numpy, - resize, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage @@ -39,7 +38,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) -from sglang.multimodal_gen.runtime.server_args import ExecutionMode, ServerArgs +from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE @@ -88,7 +87,7 @@ class ImageEncodingStage(PipelineStage): prompt_embeds = qwen_image_postprocess_text(outputs, image_inputs, 64) return prompt_embeds - @torch.inference_mode() + @torch.no_grad() def forward( self, batch: Req, @@ -112,11 +111,6 @@ class ImageEncodingStage(PipelineStage): image = batch.condition_image - # preprocess via vae_image_processor - prompt_image = server_args.pipeline_config.preprocess_image( - image, self.vae_image_processor - ) - if batch.prompt and ( isinstance(server_args.pipeline_config, QwenImageEditPipelineConfig) or isinstance(server_args.pipeline_config, QwenImagePipelineConfig) @@ -128,7 +122,7 @@ class ImageEncodingStage(PipelineStage): image_processor_kwargs = {} image_inputs = self.image_processor( - images=prompt_image, return_tensors="pt", **image_processor_kwargs + images=image, return_tensors="pt", **image_processor_kwargs ).to(cuda_device) if self.image_encoder: # if an image encoder is provided @@ -151,7 +145,7 @@ class ImageEncodingStage(PipelineStage): neg_image_processor_kwargs = {} neg_image_inputs = self.image_processor( - images=prompt_image, return_tensors="pt", **neg_image_processor_kwargs + images=image, return_tensors="pt", **neg_image_processor_kwargs ).to(get_local_torch_device()) with set_forward_context(current_timestep=0, attn_metadata=None): @@ -203,7 +197,7 @@ class ImageVAEEncodingStage(PipelineStage): Stage for encoding pixel representations into latent space. This stage handles the encoding of pixel representations into the final - input format (e.g., latents). + input format (e.g., image_latents). """ def __init__(self, vae: ParallelTiledVAE, **kwargs) -> None: @@ -229,35 +223,20 @@ class ImageVAEEncodingStage(PipelineStage): if batch.condition_image is None: return batch - if server_args.mode == ExecutionMode.INFERENCE: - assert batch.condition_image is not None and isinstance( - batch.condition_image, PIL.Image.Image - ) - assert batch.height is not None and isinstance(batch.height, int) - assert batch.width is not None and isinstance(batch.width, int) - assert batch.num_frames is not None and isinstance(batch.num_frames, int) - height = batch.height - width = batch.width - num_frames = batch.num_frames - elif server_args.mode == ExecutionMode.PREPROCESS: - assert batch.condition_image is not None and isinstance( - batch.condition_image, torch.Tensor - ) - assert batch.height is not None and isinstance(batch.height, list) - assert batch.width is not None and isinstance(batch.width, list) - assert batch.num_frames is not None and isinstance(batch.num_frames, list) - num_frames = batch.num_frames[0] - height = batch.height[0] - width = batch.width[0] + assert batch.condition_image is not None and isinstance( + batch.condition_image, PIL.Image.Image + ) + assert batch.height is not None and isinstance(batch.height, int) + assert batch.width is not None and isinstance(batch.width, int) + assert batch.num_frames is not None and isinstance(batch.num_frames, int) + + num_frames = batch.num_frames self.vae = self.vae.to(get_local_torch_device()) image = batch.condition_image image = self.preprocess( image, - vae_scale_factor=self.vae.spatial_compression_ratio, - height=height, - width=width, ).to(get_local_torch_device(), dtype=torch.float32) # (B, C, H, W) -> (B, C, 1, H, W) @@ -303,21 +282,18 @@ class ImageVAEEncodingStage(PipelineStage): video_condition ) - if server_args.mode == ExecutionMode.PREPROCESS: - latent_condition = encoder_output.mean - else: - generator = batch.generator - if generator is None: - raise ValueError("Generator must be provided") - # TODO: verify - sample_mode = ( - "argmax" - if server_args.pipeline_config.task_type == ModelTaskType.I2I - else "sample" - ) - latent_condition = self.retrieve_latents( - encoder_output, generator, sample_mode=sample_mode - ) + generator = batch.generator + if generator is None: + raise ValueError("Generator must be provided") + # TODO: verify + sample_mode = ( + "argmax" + if server_args.pipeline_config.task_type == ModelTaskType.I2I + else "sample" + ) + latent_condition = self.retrieve_latents( + encoder_output, generator, sample_mode=sample_mode + ) latent_condition = self.server_args.pipeline_config.post_process_vae_encode( latent_condition, self.vae @@ -343,97 +319,87 @@ class ImageVAEEncodingStage(PipelineStage): batch_size = batch.batch_size - if server_args.mode == ExecutionMode.PREPROCESS: - batch.image_latent = latent_condition - else: - # 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 + # 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 ) - image_latents = _pack_latents( - image_latents, - batch_size, - num_channels_latents, - image_latent_height, - image_latent_width, + 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." ) - 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 = height // self.vae.spatial_compression_ratio - latent_width = 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) + 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 - batch.image_latent = image_latents + 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) - # Offload models if needed - if hasattr(self, "maybe_free_model_hooks"): - self.maybe_free_model_hooks() + # 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() self.vae.to("cpu") @@ -455,19 +421,9 @@ class ImageVAEEncodingStage(PipelineStage): def preprocess( self, image: torch.Tensor | PIL.Image.Image, - vae_scale_factor: int, - height: int | None = None, - width: int | None = None, - resize_mode: str = "default", # "default", "fill", "crop" ) -> torch.Tensor: if isinstance(image, PIL.Image.Image): - width, height = ( - self.server_args.pipeline_config.vae_config.calculate_dimensions( - image, vae_scale_factor, width, height - ) - ) - image = resize(image, height, width, resize_mode=resize_mode) image = pil_to_numpy(image) # to np image = numpy_to_pt(image) # to pt @@ -483,14 +439,9 @@ class ImageVAEEncodingStage(PipelineStage): """Verify encoding stage inputs.""" result = VerificationResult() result.add_check("generator", batch.generator, V.generator_or_list_generators) - if server_args.mode == ExecutionMode.PREPROCESS: - result.add_check("height", batch.height, V.list_not_empty) - result.add_check("width", batch.width, V.list_not_empty) - result.add_check("num_frames", batch.num_frames, V.list_not_empty) - else: - result.add_check("height", batch.height, V.positive_int) - result.add_check("width", batch.width, V.positive_int) - result.add_check("num_frames", batch.num_frames, V.positive_int) + result.add_check("height", batch.height, V.positive_int) + result.add_check("width", batch.width, V.positive_int) + result.add_check("num_frames", batch.num_frames, V.positive_int) return result def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index 87381d384..3661a1890 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -53,6 +53,97 @@ class InputValidationStage(PipelineStage): # FIXME: the generator's in latent preparation stage seems to be different from seeds batch.generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds] + # def preprocess_condition_image(self, batch: Req, server_args: ServerArgs, condition_image_width, + # condition_image_height): + # """ + # resize condition image + # NOTE: condition image resizing is only allowed to do in InputValidationStage + # """ + # if server_args.pipeline_config.task_type == ModelTaskType.I2I: + # # calculate new condition image size + # calculated_size = ( + # server_args.pipeline_config.calculate_condition_image_size( + # batch.condition_image, + # condition_image_width, + # condition_image_height, + # ) + # ) + # + # # resize condition image if necessary + # if calculated_size is not None: + # calculated_width, calculated_height = calculated_size + # condition_image = ( + # server_args.pipeline_config.resize_condition_image( + # batch.condition_image, calculated_width, calculated_height + # ) + # ) + # batch.condition_image = condition_image + # + # # adjust output image size + # calculated_width, calculated_height = batch.condition_image.size + # width = calculated_width if batch.width_not_provided else batch.width + # height = ( + # calculated_height if batch.height_not_provided else batch.height + # ) + # multiple_of = ( + # server_args.pipeline_config.vae_config.get_vae_scale_factor() * 2 + # ) + # width = width // multiple_of * multiple_of + # height = height // multiple_of * multiple_of + # batch.width = width + # batch.height = height + # else: + # if isinstance(server_args.pipeline_config, WanI2V480PConfig): + # # TODO: could we merge with above? + # # resize image only, Wan2.1 I2V + # max_area = 720 * 1280 + # aspect_ratio = condition_image_height / condition_image_width + # mod_value = ( + # server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial + # * server_args.pipeline_config.dit_config.arch_config.patch_size[1] + # ) + # height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + # width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + # + # batch.condition_image = batch.condition_image.resize((width, height)) + # batch.height = height + # batch.width = width + # + # if ( + # server_args.pipeline_config.task_type == ModelTaskType.TI2V + # ): + # # duplicate with vae_image_processor + # # further processing for ti2v task + # img = batch.condition_image + # ih, iw = img.height, img.width + # patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size + # vae_stride = ( + # server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial + # ) + # dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride + # max_area = 704 * 1280 + # ow, oh = best_output_size(iw, ih, dw, dh, max_area) + # + # scale = max(ow / iw, oh / ih) + # img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS) + # logger.info("resized img height: %s, img width: %s", img.height, img.width) + # + # # center-crop + # x1 = (img.width - ow) // 2 + # y1 = (img.height - oh) // 2 + # img = img.crop((x1, y1, x1 + ow, y1 + oh)) + # assert img.width == ow and img.height == oh + # + # # to tensor + # img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device).unsqueeze(1) + # img = img.unsqueeze(0) + # batch.height = oh + # batch.width = ow + # # TODO: should we store in a new field: pixel values? + # height = batch.height + # width = batch.width + # batch.condition_image = resize(img, height, width) + def forward( self, batch: Req, @@ -117,24 +208,44 @@ class InputValidationStage(PipelineStage): image = load_image(batch.image_path) batch.condition_image = image condition_image_width, condition_image_height = image.width, image.height - else: - condition_image_width, condition_image_height = None, None + batch.original_condition_image_size = image.size - # NOTE: resizing needs to be bring in advance + # self.preprocess_condition_image(batch, server_args, condition_image_width, condition_image_height) + # NOTE: condition image resizing is only allowed to do in InputValidationStage if server_args.pipeline_config.task_type == ModelTaskType.I2I: if batch.condition_image is not None: - resized_image, resized_width, resized_height = ( - server_args.pipeline_config.maybe_resize_condition_image( + # calculate new condition image size + calculated_size = ( + server_args.pipeline_config.calculate_condition_image_size( + batch.condition_image, condition_image_width, condition_image_height, - batch.condition_image, ) ) - batch.condition_image = resized_image - batch.width = resized_width if batch.width_not_provided else batch.width - batch.height = ( - resized_height if batch.height_not_provided else batch.height + + # resize condition image if necessary + if calculated_size is not None: + calculated_width, calculated_height = calculated_size + condition_image = ( + server_args.pipeline_config.resize_condition_image( + image, calculated_width, calculated_height + ) + ) + batch.condition_image = condition_image + + # adjust output image size + calculated_width, calculated_height = batch.condition_image.size + width = calculated_width if batch.width_not_provided else batch.width + height = ( + calculated_height if batch.height_not_provided else batch.height ) + multiple_of = ( + server_args.pipeline_config.vae_config.get_vae_scale_factor() * 2 + ) + width = width // multiple_of * multiple_of + height = height // multiple_of * multiple_of + batch.width = width + batch.height = height elif ( server_args.pipeline_config.task_type == ModelTaskType.TI2V ) and batch.condition_image is not None: diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 9d1e274eb..4abc5cbfc 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -156,9 +156,6 @@ class ExecutionMode(str, Enum): """ INFERENCE = "inference" - PREPROCESS = "preprocess" - FINETUNING = "finetuning" - DISTILLATION = "distillation" @classmethod def from_string(cls, value: str) -> "ExecutionMode":