[Frontend] ComfyUI video & LoRA support (#1596)

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

apps/ComfyUI-vLLM-Omni/README.md

@@ -47,6 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
 
 - **Generate Image** for text-to-image and image-to-image tasks
+- **Generate Video** for text-to-video and image-to-video tasks
 - **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks
 
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
 
 - Drag a generation node onto the canvas.
 - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
-- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
 - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
     - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
     - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
@@ -104,6 +105,20 @@ You can configure per-stage sampling parameters for multi-stage models.
 > [!TIP]
 > The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
 
+### Text-to-video and image-to-video generation (e.g., Wan)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
+    <img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
+
 ### TTS (e.g., Qwen TTS series)
 
 (Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -14,17 +14,21 @@ from .comfyui_vllm_omni.nodes import (
     VLLMOmniARSampling,
     VLLMOmniDiffusionSampling,
     VLLMOmniGenerateImage,
+    VLLMOmniGenerateVideo,
     VLLMOmniQwenTTSParams,
+    VLLMOmniRemoteLoRA,
     VLLMOmniSamplingParamsList,
     VLLMOmniTTS,
     VLLMOmniUnderstanding,
     VLLMOmniVoiceClone,
+    VLLMOmniWanParams,
 )
 
 # A dictionary that contains all nodes you want to export with their names
 NODE_CLASS_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
+    "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
     "VLLMOmniUnderstanding": VLLMOmniUnderstanding,
     "VLLMOmniTTS": VLLMOmniTTS,
     "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
@@ -32,13 +36,16 @@ NODE_CLASS_MAPPINGS = {
     "VLLMOmniARSampling": VLLMOmniARSampling,
     "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
     "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
+    "VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
     "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
+    "VLLMOmniWanParams": VLLMOmniWanParams,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": "Generate Image",
+    "VLLMOmniGenerateVideo": "Generate Video",
     "VLLMOmniUnderstanding": "Multimodality Understanding",
     "VLLMOmniTTS": "TTS (Text to Speech)",
     "VLLMOmniVoiceClone": "TTS Voice Cloning",
@@ -46,7 +53,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "VLLMOmniARSampling": "AR Sampling Params",
     "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
     "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
+    "VLLMOmniRemoteLoRA": "LoRA",
     "VLLMOmniQwenTTSParams": "Qwen TTS Params",
+    "VLLMOmniWanParams": "Wan Video Params",
 }
 
 WEB_DIRECTORY = "./web"

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -11,6 +11,7 @@ from .utils.types import (
     AutoregressionSamplingParams,
     DiffusionSamplingParams,
     QwenTTSModelSpecificParams,
+    WanModelSpecificParams,
 )
 from .utils.validators import (
     add_sampling_parameters_to_stage,
@@ -56,6 +57,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                 # "video": ("VIDEO",),
                 # "audio": ("AUDIO",),
                 "sampling_params": ("SAMPLING_PARAMS",),
+                "lora": ("REMOTE_LORA",),
             },
         }
 
@@ -76,9 +78,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
         audio: AudioInput | None = None,  # Hidden & unused
         video: VideoInput | None = None,  # Hidden & unused
         sampling_params: dict | list[dict] | None = None,
+        lora: dict | None = None,
         **kwargs,
     ):
-        logger.info("Uncaught kwargs: %s", kwargs)
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
         logger.debug("Got sampling params: %s", sampling_params)
         validate_model_and_sampling_params_types(model, sampling_params)
         if image is None and mask is not None:
@@ -105,6 +109,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                     height=height,
                     negative_prompt=negative_prompt,
                     sampling_params=sampling_params,
+                    lora=lora,
                 )
                 return (output,)
             elif image is not None and audio is None and video is None:
@@ -119,6 +124,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                     negative_prompt=negative_prompt,
                     mask=mask,
                     sampling_params=sampling_params,
+                    lora=lora,
                 )
                 return (output,)
 
