[diffusion] model: support mutli-image input and qwen-image-edit-2509 (#15005)

This commit is contained in:
Yuhao Yang
2025-12-15 16:17:10 +08:00
committed by GitHub
parent 16e6bc20b0
commit 01b955ac3d
15 changed files with 577 additions and 221 deletions
@@ -208,7 +208,10 @@ class PipelineConfig:
(target_width, target_height), PIL.Image.Resampling.LANCZOS (target_width, target_height), PIL.Image.Resampling.LANCZOS
), (target_width, target_height) ), (target_width, target_height)
def prepare_image_processor_kwargs(self, batch): def prepare_calculated_size(self, image):
return self.calculate_condition_image_size(image, image.width, image.height)
def prepare_image_processor_kwargs(self, batch, neg=False):
return {} return {}
def postprocess_image_latent(self, latent_condition, batch): def postprocess_image_latent(self, latent_condition, batch):
@@ -297,6 +300,9 @@ class PipelineConfig:
latents = sequence_model_parallel_all_gather(latents, dim=2) latents = sequence_model_parallel_all_gather(latents, dim=2)
return latents return latents
def preprocess_vae_image(self, batch, vae_image_processor):
pass
def shard_latents_for_sp(self, batch, latents): def shard_latents_for_sp(self, batch, latents):
# general logic for video models # general logic for video models
sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank() sp_world_size, rank_in_sp_group = get_sp_world_size(), get_sp_parallel_rank()
@@ -107,8 +107,9 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
def prepare_sigmas(self, sigmas, num_inference_steps): def prepare_sigmas(self, sigmas, num_inference_steps):
return self._prepare_sigmas(sigmas, num_inference_steps) return self._prepare_sigmas(sigmas, num_inference_steps)
def prepare_image_processor_kwargs(self, batch): def prepare_image_processor_kwargs(self, batch, neg=False):
if batch.prompt: prompt = batch.prompt if not neg else batch.negative_prompt
if 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" 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) txt = prompt_template_encode.format(batch.prompt)
return dict(text=[txt], padding=True) return dict(text=[txt], padding=True)
@@ -328,3 +329,137 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
# remove noise over input image # remove noise over input image
noise = noise[:, : latents.size(1)] noise = noise[:, : latents.size(1)]
return noise return noise
CONDITION_IMAGE_SIZE = 384 * 384
VAE_IMAGE_SIZE = 1024 * 1024
@dataclass
class QwenImageEditPlusPipelineConfig(QwenImageEditPipelineConfig):
task_type: ModelTaskType = ModelTaskType.I2I
def _get_condition_image_sizes(self, batch) -> list[tuple[int, int]]:
image = batch.condition_image
if not isinstance(image, list):
image = [image]
condition_image_sizes = []
for img in image:
image_width, image_height = img.size
edit_width, edit_height, _ = calculate_dimensions(
VAE_IMAGE_SIZE, image_width / image_height
)
condition_image_sizes.append((edit_width, edit_height))
return condition_image_sizes
def prepare_image_processor_kwargs(self, batch, neg=False) -> dict:
prompt = batch.prompt if not neg else batch.negative_prompt
prompt_list = [prompt] if isinstance(prompt, str) else prompt
image_list = batch.condition_image
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{}<|im_end|>\n"
"<|im_start|>assistant\n"
)
img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
if isinstance(image_list, list):
base_img_prompt = ""
for i, img in enumerate(image_list):
base_img_prompt += img_prompt_template.format(i + 1)
txt = [prompt_template_encode.format(base_img_prompt + p) for p in prompt_list]
return dict(text=txt, padding=True)
def prepare_calculated_size(self, image):
return self.calculate_vae_image_size(image, image.width, image.height)
def resize_condition_image(self, images, target_width, target_height):
if not isinstance(images, list):
images = [images]
new_images = []
for img, width, height in zip(images, target_width, target_height):
new_images.append(resize(img, height, width, resize_mode="default"))
return new_images
def calculate_condition_image_size(self, image, width, height) -> tuple[int, int]:
calculated_width, calculated_height, _ = calculate_dimensions(
CONDITION_IMAGE_SIZE, width / height
)
return calculated_width, calculated_height
def calculate_vae_image_size(self, image, width, height) -> tuple[int, int]:
calculated_width, calculated_height, _ = calculate_dimensions(
VAE_IMAGE_SIZE, width / height
)
return calculated_width, calculated_height
def preprocess_vae_image(self, batch, vae_image_processor):
if not isinstance(batch.condition_image, list):
batch.condition_image = [batch.condition_image]
new_images = []
vae_image_sizes = []
for img in batch.condition_image:
width, height = self.calculate_vae_image_size(img, img.width, img.height)
new_images.append(vae_image_processor.preprocess(img, height, width))
vae_image_sizes.append((width, height))
batch.vae_image = new_images
batch.vae_image_sizes = vae_image_sizes
return batch
def _prepare_edit_cond_kwargs(
self, batch, prompt_embeds, rotary_emb, device, dtype
):
batch_size = batch.latents.shape[0]
assert batch_size == 1
height = batch.height
width = batch.width
image_size = batch.original_condition_image_size
vae_scale_factor = self.get_vae_scale_factor()
img_shapes = [
[
(1, height // vae_scale_factor // 2, width // vae_scale_factor // 2),
*[
(
1,
vae_height // vae_scale_factor // 2,
vae_width // vae_scale_factor // 2,
)
for vae_width, vae_height in batch.vae_image_sizes
],
],
] * batch_size
txt_seq_lens = [prompt_embeds[0].shape[1]]
(img_cos, img_sin), (txt_cos, txt_sin) = (
QwenImageEditPlusPipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
)
)
# perform sp shard on noisy image tokens
noisy_img_seq_len = (
1 * (height // vae_scale_factor // 2) * (width // vae_scale_factor // 2)
)
noisy_img_cos = shard_rotary_emb_for_sp(img_cos[:noisy_img_seq_len, :])
noisy_img_sin = shard_rotary_emb_for_sp(img_sin[:noisy_img_seq_len, :])
# concat back the img_cos for input image (since it is not sp-shared later)
img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to(
device=device
)
img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to(
device=device
)
return {
"txt_seq_lens": txt_seq_lens,
"freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
}
@@ -13,3 +13,11 @@ class QwenImageSamplingParams(SamplingParams):
# Denoising stage # Denoising stage
guidance_scale: float = 4.0 guidance_scale: float = 4.0
num_inference_steps: int = 50 num_inference_steps: int = 50
@dataclass
class QwenImageEditPlusSamplingParams(QwenImageSamplingParams):
# Denoising stage
guidance_scale: float = 4.0
# true_cfg_scale: float = 4.0
num_inference_steps: int = 40
+11 -1
View File
@@ -29,6 +29,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
QwenImageEditPipelineConfig, QwenImageEditPipelineConfig,
QwenImageEditPlusPipelineConfig,
QwenImagePipelineConfig, QwenImagePipelineConfig,
) )
from sglang.multimodal_gen.configs.pipeline_configs.wan import ( from sglang.multimodal_gen.configs.pipeline_configs.wan import (
@@ -43,7 +44,10 @@ from sglang.multimodal_gen.configs.sample.hunyuan import (
FastHunyuanSamplingParam, FastHunyuanSamplingParam,
HunyuanSamplingParams, HunyuanSamplingParams,
) )
from sglang.multimodal_gen.configs.sample.qwenimage import QwenImageSamplingParams from sglang.multimodal_gen.configs.sample.qwenimage import (
QwenImageEditPlusSamplingParams,
QwenImageSamplingParams,
)
from sglang.multimodal_gen.configs.sample.stepvideo import StepVideoT2VSamplingParams from sglang.multimodal_gen.configs.sample.stepvideo import StepVideoT2VSamplingParams
from sglang.multimodal_gen.configs.sample.wan import ( from sglang.multimodal_gen.configs.sample.wan import (
FastWanT2V480PConfig, FastWanT2V480PConfig,
@@ -427,5 +431,11 @@ def _register_configs():
hf_model_paths=["Qwen/Qwen-Image-Edit"], hf_model_paths=["Qwen/Qwen-Image-Edit"],
) )
register_configs(
sampling_param_cls=QwenImageEditPlusSamplingParams,
pipeline_config_cls=QwenImageEditPlusPipelineConfig,
hf_model_paths=["Qwen/Qwen-Image-Edit-2509"],
)
_register_configs() _register_configs()
@@ -82,14 +82,18 @@ def _build_sampling_params_from_request(
def _build_req_from_sampling(s: SamplingParams) -> Req: def _build_req_from_sampling(s: SamplingParams) -> Req:
# TODO: refactor this! this is so dangerous because we could forget to add a new field here!
return Req( return Req(
request_id=s.request_id, request_id=s.request_id,
data_type=s.data_type, data_type=s.data_type,
prompt=s.prompt, prompt=s.prompt,
negative_prompt=s.negative_prompt,
image_path=s.image_path, image_path=s.image_path,
height=s.height, height=s.height,
width=s.width, width=s.width,
fps=1, fps=1,
guidance_scale=s.guidance_scale,
num_inference_steps=s.num_inference_steps,
num_frames=s.num_frames, num_frames=s.num_frames,
seed=s.seed, seed=s.seed,
generator_device=s.generator_device, generator_device=s.generator_device,
@@ -172,12 +176,17 @@ async def edits(
if not images or len(images) == 0: if not images or len(images) == 0:
raise HTTPException(status_code=422, detail="Field 'image' is required") raise HTTPException(status_code=422, detail="Field 'image' is required")
# Save first input image; additional images or mask are not yet used by the pipeline # Save all input images; additional images beyond the first are saved for potential future use
uploads_dir = os.path.join("outputs", "uploads") uploads_dir = os.path.join("outputs", "uploads")
os.makedirs(uploads_dir, exist_ok=True) os.makedirs(uploads_dir, exist_ok=True)
first_image = images[0] if images is not None and not isinstance(images, list):
input_path = os.path.join(uploads_dir, f"{request_id}_{first_image.filename}") images = [images]
await _save_upload_to_path(first_image, input_path) input_paths = []
for idx, img in enumerate(images):
filename = img.filename or f"image_{idx}"
input_path = os.path.join(uploads_dir, f"{request_id}_{idx}_{filename}")
await _save_upload_to_path(img, input_path)
input_paths.append(input_path)
sampling = _build_sampling_params_from_request( sampling = _build_sampling_params_from_request(
request_id=request_id, request_id=request_id,
@@ -186,7 +195,7 @@ async def edits(
size=size, size=size,
output_format=output_format, output_format=output_format,
background=background, background=background,
image_path=input_path, image_path=input_paths,
seed=seed, seed=seed,
generator_device=generator_device, generator_device=generator_device,
) )
@@ -200,6 +209,8 @@ async def edits(
"id": request_id, "id": request_id,
"created_at": int(time.time()), "created_at": int(time.time()),
"file_path": save_file_path, "file_path": save_file_path,
"input_image_paths": input_paths, # Store all input image paths
"num_input_images": len(input_paths),
}, },
) )
@@ -1,6 +1,7 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0 # 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 import LoRAPipeline
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
@@ -45,15 +46,16 @@ def prepare_mu(batch: Req, server_args: ServerArgs):
height = batch.height height = batch.height
width = batch.width width = batch.width
vae_scale_factor = server_args.pipeline_config.vae_config.vae_scale_factor vae_scale_factor = server_args.pipeline_config.vae_config.vae_scale_factor
image_seq_len = (int(height) // vae_scale_factor) * (int(width) // vae_scale_factor) image_seq_len = (int(height) // vae_scale_factor // 2) * (
int(width) // vae_scale_factor // 2
)
mu = calculate_shift( mu = calculate_shift(
image_seq_len, image_seq_len,
# hard code, since scheduler_config is not in PipelineConfig now # hard code, since scheduler_config is not in PipelineConfig now
256, 256,
4096, 8192,
0.5, 0.5,
1.15, 0.9,
) )
return "mu", mu return "mu", mu
@@ -135,7 +137,13 @@ class QwenImageEditPipeline(LoRAPipeline, ComposedPipelineBase):
"""Set up pipeline stages with proper dependency injection.""" """Set up pipeline stages with proper dependency injection."""
self.add_stage( 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( self.add_stage(
@@ -184,4 +192,8 @@ class QwenImageEditPipeline(LoRAPipeline, ComposedPipelineBase):
) )
EntryClass = [QwenImagePipeline, QwenImageEditPipeline] class QwenImageEditPlusPipeline(QwenImageEditPipeline):
pipeline_name = "QwenImageEditPlusPipeline"
EntryClass = [QwenImagePipeline, QwenImageEditPipeline, QwenImageEditPlusPipeline]
@@ -57,6 +57,7 @@ class Req:
original_condition_image_size: tuple[int, int] = None original_condition_image_size: tuple[int, int] = None
condition_image: torch.Tensor | PIL.Image.Image | None = None condition_image: torch.Tensor | PIL.Image.Image | None = None
vae_image: torch.Tensor | PIL.Image.Image | None = None
pixel_values: torch.Tensor | PIL.Image.Image | None = None pixel_values: torch.Tensor | PIL.Image.Image | None = None
preprocessed_image: torch.Tensor | None = None preprocessed_image: torch.Tensor | None = None
@@ -933,6 +933,8 @@ class DenoisingStage(PipelineStage):
denoising_start_time = time.time() denoising_start_time = time.time()
# to avoid device-sync caused by timestep comparison # to avoid device-sync caused by timestep comparison
self.scheduler.set_begin_index(0)
timesteps_cpu = timesteps.cpu() timesteps_cpu = timesteps.cpu()
num_timesteps = timesteps_cpu.shape[0] num_timesteps = timesteps_cpu.shape[0]
with torch.autocast( with torch.autocast(
@@ -121,12 +121,11 @@ class ImageEncodingStage(PipelineStage):
elif self.text_encoder: elif self.text_encoder:
# if a text encoder is provided, e.g. Qwen-Image-Edit # if a text encoder is provided, e.g. Qwen-Image-Edit
# 1. neg prompt embeds # 1. neg prompt embeds
if batch.prompt: neg_image_processor_kwargs = (
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" server_args.pipeline_config.prepare_image_processor_kwargs(
txt = prompt_template_encode.format(batch.negative_prompt) batch, neg=True
neg_image_processor_kwargs = dict(text=[txt], padding=True) )
else: )
neg_image_processor_kwargs = {}
neg_image_inputs = self.image_processor( neg_image_inputs = self.image_processor(
images=image, return_tensors="pt", **neg_image_processor_kwargs images=image, return_tensors="pt", **neg_image_processor_kwargs
@@ -211,89 +210,98 @@ class ImageVAEEncodingStage(PipelineStage):
self.vae = self.vae.to(get_local_torch_device()) self.vae = self.vae.to(get_local_torch_device())
image = batch.condition_image images = (
image = self.preprocess( batch.vae_image if batch.vae_image is not None else batch.condition_image
image,
).to(get_local_torch_device(), dtype=torch.float32)
# (B, C, H, W) -> (B, C, 1, H, W)
image = image.unsqueeze(2)
if num_frames == 1:
video_condition = image
else:
video_condition = torch.cat(
[
image,
image.new_zeros(
image.shape[0],
image.shape[1],
num_frames - 1,
image.shape[3],
image.shape[4],
),
],
dim=2,
)
video_condition = video_condition.to(
device=get_local_torch_device(), dtype=torch.float32
) )
if not isinstance(images, list):
images = [images]
# Setup VAE precision all_image_latents = []
vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] for image in images:
vae_autocast_enabled = ( image = self.preprocess(
vae_dtype != torch.float32 image,
) and not server_args.disable_autocast ).to(get_local_torch_device(), dtype=torch.float32)
# Encode Image # (B, C, H, W) -> (B, C, 1, H, W)
with torch.autocast( image = image.unsqueeze(2)
device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled
): if num_frames == 1:
if server_args.pipeline_config.vae_tiling: video_condition = image
self.vae.enable_tiling() else:
# if server_args.vae_sp: video_condition = torch.cat(
# self.vae.enable_parallel() [
if not vae_autocast_enabled: image,
video_condition = video_condition.to(vae_dtype) image.new_zeros(
encoder_output: DiagonalGaussianDistribution = self.vae.encode( image.shape[0],
video_condition image.shape[1],
num_frames - 1,
image.shape[3],
image.shape[4],
),
],
dim=2,
)
video_condition = video_condition.to(
device=get_local_torch_device(), dtype=torch.float32
) )
generator = batch.generator # Setup VAE precision
if generator is None: vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
raise ValueError("Generator must be provided") vae_autocast_enabled = (
vae_dtype != torch.float32
) and not server_args.disable_autocast
sample_mode = server_args.pipeline_config.vae_config.encode_sample_mode() # Encode Image
with torch.autocast(
device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled
):
if server_args.pipeline_config.vae_tiling:
self.vae.enable_tiling()
# if server_args.vae_sp:
# self.vae.enable_parallel()
if not vae_autocast_enabled:
video_condition = video_condition.to(vae_dtype)
encoder_output: DiagonalGaussianDistribution = self.vae.encode(
video_condition
)
latent_condition = self.retrieve_latents( generator = batch.generator
encoder_output, generator, sample_mode=sample_mode if generator is None:
) raise ValueError("Generator must be provided")
latent_condition = server_args.pipeline_config.postprocess_vae_encode(
latent_condition, self.vae
)
scaling_factor, shift_factor = ( sample_mode = server_args.pipeline_config.vae_config.encode_sample_mode()
server_args.pipeline_config.get_decode_scale_and_shift(
device=latent_condition.device, latent_condition = self.retrieve_latents(
dtype=latent_condition.dtype, encoder_output, generator, sample_mode=sample_mode
vae=self.vae, )
latent_condition = server_args.pipeline_config.postprocess_vae_encode(
latent_condition, self.vae
) )
)
# apply shift & scale if needed scaling_factor, shift_factor = (
if isinstance(shift_factor, torch.Tensor): server_args.pipeline_config.get_decode_scale_and_shift(
shift_factor = shift_factor.to(latent_condition.device) device=latent_condition.device,
dtype=latent_condition.dtype,
vae=self.vae,
)
)
if isinstance(scaling_factor, torch.Tensor): # apply shift & scale if needed
scaling_factor = scaling_factor.to(latent_condition.device) if isinstance(shift_factor, torch.Tensor):
shift_factor = shift_factor.to(latent_condition.device)
latent_condition -= shift_factor if isinstance(scaling_factor, torch.Tensor):
latent_condition = latent_condition * scaling_factor scaling_factor = scaling_factor.to(latent_condition.device)
batch.image_latent = server_args.pipeline_config.postprocess_image_latent( latent_condition -= shift_factor
latent_condition, batch latent_condition = latent_condition * scaling_factor
)
image_latent = server_args.pipeline_config.postprocess_image_latent(
latent_condition, batch
)
all_image_latents.append(image_latent)
batch.image_latent = torch.cat(all_image_latents, dim=1)
self.maybe_free_model_hooks() self.maybe_free_model_hooks()
self.vae.to("cpu") self.vae.to("cpu")
@@ -337,6 +345,7 @@ class ImageVAEEncodingStage(PipelineStage):
assert batch.condition_image is None or ( assert batch.condition_image is None or (
isinstance(batch.condition_image, PIL.Image.Image) isinstance(batch.condition_image, PIL.Image.Image)
or isinstance(batch.condition_image, torch.Tensor) or isinstance(batch.condition_image, torch.Tensor)
or isinstance(batch.condition_image, list)
) )
assert batch.height is not None and isinstance(batch.height, int) assert batch.height is not None and isinstance(batch.height, int)
assert batch.width is not None and isinstance(batch.width, int) assert batch.width is not None and isinstance(batch.width, int)
@@ -80,41 +80,48 @@ class InputValidationStage(PipelineStage):
""" """
if server_args.pipeline_config.task_type == ModelTaskType.I2I: if server_args.pipeline_config.task_type == ModelTaskType.I2I:
# calculate new condition image size # calculate new condition image size
calculated_size = ( if not isinstance(batch.condition_image, list):
server_args.pipeline_config.calculate_condition_image_size( batch.condition_image = [batch.condition_image]
batch.condition_image,
condition_image_width,
condition_image_height,
)
)
# preprocess condition image if necessary processed_images = []
if calculated_size is not None: final_image = batch.condition_image[-1]
calculated_width, calculated_height = calculated_size config = server_args.pipeline_config
condition_image, calculated_size = ( config.preprocess_vae_image(batch, self.vae_image_processor)
server_args.pipeline_config.preprocess_condition_image(
batch.condition_image, for img in batch.condition_image:
calculated_width, size = config.calculate_condition_image_size(img, img.width, img.height)
calculated_height, if size is not None:
self.vae_image_processor, width, height = size
img, _ = config.preprocess_condition_image(
img, width, height, self.vae_image_processor
) )
)
batch.condition_image = condition_image processed_images.append(img)
batch.condition_image = processed_images
calculated_size = config.prepare_calculated_size(final_image)
# adjust output image size # adjust output image size
calculated_width, calculated_height = calculated_size if calculated_size is not None:
width = batch.width or calculated_width calculated_width, calculated_height = calculated_size
height = batch.height or calculated_height width = batch.width or calculated_width
multiple_of = ( height = batch.height or calculated_height
server_args.pipeline_config.vae_config.get_vae_scale_factor() * 2 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 width = width // multiple_of * multiple_of
batch.width = width height = height // multiple_of * multiple_of
batch.height = height batch.width = width
batch.height = height
elif server_args.pipeline_config.task_type == ModelTaskType.TI2V: elif server_args.pipeline_config.task_type == ModelTaskType.TI2V:
# duplicate with vae_image_processor # duplicate with vae_image_processor
# further processing for ti2v task # further processing for ti2v task
if isinstance(
batch.condition_image, list
): # not support multi image input yet.
batch.condition_image = batch.condition_image[0]
img = batch.condition_image img = batch.condition_image
ih, iw = img.height, img.width ih, iw = img.height, img.width
patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size
@@ -146,6 +153,11 @@ class InputValidationStage(PipelineStage):
elif isinstance(server_args.pipeline_config, WanI2V480PConfig): elif isinstance(server_args.pipeline_config, WanI2V480PConfig):
# TODO: could we merge with above? # TODO: could we merge with above?
# resize image only, Wan2.1 I2V # resize image only, Wan2.1 I2V
if isinstance(batch.condition_image, list):
batch.condition_image = batch.condition_image[
0
] # not support multi image input yet.
max_area = server_args.pipeline_config.max_area max_area = server_args.pipeline_config.max_area
aspect_ratio = condition_image_height / condition_image_width aspect_ratio = condition_image_height / condition_image_width
mod_value = ( mod_value = (
@@ -207,13 +219,33 @@ class InputValidationStage(PipelineStage):
# for i2v, get image from image_path # for i2v, get image from image_path
# @TODO(Wei) hard-coded for wan2.2 5b ti2v for now. Should put this in image_encoding stage # @TODO(Wei) hard-coded for wan2.2 5b ti2v for now. Should put this in image_encoding stage
if batch.image_path is not None: if batch.image_path is not None:
if batch.image_path.endswith(".mp4"): if isinstance(batch.image_path, list):
image = load_video(batch.image_path)[0] batch.condition_image = []
for path in batch.image_path:
if path.endswith(".mp4"):
image = load_video(path)[0]
else:
image = load_image(path)
batch.condition_image.append(image)
# Use the first image for size reference
condition_image_width = batch.condition_image[0].width
condition_image_height = batch.condition_image[0].height
batch.original_condition_image_size = (
condition_image_width,
condition_image_height,
)
else: else:
image = load_image(batch.image_path) if batch.image_path.endswith(".mp4"):
batch.condition_image = image image = load_video(batch.image_path)[0]
condition_image_width, condition_image_height = image.width, image.height else:
batch.original_condition_image_size = image.size image = load_image(batch.image_path)
batch.condition_image = image
condition_image_width, condition_image_height = (
image.width,
image.height,
)
batch.original_condition_image_size = image.size
self.preprocess_condition_image( self.preprocess_condition_image(
batch, server_args, condition_image_width, condition_image_height batch, server_args, condition_image_width, condition_image_height
@@ -454,72 +454,71 @@
"expected_median_denoise_ms": 102.03 "expected_median_denoise_ms": 102.03
}, },
"qwen_image_edit_ti2i": { "qwen_image_edit_ti2i": {
"notes": "single uploaded reference image, Qwen/Qwen-Image-Edit",
"expected_e2e_ms": 138500.0,
"expected_avg_denoise_ms": 720.0,
"expected_median_denoise_ms": 718.0,
"stages_ms": { "stages_ms": {
"InputValidationStage": 23, "InputValidationStage": 38.62,
"ImageEncodingStage": 1485.0, "ImageEncodingStage": 1174.26,
"ImageVAEEncodingStage": 400.0, "ImageVAEEncodingStage": 233.71,
"ConditioningStage": 0.13, "TimestepPreparationStage": 3.0,
"TimestepPreparationStage": 13.78, "LatentPreparationStage": 0.17,
"LatentPreparationStage": 15.0, "ConditioningStage": 0.01,
"DenoisingStage": 36000.0, "DenoisingStage": 42542.67,
"DecodingStage": 850.0 "DecodingStage": 508.39
}, },
"denoise_step_ms": { "denoise_step_ms": {
"0": 720.0, "0": 705.2,
"1": 720.0, "1": 854.4,
"2": 720.0, "2": 853.02,
"3": 720.0, "3": 852.77,
"4": 720.0, "4": 850.58,
"5": 720.0, "5": 851.46,
"6": 720.0, "6": 850.78,
"7": 720.0, "7": 851.67,
"8": 720.0, "8": 852.81,
"9": 720.0, "9": 853.98,
"10": 720.0, "10": 853.68,
"11": 720.0, "11": 852.62,
"12": 720.0, "12": 853.78,
"13": 720.0, "13": 854.42,
"14": 720.0, "14": 853.59,
"15": 720.0, "15": 853.28,
"16": 720.0, "16": 853.58,
"17": 720.0, "17": 854.02,
"18": 720.0, "18": 854.42,
"19": 720.0, "19": 854.46,
"20": 720.0, "20": 853.6,
"21": 720.0, "21": 854.18,
"22": 720.0, "22": 854.05,
"23": 720.0, "23": 854.45,
"24": 720.0, "24": 855.4,
"25": 720.0, "25": 851.82,
"26": 720.0, "26": 855.31,
"27": 720.0, "27": 854.42,
"28": 720.0, "28": 854.2,
"29": 720.0, "29": 854.43,
"30": 720.0, "30": 855.49,
"31": 720.0, "31": 854.51,
"32": 720.0, "32": 855.26,
"33": 720.0, "33": 852.42,
"34": 720.0, "34": 853.82,
"35": 720.0, "35": 856.22,
"36": 720.0, "36": 854.53,
"37": 720.0, "37": 854.44,
"38": 720.0, "38": 854.07,
"39": 720.0, "39": 852.74,
"40": 720.0, "40": 854.56,
"41": 720.0, "41": 854.24,
"42": 720.0, "42": 853.56,
"43": 720.0, "43": 854.74,
"44": 720.0, "44": 855.34,
"45": 720.0, "45": 853.93,
"46": 720.0, "46": 854.36,
"47": 720.0, "47": 852.65,
"48": 720.0, "48": 851.19,
"49": 720.0 "49": 851.89
} },
"expected_e2e_ms": 44503.12,
"expected_avg_denoise_ms": 850.73,
"expected_median_denoise_ms": 854.0
}, },
"wan2_1_t2v_1.3b": { "wan2_1_t2v_1.3b": {
"stages_ms": { "stages_ms": {
@@ -655,6 +654,63 @@
"expected_avg_denoise_ms": 2608.84, "expected_avg_denoise_ms": 2608.84,
"expected_median_denoise_ms": 2601.59 "expected_median_denoise_ms": 2601.59
}, },
"qwen_image_edit_2509_ti2i": {
"stages_ms": {
"InputValidationStage": 213.24,
"ImageEncodingStage": 1089.12,
"ImageVAEEncodingStage": 304.56,
"TimestepPreparationStage": 2.94,
"LatentPreparationStage": 0.2,
"ConditioningStage": 0.01,
"DenoisingStage": 50724.5,
"DecodingStage": 601.02
},
"denoise_step_ms": {
"0": 1057.09,
"1": 1267.06,
"2": 1268.33,
"3": 1268.94,
"4": 1270.36,
"5": 1270.44,
"6": 1268.61,
"7": 1270.21,
"8": 1274.98,
"9": 1271.57,
"10": 1273.15,
"11": 1271.56,
"12": 1272.69,
"13": 1271.62,
"14": 1274.04,
"15": 1276.81,
"16": 1272.2,
"17": 1269.33,
"18": 1275.96,
"19": 1274.43,
"20": 1272.57,
"21": 1275.28,
"22": 1273.63,
"23": 1275.06,
"24": 1277.39,
"25": 1277.27,
"26": 1274.74,
"27": 1273.38,
"28": 1276.77,
"29": 1275.59,
"30": 1275.51,
"31": 1274.9,
"32": 1274.8,
"33": 1279.03,
"34": 1272.9,
"35": 1274.67,
"36": 1272.61,
"37": 1272.82,
"38": 1276.41,
"39": 1273.55
},
"expected_e2e_ms": 52938.04,
"expected_avg_denoise_ms": 1267.96,
"expected_median_denoise_ms": 1273.46
},
"fastwan2_2_ti2v_5b": { "fastwan2_2_ti2v_5b": {
"stages_ms": { "stages_ms": {
"InputValidationStage": 88.86, "InputValidationStage": 88.86,
@@ -90,15 +90,25 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
and sampling_params.image_path and sampling_params.image_path
): ):
# Handle URL or local path # Handle URL or local path
if is_image_url(sampling_params.image_path): image_path_list = sampling_params.image_path
image_path = download_image_from_url(str(sampling_params.image_path)) if not isinstance(image_path_list, list):
else: image_path_list = [image_path_list]
image_path = Path(sampling_params.image_path)
new_image_path_list = []
for image_path in image_path_list:
if is_image_url(image_path):
new_image_path_list.append(download_image_from_url(str(image_path)))
else:
new_image_path_list.append(Path(image_path))
if not image_path.exists():
pytest.skip(f"{case.id}: file missing: {image_path}")
image_path_list = new_image_path_list
warmup.run_edit_warmups( warmup.run_edit_warmups(
count=server_args.warmup_edit, count=server_args.warmup_edit,
edit_prompt=sampling_params.prompt, edit_prompt=sampling_params.prompt,
image_path=image_path, image_path=image_path_list,
) )
except Exception as exc: except Exception as exc:
logger.error("Warm-up failed for %s: %s", case.id, exc) logger.error("Warm-up failed for %s: %s", case.id, exc)
@@ -278,23 +278,31 @@ class WarmupRunner:
if count <= 0: if count <= 0:
return return
if not image_path.exists(): if not isinstance(image_path, list):
logger.warning( image_path = [image_path]
"[server-test] Skipping edit warmup: image missing at %s", image_path
) for image in image_path:
return if not image.exists():
logger.warning(
"[server-test] Skipping edit warmup: image missing at %s", image
)
return
logger.info("[server-test] Running %s edit warm-up(s)", count) logger.info("[server-test] Running %s edit warm-up(s)", count)
for _ in range(count): for _ in range(count):
with image_path.open("rb") as fh: images = [open(image, "rb") for image in image_path]
try:
result = self.client.images.edit( result = self.client.images.edit(
model=self.model, model=self.model,
image=fh, image=images,
prompt=edit_prompt, prompt=edit_prompt,
n=1, n=1,
size=self.output_size, size=self.output_size,
response_format="b64_json", response_format="b64_json",
) )
finally:
for img in images:
img.close()
validate_image(result.data[0].b64_json) validate_image(result.data[0].b64_json)
@@ -585,22 +593,36 @@ def get_generate_fn(
if not sampling_params.prompt or not sampling_params.image_path: if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config") pytest.skip(f"{id}: no edit config")
if is_image_url(sampling_params.image_path): image_paths = sampling_params.image_path
image_path = download_image_from_url(str(sampling_params.image_path))
else:
image_path = Path(sampling_params.image_path)
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
with image_path.open("rb") as fh: if not isinstance(image_paths, list):
image_paths = [image_paths]
new_image_paths = []
for image_path in image_paths:
if is_image_url(image_path):
new_image_paths.append(download_image_from_url(str(image_path)))
else:
new_image_paths.append(Path(image_path))
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
image_paths = new_image_paths
images = [open(image_path, "rb") for image_path in image_paths]
try:
response = client.images.with_raw_response.edit( response = client.images.with_raw_response.edit(
model=model_path, model=model_path,
image=fh, image=images,
prompt=sampling_params.prompt, prompt=sampling_params.prompt,
n=1, n=1,
size=sampling_params.output_size, size=sampling_params.output_size,
response_format="b64_json", response_format="b64_json",
) )
finally:
for img in images:
img.close()
rid = response.headers.get("x-request-id", "") rid = response.headers.get("x-request-id", "")
result = response.parse() result = response.parse()
@@ -226,6 +226,14 @@ TI2I_sampling_params = DiffusionSamplingParams(
image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg", image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
) )
MULTI_IMAGE_TI2I_sampling_params = DiffusionSamplingParams(
prompt="The magician bear is on the left, the alchemist bear is on the right, facing each other in the central park square.",
image_path=[
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_1.jpg",
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_2.jpg",
],
)
T2V_PROMPT = "A curious raccoon" T2V_PROMPT = "A curious raccoon"
TI2V_sampling_params = DiffusionSamplingParams( TI2V_sampling_params = DiffusionSamplingParams(
@@ -288,6 +296,16 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
), ),
TI2I_sampling_params, TI2I_sampling_params,
), ),
DiffusionTestCase(
"qwen_image_edit_2509_ti2i",
DiffusionServerArgs(
model_path="Qwen/Qwen-Image-Edit-2509",
modality="image",
warmup_text=0,
warmup_edit=1,
),
MULTI_IMAGE_TI2I_sampling_params,
),
] ]
ONE_GPU_CASES_B: list[DiffusionTestCase] = [ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
@@ -6,6 +6,7 @@ import logging
import os import os
import tempfile import tempfile
from datetime import datetime from datetime import datetime
from typing import List, Union
from urllib.parse import urlparse from urllib.parse import urlparse
from urllib.request import urlopen from urllib.request import urlopen
@@ -100,9 +101,9 @@ def upload_file_to_slack(
model: str = None, model: str = None,
prompt: str = None, prompt: str = None,
file_path: str = None, file_path: str = None,
origin_file_path: str = None, origin_file_path: Union[str, List[str]] = None,
) -> bool: ) -> bool:
temp_path = None temp_paths = []
try: try:
from slack_sdk import WebClient from slack_sdk import WebClient
@@ -117,17 +118,39 @@ def upload_file_to_slack(
logger.info(f"Slack upload failed: no file path") logger.info(f"Slack upload failed: no file path")
return False return False
if origin_file_path and origin_file_path.startswith(("http", "https")): origin_paths = []
suffix = os.path.splitext(urlparse(origin_file_path).path)[1] or ".tmp" if isinstance(origin_file_path, str):
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf: if origin_file_path:
with urlopen(origin_file_path) as response: origin_paths.append(origin_file_path)
tf.write(response.read()) elif isinstance(origin_file_path, list):
temp_path = tf.name origin_paths = [p for p in origin_file_path if p]
origin_file_path = temp_path
uploads = [{"file": file_path, "title": "Generated Image"}] final_origin_paths = []
if origin_file_path and os.path.exists(origin_file_path): for path in origin_paths:
uploads.insert(0, {"file": origin_file_path, "title": "Original Image"}) if path.startswith(("http", "https")):
try:
suffix = os.path.splitext(urlparse(path).path)[1] or ".tmp"
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf:
with urlopen(path) as response:
tf.write(response.read())
temp_paths.append(tf.name)
final_origin_paths.append(tf.name)
except Exception as e:
logger.warning(f"Failed to download {path}: {e}")
else:
final_origin_paths.append(path)
uploads = []
for i, path in enumerate(final_origin_paths):
if os.path.exists(path):
title = (
"Original Image"
if len(final_origin_paths) == 1
else f"Original Image {i+1}"
)
uploads.append({"file": path, "title": title})
uploads.append({"file": file_path, "title": "Generated Image"})
message = ( message = (
f"*Case ID:* `{case_id}`\n" f"*Model:* `{model}`\n" f"*Prompt:* {prompt}" f"*Case ID:* `{case_id}`\n" f"*Model:* `{model}`\n" f"*Prompt:* {prompt}"
@@ -189,5 +212,6 @@ def upload_file_to_slack(
logger.info(f"Slack upload failed: {e}") logger.info(f"Slack upload failed: {e}")
return False return False
finally: finally:
if temp_path and os.path.exists(temp_path): for p in temp_paths:
os.remove(temp_path) if os.path.exists(p):
os.remove(p)