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

[diffusion] feat: support multi-frame image output (#15878)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
chhnb
2025-12-31 13:50:02 +08:00
committed by GitHub
parent db6b51a838
commit 3619ec61b4
7 changed files with 157 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -500,31 +500,25 @@ class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
img_shapes = batch.img_shapes
txt_seq_lens = batch.txt_seq_lens
(img_cos, img_sin), (txt_cos, txt_sin) = (
QwenImageEditPlusPipelineConfig.get_freqs_cis(
img_shapes, txt_seq_lens, rotary_emb, device, dtype
)
freqs_cis = 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
)
img_cache, txt_cache = freqs_cis
noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
img_cache = torch.cat(
[noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
).to(device=device)
return {
"txt_seq_lens": txt_seq_lens,
"img_shapes": img_shapes,
"freqs_cis": ((img_cos, img_sin), (txt_cos, txt_sin)),
"freqs_cis": (img_cache, txt_cache),
"additional_t_cond": torch.tensor([0], device=device, dtype=torch.long),
}

28
python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py
View File

@@ -63,6 +63,7 @@ def _build_sampling_params_from_request(
guidance_scale: Optional[float] = None,
num_inference_steps: Optional[int] = None,
enable_teacache: Optional[bool] = None,
num_frames: int = 1,
) -> SamplingParams:
if size is None:
width, height = None, None
@@ -76,7 +77,7 @@ def _build_sampling_params_from_request(
request_id=request_id,
prompt=prompt,
image_path=image_path,
num_frames=1, # image
num_frames=num_frames, # image
width=width,
height=height,
num_outputs_per_prompt=max(1, min(int(n or 1), 10)),
@@ -121,9 +122,12 @@ async def generations(
server_args=get_global_server_args(),
sampling_params=sampling,
)
save_file_path, result = await process_generation_batch(
# Run synchronously for images and save to disk
save_file_path_list, result = await process_generation_batch(
async_scheduler_client, batch
)
save_file_path = save_file_path_list[0]
await IMAGE_STORE.upsert(
request_id,
@@ -178,6 +182,7 @@ async def edits(
guidance_scale: Optional[float] = Form(None),
num_inference_steps: Optional[int] = Form(None),
enable_teacache: Optional[bool] = Form(False),
num_frames: int = Form(1),
):
request_id = generate_request_id()
# Resolve images from either `image` or `image[]` (OpenAI SDK sends `image[]` when list is provided)
@@ -221,13 +226,14 @@ async def edits(
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
enable_teacache=enable_teacache,
num_frames=num_frames, # image
)
batch = _build_req_from_sampling(sampling)
save_file_path, result = await process_generation_batch(
save_file_path_list, result = await process_generation_batch(
async_scheduler_client, batch
)
save_file_path = save_file_path_list[0]
await IMAGE_STORE.upsert(
request_id,
{
@@ -241,11 +247,15 @@ async def edits(
# Default to b64_json to align with gpt-image-1 behavior in OpenAI examples
if (response_format or "b64_json").lower() == "b64_json":
with open(save_file_path, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
response_kwargs = {
"data": [ImageResponseData(b64_json=b64, revised_prompt=prompt)],
}
response_kwargs = {"data": []}
for save_file_path in save_file_path_list:
with open(save_file_path, "rb") as f:
b64 = base64.b64encode(f.read()).decode("utf-8")
response_kwargs["data"].append(
ImageResponseData(b64_json=b64, revised_prompt=prompt)
)
if result.peak_memory_mb and result.peak_memory_mb > 0:
response_kwargs["peak_memory_mb"] = result.peak_memory_mb
else:
url = f"/v1/images/{request_id}/content"
response_kwargs = {

41
python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
View File

@@ -9,6 +9,7 @@ from typing import Any, List, Optional, Union
import httpx
from fastapi import UploadFile
from sglang.multimodal_gen.configs.sample.sampling_params import DataType
from sglang.multimodal_gen.runtime.entrypoints.utils import post_process_sample
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient
@@ -183,15 +184,35 @@ async def process_generation_batch(
raise RuntimeError(
f"Model generation returned no output. Error from scheduler: {error_msg}"
)
save_file_path = str(os.path.join(batch.output_path, batch.output_file_name))
post_process_sample(
result.output[0],
batch.data_type,
batch.fps,
batch.save_output,
save_file_path,
)
save_file_path_list = []
if batch.data_type == DataType.VIDEO:
for idx, output in enumerate(result.output):
save_file_path = str(
os.path.join(batch.output_path, batch.output_file_name)
)
post_process_sample(
result.output[idx],
batch.data_type,
batch.fps,
batch.save_output,
save_file_path,
)
save_file_path_list.append(save_file_path)
else:
for idx, output in enumerate(result.output):
save_file_path = str(
os.path.join(
batch.output_path, f"sample_{idx}_" + batch.output_file_name
)
)
post_process_sample(
output,
batch.data_type,
batch.fps,
batch.save_output,
save_file_path,
)
save_file_path_list.append(save_file_path)
total_time = time.perf_counter() - total_start_time
log_batch_completion(logger, 1, total_time)
@@ -199,7 +220,7 @@ async def process_generation_batch(
if result.peak_memory_mb and result.peak_memory_mb > 0:
logger.info(f"Peak memory usage: {result.peak_memory_mb:.2f} MB")
return save_file_path, result
return save_file_path_list, result
def merge_image_input_list(*inputs: Union[List, Any, None]) -> List:

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

@@ -788,6 +788,69 @@
"expected_avg_denoise_ms": 1267.96,
"expected_median_denoise_ms": 1273.46
},
"qwen_image_layered_i2i": {
"stages_ms": {
"QwenImageLayeredBeforeDenoisingStage": 3240.24,
"TimestepPreparationStage": 3.14,
"DenoisingStage": 41551.18,
"DecodingStage": 312.93
},
"denoise_step_ms": {
"0": 809.93,
"1": 836.69,
"2": 834.98,
"3": 826.84,
"4": 827.15,
"5": 827.28,
"6": 830.97,
"7": 827.7,
"8": 829.4,
"9": 832.14,
"10": 825.99,
"11": 831.65,
"12": 829.31,
"13": 829.46,
"14": 828.33,
"15": 831.14,
"16": 830.44,
"17": 831.6,
"18": 829.18,
"19": 831.64,
"20": 828.21,
"21": 831.02,
"22": 831.39,
"23": 830.16,
"24": 832.21,
"25": 831.04,
"26": 830.48,
"27": 831.88,
"28": 833.5,
"29": 837.31,
"30": 828.16,
"31": 832.24,
"32": 833.56,
"33": 829.08,
"34": 833.11,
"35": 831.07,
"36": 832.71,
"37": 833.12,
"38": 830.65,
"39": 831.59,
"40": 833.24,
"41": 831.92,
"42": 832.77,
"43": 830.88,
"44": 833.75,
"45": 831.29,
"46": 834.48,
"47": 832.6,
"48": 835.24,
"49": 832.49
},
"expected_e2e_ms": 45109.63,
"expected_avg_denoise_ms": 830.86,
"expected_median_denoise_ms": 831.34
},
"fastwan2_2_ti2v_5b": {
"stages_ms": {
"InputValidationStage": 88.86,

1
python/sglang/multimodal_gen/test/server/test_server_common.py
View File

@@ -103,6 +103,7 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
model=server_args.model_path,
prompt=sampling_params.prompt or "A colorful raccoon icon",
output_size=output_size,
output_format=sampling_params.output_format,
)
warmup.run_text_warmups(server_args.warmup_text)

17
python/sglang/multimodal_gen/test/server/test_server_utils.py
View File

@@ -455,6 +455,7 @@ class WarmupRunner:
model: str,
prompt: str,
output_size: str,
output_format: str = None,
):
self.client = OpenAI(
api_key="sglang-anything",
@@ -463,6 +464,7 @@ class WarmupRunner:
self.model = model
self.prompt = prompt
self.output_size = output_size
self.output_format = output_format
def run_text_warmups(self, count: int) -> None:
"""Run text-to-image warmup requests."""
@@ -511,6 +513,7 @@ class WarmupRunner:
n=1,
size=self.output_size,
response_format="b64_json",
output_format=self.output_format,
)
finally:
for img in images:
@@ -898,7 +901,9 @@ def get_generate_fn(
image_paths = new_image_paths
# Request parameters that affect output format
req_output_format = None # Not specified in current request
req_output_format = (
sampling_params.output_format
) # Not specified in current request
req_background = None # Not specified in current request
images = [open(image_path, "rb") for image_path in image_paths]
@@ -910,6 +915,8 @@ def get_generate_fn(
n=n,
size=output_size,
response_format="b64_json",
output_format=req_output_format,
extra_body={"num_frames": sampling_params.num_frames},
)
finally:
for img in images:
@@ -954,7 +961,6 @@ def get_generate_fn(
"""TI2I: Text + Image ? Image edit using direct URL transfer (no pre-download)."""
if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config")
# Handle both single URL and list of URLs
image_urls = sampling_params.image_path
if not isinstance(image_urls, list):
@@ -968,7 +974,9 @@ def get_generate_fn(
)
# Request parameters that affect output format
req_output_format = None # Not specified in current request
req_output_format = (
sampling_params.output_format
) # Not specified in current request
req_background = None # Not specified in current request
response = client.images.with_raw_response.edit(
@@ -978,7 +986,8 @@ def get_generate_fn(
n=n,
size=sampling_params.output_size,
response_format="b64_json",
extra_body={"url": image_urls},
output_format=req_output_format,
extra_body={"url": image_urls, "num_frames": sampling_params.num_frames},
)
result = response.parse()

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

@@ -173,6 +173,9 @@ class DiffusionSamplingParams:
# URL direct test flag - if True, don't pre-download URL images
direct_url_test: bool = False
# output format
output_format: str | None = None # "png", "jpeg", "mp4", etc.
num_outputs_per_prompt: int = 1
@@ -269,6 +272,15 @@ MULTI_IMAGE_TI2I_sampling_params = DiffusionSamplingParams(
],
direct_url_test=True,
)
MULTI_FRAME_I2I_sampling_params = DiffusionSamplingParams(
prompt="a high quality, cute halloween themed illustration, consistent style and lighting",
image_path=[
"https://raw.githubusercontent.com/QwenLM/Qwen-Image-Layered/main/assets/test_images/4.png"
],
num_frames=4,
direct_url_test=True,
output_format="png",
)
T2V_PROMPT = "A curious raccoon"
@@ -356,6 +368,16 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
),
MULTI_IMAGE_TI2I_sampling_params,
),
DiffusionTestCase(
"qwen_image_layered_i2i",
DiffusionServerArgs(
model_path="Qwen/Qwen-Image-Layered",
modality="image",
warmup_text=0,
warmup_edit=1,
),
MULTI_FRAME_I2I_sampling_params,
),
]
ONE_GPU_CASES_B: list[DiffusionTestCase] = [