@@ -136,11 +142,91 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
             audio=audio,
             video=video,
             sampling_params=sampling_params,
+            lora=lora,
         )
 
         return (output,)
 
 
+class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "width": ("INT", {"default": 832, "min": 1}),
+                "height": ("INT", {"default": 480, "min": 1}),
+                "fps": ("INT", {"default": 16, "min": 1}),
+                "num_frames": ("INT", {"default": 41, "min": 1}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+                "lora": ("REMOTE_LORA",),
+                "model_params": ("VIDEO_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    RETURN_NAMES = ("video",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        fps: int,
+        num_frames: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        model_params: dict | None = None,
+        lora: dict | None = None,
+        **kwargs,
+    ):
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        logger.debug("Got model params: %s", model_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+
+        # Currently, all video generation models are single-stage diffusion models
+        if isinstance(sampling_params, list):
+            if len(sampling_params) != 1:
+                raise ValueError(
+                    "Video generation expects a single sampling params group. "
+                    "Please provide one Diffusion sampling node."
+                )
+            sampling_params = sampling_params[0]
+
+        if sampling_params is not None:
+            sampling_params.pop("type", None)  # internal fields
+        if model_params is not None:
+            model_params.pop("type", None)  # internal fields
+
+        client = VLLMOmniClient(url)
+        output = await client.generate_video(
+            model=model,
+            prompt=prompt,
+            image=image,  # image present => i2v, absent => t2v
+            width=width,
+            height=height,
+            num_frames=num_frames,
+            fps=fps,
+            negative_prompt=negative_prompt,
+            sampling_params=sampling_params,
+            lora=lora,
+            model_params=model_params,
+        )
+        return (output,)
+
+
 class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
     @classmethod
     def INPUT_TYPES(cls):
@@ -188,7 +274,8 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
         use_audio_in_video: bool = True,
         **kwargs,
     ) -> tuple[str, AudioInput]:
-        logger.info("Uncaught kwargs: %s", kwargs)
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
         logger.debug("Got sampling params: %s", sampling_params)
         validate_model_and_sampling_params_types(model, sampling_params)
 
@@ -547,6 +634,52 @@ class VLLMOmniSamplingParamsList:
         return (params,)
 
 
+class VLLMOmniRemoteLoRA:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "local_path": ("STRING", {"default": ""}),
+                "name": ("STRING", {"default": ""}),
+                "scale": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
+                ),
+                "int_id": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "step": 1,
+                        "tooltip": "0 means it is not set and the server can derive it.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("REMOTE_LORA",)
+    RETURN_NAMES = ("lora",)
+    FUNCTION = "get_lora"
+    CATEGORY = "vLLM-Omni"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
+        if not local_path.strip() or not name.strip():
+            return "Both local_path and name must be provided."
+        return True
+
+    def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
+        local_path = local_path.strip()
+        name = name.strip()
+        lora = {
+            "local_path": local_path or None,
+            "name": name or None,
+            "scale": float(scale),
+            "int_id": int(int_id) if int_id > 0 else None,
+        }
+        return (lora,)
+
+
 class VLLMOmniQwenTTSParams:
     @classmethod
     def INPUT_TYPES(cls):
@@ -572,3 +705,32 @@ class VLLMOmniQwenTTSParams:
 
     def get_params(self, **kwargs):
         return (QwenTTSModelSpecificParams(kwargs),)
