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

[diffusion] fix: fix Flux.2 condition image resize (#14232)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Mick
2025-12-02 10:05:44 +08:00
committed by GitHub
parent 03888b9de5
commit 3ab8ae6847
15 changed files with 235 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/multimodal_gen/configs/models/vaes/base.py
View File

@@ -143,6 +143,9 @@ class VAEConfig(ModelConfig):
def get_vae_scale_factor(self):
return 2 ** (len(self.arch_config.block_out_channels) - 1)
def encode_sample_mode(self):
return "argmax"
@classmethod
def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig":
kwargs = {}

11
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
View File

@@ -190,8 +190,15 @@ class PipelineConfig:
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 preprocess_condition_image(
self, image, target_width, target_height, _vae_image_processor
):
"""
preprocess the condition image, returns (image, final_image_width, final_image_height)
"""
return image.resize(
(target_width, target_height), PIL.Image.Resampling.LANCZOS
), (target_width, target_height)
def prepare_image_processor_kwargs(self, batch):
return {}

16
python/sglang/multimodal_gen/configs/pipeline_configs/flux.py
View File

@@ -4,6 +4,7 @@ from typing import Callable, List, Optional
import PIL
import torch
from diffusers.image_processor import VaeImageProcessor
from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
@@ -465,8 +466,19 @@ class Flux2PipelineConfig(FluxPipelineConfig):
return None
def resize_condition_image(self, image, target_width, target_height):
return image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS)
def preprocess_condition_image(
self, image, target_width, target_height, vae_image_processor: VaeImageProcessor
):
img = image.resize((target_width, target_height), PIL.Image.Resampling.LANCZOS)
image_width, image_height = img.size
vae_scale_factor = self.vae_config.arch_config.vae_scale_factor
multiple_of = vae_scale_factor * 2
image_width = (image_width // multiple_of) * multiple_of
image_height = (image_height // multiple_of) * multiple_of
img = vae_image_processor.preprocess(
img, height=image_height, width=image_width, resize_mode="crop"
)
return img, (image_width, image_height)
def postprocess_image_latent(self, latent_condition, batch):
batch_size = batch.batch_size

9
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
View File

@@ -273,8 +273,13 @@ 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 preprocess_condition_image(
self, image, target_width, target_height, _vae_image_processor
):
return resize(image, target_height, target_width, resize_mode="default"), (
target_width,
target_height,
)
def postprocess_image_latent(self, latent_condition, batch):
batch_size = batch.batch_size

3
python/sglang/multimodal_gen/configs/sample/sampling_params.py
View File

@@ -299,7 +299,6 @@ class SamplingParams:
from sglang.multimodal_gen.registry import get_model_info
model_info = get_model_info(model_path)
logger.debug(f"Found model info: {model_info}")
if model_info is not None:
sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs)
else:
@@ -317,8 +316,6 @@ class SamplingParams:
user_sampling_params = SamplingParams(*args, **kwargs)
# TODO: refactor
sampling_params._merge_with_user_params(user_sampling_params)
sampling_params.width_not_provided = user_sampling_params.width is None
sampling_params.height_not_provided = user_sampling_params.height is None
sampling_params._adjust(server_args)
return sampling_params

7
python/sglang/multimodal_gen/registry.py
View File

@@ -260,12 +260,15 @@ def get_model_info(model_path: str) -> Optional[ModelInfo]:
)
return None
# 4. Combine and return the complete model info
return ModelInfo(
# 4. Combine the complete model info
model_info = ModelInfo(
pipeline_cls=pipeline_cls,
sampling_param_cls=config_info.sampling_param_cls,
pipeline_config_cls=config_info.pipeline_config_cls,
)
logger.info(f"Found model info: {model_info}")
return model_info
# Registration of model configs

1
python/sglang/multimodal_gen/runtime/loader/__init__.py
View File

@@ -1 +0,0 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

46
python/sglang/multimodal_gen/runtime/models/utils.py
View File

@@ -8,52 +8,6 @@ from typing import Any
import torch
# TODO(PY): move it elsewhere
def auto_attributes(init_func):
"""
Decorator that automatically adds all initialization arguments as object attributes.
Example:
@auto_attributes
def __init__(self, a=1, b=2):
pass
# This will automatically set:
# - self.a = 1 and self.b = 2
# - self.config.a = 1 and self.config.b = 2
"""
def wrapper(self, *args, **kwargs):
# Get the function signature
import inspect
signature = inspect.signature(init_func)
parameters = signature.parameters
# Get parameter names (excluding 'self')
param_names = list(parameters.keys())[1:]
# Bind arguments to parameters
bound_args = signature.bind(self, *args, **kwargs)
bound_args.apply_defaults()
# Create config object if it doesn't exist
if not hasattr(self, "config"):
self.config = type("Config", (), {})()
# Set attributes on self and self.config
for name in param_names:
if name in bound_args.arguments:
value = bound_args.arguments[name]
setattr(self, name, value)
setattr(self.config, name, value)
# Call the original __init__ function
return init_func(self, *args, **kwargs)
return wrapper
def set_weight_attrs(
weight: torch.Tensor,
weight_attrs: dict[str, Any] | None,

8
python/sglang/multimodal_gen/runtime/pipelines/flux_2.py
View File

@@ -58,7 +58,13 @@ class Flux2Pipeline(LoRAPipeline, ComposedPipelineBase):
"""Set up pipeline stages with proper dependency injection."""
self.add_stage(
stage_name="input_validation_stage", stage=InputValidationStage()
stage_name="input_validation_stage",
stage=InputValidationStage(
vae_image_processor=VaeImageProcessor(
vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor
* 2
),
),
)
self.add_stage(

5
python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py
View File

@@ -1,7 +1,6 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from diffusers.image_processor import VaeImageProcessor
from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
@@ -144,10 +143,6 @@ class QwenImageEditPipeline(LoRAPipeline, ComposedPipelineBase):
stage=ImageEncodingStage(
image_processor=self.get_module("processor"),
text_encoder=self.get_module("text_encoder"),
vae_image_processor=VaeImageProcessor(
vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor
* 2
),
),
)

7
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
View File

@@ -27,7 +27,6 @@ from sglang.multimodal_gen.configs.sample.teacache import (
from sglang.multimodal_gen.runtime.server_args import ServerArgs
if TYPE_CHECKING:
from torchcodec.decoders import VideoDecoder
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestTimings
@@ -240,9 +239,3 @@ class OutputBatch:
# logged timings info, directly from Req.timings
timings: Optional["RequestTimings"] = None
@dataclass
class PreprocessBatch(Req):
video_loader: list["VideoDecoder"] | list[str] = field(default_factory=list)
video_file_name: list[str] = field(default_factory=list)

29
python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py
View File

@@ -11,7 +11,6 @@ import PIL
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.qwen_image import (
qwen_image_postprocess_text,
)
@@ -51,7 +50,6 @@ class ImageEncodingStage(PipelineStage):
image_processor,
image_encoder=None,
text_encoder=None,
vae_image_processor=None,
) -> None:
"""
Initialize the prompt encoding stage.
@@ -61,7 +59,6 @@ class ImageEncodingStage(PipelineStage):
"""
super().__init__()
self.image_processor = image_processor
self.vae_image_processor = vae_image_processor
self.image_encoder = image_encoder
self.text_encoder = text_encoder
@@ -210,13 +207,6 @@ class ImageVAEEncodingStage(PipelineStage):
if batch.condition_image is None:
return batch
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())
@@ -272,16 +262,12 @@ class ImageVAEEncodingStage(PipelineStage):
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"
)
sample_mode = server_args.pipeline_config.vae_config.encode_sample_mode()
latent_condition = self.retrieve_latents(
encoder_output, generator, sample_mode=sample_mode
)
latent_condition = server_args.pipeline_config.postprocess_vae_encode(
latent_condition, self.vae
)
@@ -347,6 +333,15 @@ class ImageVAEEncodingStage(PipelineStage):
def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
"""Verify encoding stage inputs."""
result = VerificationResult()
assert batch.condition_image is None or (
isinstance(batch.condition_image, PIL.Image.Image)
or isinstance(batch.condition_image, torch.Tensor)
)
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)
result.add_check("generator", batch.generator, V.generator_or_list_generators)
result.add_check("height", batch.height, V.positive_int)
result.add_check("width", batch.width, V.positive_int)

