[Frontend][Model] Support batch request with refined OmniDiffusionReq… (#797)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

examples/offline_inference/bagel/end2end.py

@@ -1,5 +1,8 @@
 import argparse
 import os
+from typing import cast
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
 
 def parse_args():
@@ -50,21 +53,24 @@ def parse_args():
 def main():
     args = parse_args()
     model_name = args.model
+    prompts: list[OmniPromptType] = []
     try:
         # Preferred: load from txt file (one prompt per line)
         if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
             with open(args.txt_prompts, encoding="utf-8") as f:
                 lines = [ln.strip() for ln in f.readlines()]
-            args.prompts = [ln for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+            prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+        else:
+            prompts = args.prompts
     except Exception as e:
         print(f"[Error] Failed to load prompts: {e}")
         raise
 
-    if args.prompts is None:
+    if not prompts:
         # Default prompt for text2img test if none provided
-        args.prompts = ["<|im_start|>A cute cat<|im_end|>"]
-        print(f"[Info] No prompts provided, using default: {args.prompts}")
+        prompts = ["<|im_start|>A cute cat<|im_end|>"]
+        print(f"[Info] No prompts provided, using default: {prompts}")
     omni_outputs = []
 
     from PIL import Image
@@ -77,21 +83,27 @@ def main():
         print("[Info] Running in img2img mode (Stage 1 only)")
         client = OmniDiffusion(model=model_name)
 
-        generate_kwargs = {
-            "prompt": args.prompts,
-            "seed": 52,
-            "need_kv_receive": False,
-            "num_inference_steps": args.steps,
-        }
-
         if args.image_path:
             if os.path.exists(args.image_path):
                 loaded_image = Image.open(args.image_path).convert("RGB")
-                generate_kwargs["pil_image"] = loaded_image
+                prompts = [
+                    {
+                        "prompt": cast(str, p),
+                        "multi_modal_data": {"image": loaded_image},
+                    }
+                    for p in prompts
+                ]
             else:
                 print(f"[Warning] Image path {args.image_path} does not exist.")
 
-        result = client.generate(**generate_kwargs)
+        result = client.generate(
+            prompts,
+            OmniDiffusionSamplingParams(
+                seed=52,
+                need_kv_receive=False,
+                num_inference_steps=args.steps,
+            ),
+        )
 
         # Ensure result is a list for iteration
         if not isinstance(result, list):
@@ -100,8 +112,6 @@ def main():
             omni_outputs = result
 
     else:
-        import copy
-
         from vllm_omni.entrypoints.omni import Omni
 
         omni_kwargs = {}
@@ -144,11 +154,11 @@ def main():
                 prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
                 formatted_prompts.append(prompt_dict)
 
-        params_list = copy.deepcopy(omni.default_sampling_params_list)
+        params_list = omni.default_sampling_params_list
         if args.modality == "text2img":
-            params_list[0]["max_tokens"] = 1
+            params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
             if len(params_list) > 1:
-                params_list[1]["num_inference_steps"] = args.steps
+                params_list[1].num_inference_steps = args.steps  # type: ignore # The second stage is an OmniDiffusionSamplingParam
 
         omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 

examples/offline_inference/image_to_image/image_edit.py

@@ -79,6 +79,7 @@ from PIL import Image
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -374,25 +375,28 @@ def main():
     print(f"{'=' * 60}\n")
 
     generation_start = time.perf_counter()
-    # Generate edited image
-    generate_kwargs = {
-        "prompt": args.prompt,
-        "pil_image": input_image,
-        "negative_prompt": args.negative_prompt,
-        "generator": generator,
-        "true_cfg_scale": args.cfg_scale,
-        "guidance_scale": args.guidance_scale,
-        "num_inference_steps": args.num_inference_steps,
-        "num_outputs_per_prompt": args.num_outputs_per_prompt,
-        "layers": args.layers,
-        "resolution": args.resolution,
-    }
 
     if profiler_enabled:
         print("[Profiler] Starting profiling...")
         omni.start_profile()
 
-    outputs = omni.generate(**generate_kwargs)
+    # Generate edited image
+    outputs = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": input_image},
+        },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_outputs_per_prompt,
+            layers=args.layers,
+            resolution=args.resolution,
+        ),
+    )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -27,6 +27,7 @@ import PIL.Image
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -119,16 +120,20 @@ def main():
 
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        pil_image=image,
-        height=height,
-        width=width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        guidance_scale_2=args.guidance_scale_high,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": image},
+        },
+        OmniDiffusionSamplingParams(
+            height=height,
+            width=width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_high,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
     )
 
     # Extract video frames from OmniRequestOutput