+
+
+class VLLMOmniWanParams:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "guidance_scale_2": (
+                    "FLOAT",
+                    {"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
+                ),
+                "boundary_ratio": (
+                    "FLOAT",
+                    {"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
+                ),
+                "flow_shift": (
+                    "FLOAT",
+                    {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("VIDEO_PARAMS",)
+    RETURN_NAMES = ("Wan video params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Video Params"
+
+    def get_params(self, **kwargs):
+        return (WanModelSpecificParams(kwargs),)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -6,6 +6,7 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
 Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
 """
 
+import json
 from typing import Any
 
 import aiohttp
@@ -17,6 +18,7 @@ from .format import (
     audio_to_base64,
     base64_to_audio,
     base64_to_image_tensor,
+    base64_to_video,
     bytes_to_audio,
     image_tensor_to_base64,
     image_tensor_to_png_bytes,
@@ -30,7 +32,7 @@ logger = get_logger(__name__)
 
 
 class VLLMOmniClient:
-    def __init__(self, base_url: str, timeout: float = 300.0):
+    def __init__(self, base_url: str, timeout: float | None = None):
         self.base_url = base_url
         self.timeout = aiohttp.ClientTimeout(total=timeout)
 
@@ -43,12 +45,13 @@ class VLLMOmniClient:
         height: int,
         negative_prompt: str | None = None,
         sampling_params: dict | None = None,
+        lora: dict | None = None,
     ) -> torch.Tensor:
         """Run text-to-image generation via DALLE API"""
         await self._check_model_exist(model)
 
         size = f"{width}x{height}"
-        payload = {
+        payload: dict[str, Any] = {
             "model": model,
             "prompt": prompt,
             "size": size,
@@ -58,6 +61,8 @@ class VLLMOmniClient:
             payload["negative_prompt"] = negative_prompt
         if sampling_params is not None:
             payload.update(sampling_params)
+        if lora is not None:
+            payload["lora"] = lora
         logger.debug("img gen payload: %s", payload)
 
         url = self.base_url + "/images/generations"
@@ -110,6 +115,7 @@ class VLLMOmniClient:
         negative_prompt: str | None = None,
         mask: torch.Tensor | None = None,
         sampling_params: dict | None = None,
+        lora: dict | None = None,
     ) -> torch.Tensor:
         """Run image editing via DALLE API"""
         await self._check_model_exist(model)
@@ -131,6 +137,8 @@ class VLLMOmniClient:
         if sampling_params is not None:
             for k, v in sampling_params.items():
                 form.add_field(k, str(v))
+        if lora is not None:
+            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
         if mask is not None:
             mask_filename = "mask.png"
             form.add_field(
@@ -183,6 +191,7 @@ class VLLMOmniClient:
         audio: AudioInput | None = None,
         video: VideoInput | None = None,
         sampling_params: dict | list[dict] | None = None,
+        lora: dict | None = None,
     ) -> torch.Tensor:
         payload = VLLMOmniClient._prepare_chat_completion_messages(
             model=model,
@@ -193,6 +202,8 @@ class VLLMOmniClient:
             video=video,
             sampling_params=sampling_params,
             modalities=["image"],
+            # === below are additional `extra_body` fields, handled by **kwargs ===
+            lora=lora,
         )
         choices = await self._generate_base_chat_completion(model, payload)
 
@@ -206,6 +217,79 @@ class VLLMOmniClient:
 
         return torch.stack(image_tensors, dim=0)
 
+    async def generate_video(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        num_frames: int,
+        fps: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        sampling_params: dict | None = None,
+        model_params: dict | None = None,
+        lora: dict | None = None,
+        **extra_body,
+    ) -> VideoInput:
+        form = aiohttp.FormData()
+        form.add_field("model", model)
+        form.add_field("prompt", prompt)
+        form.add_field("width", str(width))
+        form.add_field("height", str(height))
+        form.add_field("num_frames", str(num_frames))
+        form.add_field("fps", str(fps))
+        if negative_prompt:
+            form.add_field("negative_prompt", negative_prompt)
+        if sampling_params is not None:
+            for k, v in sampling_params.items():
+                form.add_field(k, str(v))
+        if model_params is not None:
+            for k, v in model_params.items():
+                form.add_field(k, str(v))
+        if lora is not None:
+            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
+        if extra_body:
+            form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
+
+        if image is not None:
+            image_filename = "image.png"  # Required for multipart form
+            form.add_field(
+                "input_reference",
+                image_tensor_to_png_bytes(image, image_filename),
+                filename=image_filename,
+                content_type="image/png",
+            )
+
+        url = self.base_url + "/videos"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(url, data=form) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+        if "data" not in data:
+            raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
+        if not data["data"]:
+            raise RuntimeError("API returned empty data array")
+        try:
+            base64_str = data["data"][0]["b64_json"]
+        except (KeyError, IndexError):
+            raise RuntimeError("API response missing 'b64_json' field in first data item")
+
+        return base64_to_video(base64_str)
+
     async def generate_understanding_chat_completion(
         self,
         *,

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py

@@ -6,12 +6,17 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
 
 import base64
 import mimetypes
+from fractions import Fraction
 from io import BytesIO
 
 import av
 import numpy as np
 import torch
+from av.audio.frame import AudioFrame
+from av.audio.resampler import AudioResampler
+from av.video.frame import VideoFrame
 from comfy_api.input import AudioInput, VideoInput
+from comfy_api.latest import InputImpl, Types
 from comfy_extras import nodes_audio
 from PIL import Image
 
@@ -138,6 +143,71 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
     return f"data:{mime_type};base64,{base64_str}"
 
 
+def base64_to_video(base64_str: str) -> VideoInput:
+    if base64_str.startswith("data:video"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        video_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    video_buffer = BytesIO(video_bytes)
+
+    try:
+        with av.open(video_buffer, mode="r") as container:
+            video_stream = next((s for s in container.streams if s.type == "video"), None)
+            if video_stream is None:
+                raise ValueError("No video stream found in decoded payload.")
+
+            frames = []
+            for frame in container.decode(video_stream):
+                if not isinstance(frame, VideoFrame):
+                    continue
+                image = frame.to_ndarray(format="rgb24")
+                frames.append(torch.from_numpy(image).float() / 255.0)
+
+            if len(frames) == 0:
+                raise ValueError("No video frames found in decoded payload.")
+
+            images = torch.stack(frames, dim=0)
+            frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
+
+            audio: AudioInput | None = None
+            if len(container.streams.audio):
+                audio_stream = container.streams.audio[-1]
+                audio_frames = []
+                resampler = AudioResampler(format="fltp")
+                for frame in container.decode(audio_stream):
+                    if not isinstance(frame, AudioFrame):
+                        continue
+                    resampled = resampler.resample(frame)
+                    if not isinstance(resampled, list):
+                        resampled = [resampled]
+                    for audio_frame in resampled:
+                        if audio_frame is not None:
+                            audio_frames.append(audio_frame.to_ndarray())
+
+                if len(audio_frames) > 0:
+                    audio_data = np.concatenate(audio_frames, axis=1)
+                    sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
+                    audio = {
+                        "waveform": torch.from_numpy(audio_data).unsqueeze(0),
+                        "sample_rate": sample_rate,
+                    }
+
+            components = Types.VideoComponents(
+                images=images,
+                frame_rate=frame_rate,
+                audio=audio,
+                metadata=container.metadata if container.metadata else None,
+            )
+    except Exception as e:
+        raise RuntimeError(f"Failed to decode video: {e}")
+
+    return InputImpl.VideoFromComponents(components)
+
+
 def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
     waveform = audio["waveform"][0]  # Shape: (C, T)
     sample_rate = audio["sample_rate"]

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py

@@ -23,6 +23,10 @@ class QwenTTSModelSpecificParams(dict):
     pass
 
 
+class WanModelSpecificParams(dict):
+    pass
+
+
 class ModelMode(Enum):
     IMAGE_GENERATION = auto()
     VIDEO_GENERATION = auto()

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py

@@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types(
     # Skip if no spec or no sampling params
     pipeline_spec, _ = lookup_model_spec(model_name)
     if pipeline_spec is None:
-        logger.info("skipping sampling params check because spec is not found")
+        logger.info(f"skipping sampling params check because spec for {model_name} is not found")
         return
     if sampling_param_list is None:
         return

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json

@@ -0,0 +1,513 @@
+{
+    "id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
+    "revision": 0,
+    "last_node_id": 14,
+    "last_link_id": 28,
+    "nodes": [
+        {
+            "id": 10,
+            "type": "MarkdownNote",
+            "pos": [
+                227.99306819575278,
+                -231.24143306069843
+            ],
+            "size": [
+                240.3326478922474,
+                136.3505820791881
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Task and Input",
+            "properties": {},
+            "widgets_values": [
+                "vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 4,
+            "type": "LoadImage",
+            "pos": [
+                496.31859627609606,
+                -229.71277089860084
+            ],
+            "size": [
+                270,
+                314
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "image"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "IMAGE",
+                    "name": "IMAGE",
+                    "type": "IMAGE",
+                    "links": [
+                        25
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "cute-cat.jpg",
+                "image"
+            ]
+        },
+        {
+            "id": 8,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                465.5140218220927,
+                148.80072646828987
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        27
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1992,
+                "randomize"
+            ]
+        },
+        {
+            "id": 14,
+            "type": "VLLMOmniWanParams",
+            "pos": [
+                518.6005521845747,
+                428.4455663719664
+            ],
+            "size": [
+                270,
+                106
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "guidance_scale_2",
+                    "name": "guidance_scale_2",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale_2"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "boundary_ratio",
+                    "name": "boundary_ratio",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "boundary_ratio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "flow_shift",
+                    "name": "flow_shift",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "flow_shift"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "Wan video params",
+                    "name": "Wan video params",
+                    "type": "VIDEO_PARAMS",
+                    "links": [
+                        28
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniWanParams"
+            },
+            "widgets_values": [
+                4,
+                0.875,
+                5
+            ]
+        },
+        {
+            "id": 12,
+            "type": "VLLMOmniGenerateVideo",
+            "pos": [
+                827.2566336087867,
+                -73.77449831827545
+            ],
+            "size": [
+                400,
+                346
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 25
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 27
+                },
+                {
+                    "localized_name": "lora",
+                    "name": "lora",
+                    "shape": 7,
+                    "type": "REMOTE_LORA",
+                    "link": null
+                },
+                {
+                    "localized_name": "model_params",
+                    "name": "model_params",
+                    "shape": 7,
+                    "type": "VIDEO_PARAMS",
+                    "link": 28
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "fps",
+                    "name": "fps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "fps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_frames",
+                    "name": "num_frames",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_frames"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "type": "VIDEO",
+                    "links": [
+                        26
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateVideo"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+                "Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
+                "Cartoonish",
+                900,
+                1200,
+                12,
+                41
+            ]
+        },
+        {
+            "id": 13,
+            "type": "SaveVideo",
+            "pos": [
+                1257.5441369501657,
+                -44.12956394419863
+            ],
+            "size": [
+                374,
+                372.34896826553796
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "type": "VIDEO",
+                    "link": 26
+                },
+                {
+                    "localized_name": "filename_prefix",
+                    "name": "filename_prefix",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "filename_prefix"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "format",
+                    "name": "format",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "format"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "codec",
+                    "name": "codec",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "codec"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {},
+            "widgets_values": [
+                "video/ComfyUI",
+                "auto",
+                "auto"
+            ]
+        }
+    ],
+    "links": [
+        [
+            25,
+            4,
+            0,
+            12,
+            0,
+            "IMAGE"
+        ],
+        [
+            26,
+            12,
+            0,
+            13,
+            0,
+            "VIDEO"
+        ],
+        [
+            27,
+            8,
+            0,
+            12,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            28,
+            14,
+            0,
+            12,
+            3,
+            "VIDEO_PARAMS"
+        ]
+    ],
+    "groups": [
+        {
+            "id": 1,
+            "title": "Input",
+            "bounding": [
+                213.1706010087147,
+                -313.20095750667554,
+                560.8246471437027,
+                407.4246532472182
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        }
+    ],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 1.1469075819486894,
+            "offset": [
+                101.7632950847089,
+                492.29121889347806
+            ]
+        }
+    },
+    "version": 0.4
+}

docs/features/comfyui.md

@@ -47,7 +47,8 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities:
 
 - **Generate Image** for text-to-image and image-to-image tasks
-- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **Generate Video** for text-to-video and image-to-video tasks
+- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks
 
 This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
 
 - Drag a generation node onto the canvas.
 - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
-- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
 - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
     - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
     - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.

tests/comfyui/conftest.py

@@ -63,6 +63,12 @@ def _setup_comfyui_test_environment():
     mock_comfy_api_input.AudioInput = AudioInput
     mock_comfy_api_input.VideoInput = VideoInput
     mock_comfy_api.input = mock_comfy_api_input
+    mock_comfy_api_latest = MagicMock()
+    mock_comfy_api_latest.Types.VideoComponents = MagicMock(side_effect=lambda **kwargs: kwargs)
+    mock_comfy_api_latest.InputImpl.VideoFromComponents = MagicMock(
+        side_effect=lambda _: VideoInput(b"mock_video_from_components")
+    )
+    mock_comfy_api.latest = mock_comfy_api_latest
 
     def mock_load(_: str | BinaryIO):
         """Mock nodes_audio.load that returns a waveform tensor (channels, samples) and sample rate."""
@@ -78,5 +84,6 @@ def _setup_comfyui_test_environment():
     # Install mock modules BEFORE importing any comfyui_vllm_omni code
     sys.modules["comfy_api"] = mock_comfy_api
     sys.modules["comfy_api.input"] = mock_comfy_api_input
+    sys.modules["comfy_api.latest"] = mock_comfy_api_latest
     sys.modules["comfy_extras"] = mock_comfy_extras
     sys.modules["comfy_extras.nodes_audio"] = mock_nodes_audio

tests/comfyui/test_comfyui_integration.py

@@ -20,11 +20,12 @@ import torch
 from comfy_api.input import AudioInput, VideoInput
 from comfyui_vllm_omni.nodes import (
     VLLMOmniGenerateImage,
+    VLLMOmniGenerateVideo,
     VLLMOmniTTS,
     VLLMOmniUnderstanding,
     VLLMOmniVoiceClone,
 )
-from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams
+from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams, WanModelSpecificParams
 from PIL import Image
 from vllm import SamplingParams
 from vllm.outputs import CompletionOutput, RequestOutput
@@ -34,6 +35,8 @@ from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class ServerCase(NamedTuple):
     """Parametrizing the model to serve."""
@@ -49,6 +52,7 @@ class SamplingCase(NamedTuple):
 
     kind: "SamplingKind"
     sampling_params: dict | list[dict] | None
+    lora: dict | None = None
 
 
 class SamplingKind(StrEnum):
@@ -58,11 +62,17 @@ class SamplingKind(StrEnum):
     UNDERSTANDING_AR_LIST = auto()
     TTS_NONE = auto()
     TTS_DIFFUSION_SINGLE = auto()
+    VIDEO_NONE = auto()
+    VIDEO_DIFFUSION_SINGLE = auto()
 
 
 # Pre-defined arguments to be used in function calls during the tests
 IMAGE_WIDTH = 64
 IMAGE_HEIGHT = 64
+VIDEO_WIDTH = 32
+VIDEO_HEIGHT = 32
+VIDEO_FPS = 8
+VIDEO_NUM_FRAMES = 5
 DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
     {
         "n": 2,
@@ -102,6 +112,10 @@ AR_LIST_SAMPLING_PARAMS = [
     ),
 ]
 
+VIDEO_MODEL_PARAMS = WanModelSpecificParams({"guidance_scale_2": 5.0, "boundary_ratio": 0.98, "flow_shift": 12.0})
+
+LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10}
+
 
 def _build_image_output(size: tuple[int, int] = (IMAGE_WIDTH, IMAGE_HEIGHT), color: str = "red") -> Image.Image:
     return Image.new("RGB", size, color=color)
@@ -181,6 +195,16 @@ def _build_diffusion_image_output_for_images_endpoint() -> OmniRequestOutput:
     )
 
 
+def _build_diffusion_video_output() -> OmniRequestOutput:
+    # Small video: VIDEO_NUM_FRAMES frames of (VIDEO_HEIGHT x VIDEO_WIDTH) RGB, shape (F, H, W, C)
+    video_frames = torch.zeros((VIDEO_NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
+    return OmniRequestOutput.from_diffusion(
+        request_id="test_req_video",
+        images=[video_frames],
+        final_output_type="video",
+    )
+
+
 def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
     request_output = MagicMock()
     request_output.images = [_build_image_output(color="blue")]
@@ -193,12 +217,40 @@ def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
     )
 
 
-def _assert_sampling_param_values(received: OmniSamplingParams, expected: dict[str, Any]):
+def _assert_sampling_param_values(
+    received: OmniSamplingParams, expected: dict[str, Any], expected_lora: dict | None = None
+):
     for key, expected_value in expected.items():
-        actual_value = getattr(received, key, None)
+        actual_value = getattr(received, key)
         assert actual_value == expected_value, (
             f"Expected sampling param '{key}'={expected_value}, got {actual_value}. The received sampling params: {received}"
         )
+    if expected_lora:
+        assert received.lora_request.lora_name == expected_lora["name"], (
+            f"Expected lora name={(expected_lora['name'])}, got {received.lora_request.lora_name}. The received sampling params: {received}"
+        )
+        assert received.lora_request.lora_int_id == expected_lora["int_id"], (
+            f"Expected lora int_id={expected_lora['int_id']}, got {received.lora_request.lora_int_id}. The received sampling params: {received}"
+        )
+        assert received.lora_request.lora_path == expected_lora["local_path"], (
+            f"Expected lora path={expected_lora['local_path']}, got {received.lora_request.lora_path}. The received sampling params: {received}"
+        )
+        assert received.lora_scale == expected_lora["scale"], (
+            f"Expected lora scale={expected_lora['scale']}, got {received.lora_scale}. The received sampling params: {received}"
+        )
+
+
+def _assert_model_param_values(received: OmniSamplingParams, expected: dict):
+    for key, expected_value in expected.items():
+        try:
+            actual_value = getattr(received, key)
+            expected_param_name = key
+        except AttributeError:
+            actual_value = received.extra_args.get(key, None)
+            expected_param_name = f'extra_args["{key}"]'
+        assert actual_value == expected_value, (
+            f"Expected model param '{expected_param_name}'={expected_value}, got {actual_value}. The received sampling params: {received}"
+        )
 
 
 def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase):
@@ -235,6 +287,7 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
                     "height": IMAGE_HEIGHT,
                     **expected,
                 },
+                LORA_PARAMS,
             )
         elif sampling_case.kind is SamplingKind.UNDERSTANDING_NONE:
             assert len(received_sampling_params_list) == 3
@@ -244,6 +297,33 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
                 _assert_sampling_param_values(received_sampling_params_list[i], expected)
         elif sampling_case.kind in {SamplingKind.TTS_NONE, SamplingKind.TTS_DIFFUSION_SINGLE}:
             assert len(received_sampling_params_list) == 1
+        elif sampling_case.kind is SamplingKind.VIDEO_NONE:
+            assert len(received_sampling_params_list) == 1
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": VIDEO_WIDTH,
+                    "height": VIDEO_HEIGHT,
+                    "num_frames": VIDEO_NUM_FRAMES,
+                    "fps": VIDEO_FPS,
+                },
+            )
+        elif sampling_case.kind is SamplingKind.VIDEO_DIFFUSION_SINGLE:
+            assert len(received_sampling_params_list) == 1
+            expected = DIFFUSION_SINGLE_SAMPLING_PARAMS.copy()
+            expected["num_outputs_per_prompt"] = expected.pop("n")  # convert from n to num_outputs_per_prompt
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": VIDEO_WIDTH,
+                    "height": VIDEO_HEIGHT,
+                    "num_frames": VIDEO_NUM_FRAMES,
+                    "fps": VIDEO_FPS,
+                    **expected,
+                },
+                LORA_PARAMS,
+            )
+            _assert_model_param_values(received_sampling_params_list[0], VIDEO_MODEL_PARAMS)
         else:
             raise AssertionError(f"Unknown sampling case: {sampling_case.kind}")
 
@@ -401,9 +481,15 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
 @pytest.mark.parametrize(
     "sampling_case",
     [
-        pytest.param(SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None), id="no-sampling-params"),
         pytest.param(
-            SamplingCase(kind=SamplingKind.IMAGE_DIFFUSION_SINGLE, sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS),
+            SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None, lora=None), id="no-sampling-params"
+        ),
+        pytest.param(
+            SamplingCase(
+                kind=SamplingKind.IMAGE_DIFFUSION_SINGLE,
+                sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS,
+                lora=LORA_PARAMS,
+            ),
             id="single-diffusion-sampling-params",
         ),
     ],
@@ -423,7 +509,8 @@ async def test_image_generation_node(api_server: str, model: str, image_input: b
         kwargs["image"] = torch.zeros((1, IMAGE_WIDTH, IMAGE_HEIGHT, 3), dtype=torch.float32)
     if sampling_case.sampling_params is not None:
         kwargs["sampling_params"] = sampling_case.sampling_params
-    print(f"!!!!!! Calling {model} node.generate with kwargs: {sampling_case.sampling_params}")
+    if sampling_case.lora:
+        kwargs["lora"] = sampling_case.lora
 
     result = await node.generate(**kwargs)
 
@@ -565,3 +652,75 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
     assert result[0]["sample_rate"] == 24000
     assert isinstance(result[0]["waveform"], torch.Tensor)
     assert result[0]["waveform"].shape == (1, 1, 24000)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case,model,image_input",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            False,
+            id="text-to-video",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+            True,
+            id="image-to-video",
+        ),
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(SamplingCase(kind=SamplingKind.VIDEO_NONE, sampling_params=None), id="no-sampling-params"),
+        pytest.param(
+            SamplingCase(
+                kind=SamplingKind.VIDEO_DIFFUSION_SINGLE,
+                sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS,
+                lora=LORA_PARAMS,
+            ),
+            id="single-diffusion-sampling-params",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_video_generation_node(api_server: str, model: str, image_input: bool, sampling_case: SamplingCase):
+    node = VLLMOmniGenerateVideo()
+
+    kwargs = {
+        "url": api_server,
+        "model": model,
+        "prompt": "A beautiful sunset timelapse",
+        "negative_prompt": "",
+        "width": VIDEO_WIDTH,
+        "height": VIDEO_HEIGHT,
+        "fps": VIDEO_FPS,
+        "num_frames": VIDEO_NUM_FRAMES,
+        "model_params": VIDEO_MODEL_PARAMS,
+    }
+    if image_input:
+        kwargs["image"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
+    if sampling_case.sampling_params is not None:
+        kwargs["sampling_params"] = sampling_case.sampling_params
+    if sampling_case.lora:
+        kwargs["lora"] = sampling_case.lora
+
+    result = await node.generate(**kwargs)
+
+    assert isinstance(result, tuple)
+    assert len(result) == 1
+    assert isinstance(result[0], VideoInput)

tests/entrypoints/openai_api/test_image_server.py

@@ -136,7 +136,6 @@ def mock_async_diffusion(mocker: MockerFixture):
 
     async def generate(**kwargs):
         # Return n PIL images wrapped in result object
-        print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
         n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
         mock.captured_sampling_params_list = kwargs["sampling_params_list"]
         mock.captured_prompt = kwargs["prompt"]