diff --git a/apps/ComfyUI-vLLM-Omni/README.md b/apps/ComfyUI-vLLM-Omni/README.md index 23bbe5b72..54f2fdf2e 100644 --- a/apps/ComfyUI-vLLM-Omni/README.md +++ b/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**) + +

+ + + vLLM-Omni video generation + +

+ +> [!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**) diff --git a/apps/ComfyUI-vLLM-Omni/__init__.py b/apps/ComfyUI-vLLM-Omni/__init__.py index 641824add..e89b2ca1e 100644 --- a/apps/ComfyUI-vLLM-Omni/__init__.py +++ b/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" diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py index 5caa3869e..bfea93998 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py +++ b/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),) diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py index d340073bc..2c936d33f 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py +++ b/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, *, diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py index 0ddf86f9c..642c83c8d 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py +++ b/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"] diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py index 6cc79f3ed..c7d254eb9 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py +++ b/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() diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py index 2abdd2100..c5ac4bdf4 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py +++ b/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 diff --git a/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg new file mode 100644 index 000000000..a449a4eee Binary files /dev/null and b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg differ diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json new file mode 100644 index 000000000..636360ab8 --- /dev/null +++ b/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 +} diff --git a/docs/features/comfyui.md b/docs/features/comfyui.md index 3b7a9551e..ff766cbb3 100644 --- a/docs/features/comfyui.md +++ b/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. diff --git a/tests/comfyui/conftest.py b/tests/comfyui/conftest.py index 7c9770b8e..0b4565e94 100644 --- a/tests/comfyui/conftest.py +++ b/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 diff --git a/tests/comfyui/test_comfyui_integration.py b/tests/comfyui/test_comfyui_integration.py index 463596326..d28818ab0 100644 --- a/tests/comfyui/test_comfyui_integration.py +++ b/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) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 36c58b3d8..743e36a05 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/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"]