chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] refactor: refactor condition image resize logic (#14079)

Browse Source
This commit is contained in:
Mick
2025-11-28 14:06:34 +08:00
committed by GitHub
parent e12c78aab6
commit 3543a04a48
13 changed files with 268 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/vaes/base.py
+3 -8
View File
@@ -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 = {}
python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py
+2 -9
View File
@@ -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 = (
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+13 -11
View File
@@ -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
python/sglang/multimodal_gen/configs/pipeline_configs/flux.py
+9 -6
View File
@@ -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):
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
+8 -20
View File
@@ -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