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

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -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) = (
+        freqs_cis = QwenImageEditPlusPipelineConfig.get_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_cache, txt_cache = freqs_cis
-        img_cos = torch.cat([noisy_img_cos, img_cos[noisy_img_seq_len:, :]], dim=0).to(
+        noisy_img_cache = shard_rotary_emb_for_sp(img_cache[:noisy_img_seq_len, :])
-            device=device
+        img_cache = torch.cat(
-        )
+            [noisy_img_cache, img_cache[noisy_img_seq_len:, :]], dim=0
-        img_sin = torch.cat([noisy_img_sin, img_sin[noisy_img_seq_len:, :]], dim=0).to(
+        ).to(device=device)
-            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),
         }
 

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -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":
+        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 = {
+                response_kwargs["data"].append(
-            "data": [ImageResponseData(b64_json=b64, revised_prompt=prompt)],
+                    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 = {

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -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_list = []
-        save_file_path = str(os.path.join(batch.output_path, batch.output_file_name))
+        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[0],
+                    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:

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -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,

python/sglang/multimodal_gen/test/server/test_server_common.py

@@ -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)
 

python/sglang/multimodal_gen/test/server/test_server_utils.py

@@ -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()

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -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] = [