270
python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py
View File

@@ -41,6 +41,10 @@ class InputValidationStage(PipelineStage):
In this stage, input image and output image may be resized
"""
def __init__(self, vae_image_processor=None):
super().__init__()
self.vae_image_processor = vae_image_processor
def _generate_seeds(self, batch: Req, server_args: ServerArgs):
"""Generate seeds for the inference"""
seed = batch.seed
@@ -53,96 +57,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 preprocess_condition_image(
self,
batch: Req,
server_args: ServerArgs,
condition_image_width,
condition_image_height,
):
"""
preprocess condition image
NOTE: condition image resizing is only allowed 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,
)
)
# preprocess condition image if necessary
if calculated_size is not None:
calculated_width, calculated_height = calculated_size
condition_image, calculated_size = (
server_args.pipeline_config.preprocess_condition_image(
batch.condition_image,
calculated_width,
calculated_height,
self.vae_image_processor,
)
)
batch.condition_image = condition_image
# adjust output image size
calculated_width, calculated_height = calculated_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:
# 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?
batch.condition_image = img
elif 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
def forward(
self,
@@ -210,90 +215,9 @@ class InputValidationStage(PipelineStage):
condition_image_width, condition_image_height = image.width, image.height
batch.original_condition_image_size = image.size
# 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:
# 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(
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:
# 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
self.preprocess_condition_image(
batch, server_args, condition_image_width, condition_image_height
)
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?
batch.condition_image = img
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 = image.height / 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
return batch

67
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -295,6 +295,73 @@
"expected_avg_denoise_ms": 520.09,
"expected_median_denoise_ms": 528.0
},
"flux_2_ti2i": {
"stages_ms": {
"InputValidationStage": 99.82,
"TextEncodingStage": 519.88,
"ImageVAEEncodingStage": 254.56,
"ConditioningStage": 0.01,
"LatentPreparationStage": 12.4,
"TimestepPreparationStage": 2.71,
"DenoisingStage": 54705.41,
"DecodingStage": 311.13
},
"denoise_step_ms": {
"0": 1067.03,
"1": 271.58,
"2": 1073.07,
"3": 1071.93,
"4": 1100.0,
"5": 1102.28,
"6": 1088.3,
"7": 1089.09,
"8": 1086.95,
"9": 1089.33,
"10": 1089.28,
"11": 1096.51,
"12": 1098.88,
"13": 1080.84,
"14": 1098.44,
"15": 1100.88,
"16": 1086.83,
"17": 1090.58,
"18": 1096.35,
"19": 1086.25,
"20": 1082.71,
"21": 1097.6,
"22": 1098.72,
"23": 1100.9,
"24": 1099.02,
"25": 1101.52,
"26": 1098.75,
"27": 1101.41,
"28": 1091.75,
"29": 1087.2,
"30": 1101.33,
"31": 1098.14,
"32": 1100.14,
"33": 1098.91,
"34": 1100.05,
"35": 1099.12,
"36": 1100.22,
"37": 1103.29,
"38": 1092.79,
"39": 1086.59,
"40": 1094.81,
"41": 1105.6,
"42": 1100.54,
"43": 1099.95,
"44": 1096.5,
"45": 1086.69,
"46": 1095.85,
"47": 1092.85,
"48": 1086.17,
"49": 1099.67
},
"expected_e2e_ms": 56308.23,
"expected_avg_denoise_ms": 1077.26,
"expected_median_denoise_ms": 1096.5
},
"flux_image_t2i_2_gpus": {
"stages_ms": {
"InputValidationStage": 0.03,

15
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -125,7 +125,7 @@ class DiffusionServerArgs:
class DiffusionSamplingParams:
"""Configuration for a single model/scenario test case."""
output_size: str = "1024x1024" # output image dimensions (or video resolution)
output_size: str = ""
# inputs and conditioning
prompt: str | None = None # text prompt for generation
@@ -310,6 +310,19 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
# warmup_edit=0,
# custom_validator="video",
# ),
DiffusionTestCase(
"flux_2_ti2i",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-dev",
modality="image",
warmup_text=0,
warmup_edit=1,
),
DiffusionSamplingParams(
prompt="Convert 2D style to 3D style",
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
),
),
DiffusionTestCase(
"fast_hunyuan_video",
DiffusionServerArgs(