examples/offline_inference/lora_inference/lora_inference.py

@@ -5,6 +5,7 @@ import argparse
 from pathlib import Path
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 
@@ -95,18 +96,17 @@ def main():
         )
         print(f"Activating pre-loaded LoRA: id={lora_request_id}, scale={args.lora_scale}")
 
-    gen_kwargs = {
-        "prompt": args.prompt,
-        "height": args.height,
-        "width": args.width,
-        "num_inference_steps": args.num_inference_steps,
-    }
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+    )
 
     if lora_request:
-        gen_kwargs["lora_request"] = lora_request
-        gen_kwargs["lora_scale"] = args.lora_scale
+        sampling_params.lora_request = lora_request
+        sampling_params.lora_scale = args.lora_scale
 
-    outputs = omni.generate(**gen_kwargs)
+    outputs = omni.generate(args.prompt, sampling_params)
 
     if not outputs or len(outputs) == 0:
         raise ValueError("No output generated from omni.generate()")

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -21,6 +21,7 @@ import numpy as np
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.utils.platform_utils import detect_device_type
 
 
@@ -143,16 +144,20 @@ def main():
 
     # Generate audio
     outputs = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_waveforms,
-        extra={
-            "audio_start_in_s": args.audio_start,
-            "audio_end_in_s": audio_end_in_s,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
         },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_waveforms,
+            extra_args={
+                "audio_start_in_s": args.audio_start,
+                "audio_end_in_s": audio_end_in_s,
+            },
+        ),
     )
 
     generation_end = time.perf_counter()

examples/offline_inference/text_to_image/README.md

@@ -20,7 +20,7 @@ if __name__ == "__main__":
     images[0].save("coffee.png")
 ```
 
-Or put more than one prompt in a request, processing them sequentially.
+Or put more than one prompt in a request.
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
@@ -37,6 +37,41 @@ if __name__ == "__main__":
       image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
+Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
+This is useful for models that support negative prompts.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Qwen/Qwen-Image")
+    outputs = omni.generate([
+      {
+        "prompt": "a cup of coffee on a table"，
+        "negative_prompt": "low resolution"
+      },
+      {
+        "prompt": "a toy dinosaur on a sandy beach"，
+        "negative_prompt": "cinematic, realistic"
+      }
+    ])
+    for i, output in enumerate(outputs):
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
+```
+
 ## Local CLI Usage
 
 ```bash

examples/offline_inference/text_to_image/gradio_demo.py

@@ -5,6 +5,7 @@ import gradio as gr
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -101,12 +102,14 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
         generator = torch.Generator(device=device).manual_seed(seed)
         outputs = omni.generate(
             prompt.strip(),
-            height=height,
-            width=width,
-            generator=generator,
-            true_cfg_scale=float(cfg_scale_value),
-            num_inference_steps=num_steps,
-            num_outputs_per_prompt=num_images,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                generator=generator,
+                true_cfg_scale=float(cfg_scale_value),
+                num_inference_steps=num_steps,
+                num_outputs_per_prompt=num_images,
+            ),
         )
         images_outputs = []
         for output in outputs:

examples/offline_inference/text_to_image/text_to_image.py

@@ -10,6 +10,7 @@ import torch
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -189,15 +190,19 @@ def main():
 
     generation_start = time.perf_counter()
     outputs = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        true_cfg_scale=args.cfg_scale,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_images_per_prompt,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_images_per_prompt,
+        ),
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start

examples/offline_inference/text_to_video/text_to_video.py

@@ -11,6 +11,7 @@ import torch
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
@@ -108,15 +109,19 @@ def main():
 
     generation_start = time.perf_counter()
     frames = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        guidance_scale_2=args.guidance_scale_high,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_high,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start