comfy/model_patcher.py

@@ -704,7 +704,7 @@ class ModelPatcher:
 
                 lowvram_weight = False
 
-                potential_offload = max(offload_buffer, module_offload_mem + (comfy.model_management.NUM_STREAMS * module_mem))
+                potential_offload = max(offload_buffer, module_offload_mem * (comfy.model_management.NUM_STREAMS + 1))
                 lowvram_fits = mem_counter + module_mem + potential_offload < lowvram_model_memory
 
                 weight_key = "{}.weight".format(n)
@@ -883,7 +883,7 @@ class ModelPatcher:
                     break
                 module_offload_mem, module_mem, n, m, params = unload
 
-                potential_offload = module_offload_mem + (comfy.model_management.NUM_STREAMS * module_mem)
+                potential_offload = (comfy.model_management.NUM_STREAMS + 1) * module_offload_mem
 
                 lowvram_possible = hasattr(m, "comfy_cast_weights")
                 if hasattr(m, "comfy_patched_weights") and m.comfy_patched_weights == True:

comfy/ops.py

@@ -111,24 +111,22 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
     if s.bias is not None:
         bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=bias_has_function, stream=offload_stream)
 
-    comfy.model_management.sync_stream(device, offload_stream)
-
-    bias_a = bias
-    weight_a = weight
-
-    if s.bias is not None:
-        for f in s.bias_function:
-            bias = f(bias)
+        if bias_has_function:
+            with wf_context:
+                for f in s.bias_function:
+                    bias = f(bias)
 
     if weight_has_function or weight.dtype != dtype:
-        weight = weight.to(dtype=dtype)
-        if isinstance(weight, QuantizedTensor):
-            weight = weight.dequantize()
-        for f in s.weight_function:
-            weight = f(weight)
+        with wf_context:
+            weight = weight.to(dtype=dtype)
+            if isinstance(weight, QuantizedTensor):
+                weight = weight.dequantize()
+            for f in s.weight_function:
+                weight = f(weight)
 
+    comfy.model_management.sync_stream(device, offload_stream)
     if offloadable:
-        return weight, bias, (offload_stream, weight_a, bias_a)
+        return weight, bias, offload_stream
     else:
         #Legacy function signature
         return weight, bias
@@ -137,16 +135,13 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
 def uncast_bias_weight(s, weight, bias, offload_stream):
     if offload_stream is None:
         return
-    os, weight_a, bias_a = offload_stream
-    if os is None:
-        return
-    if weight_a is not None:
-        device = weight_a.device
+    if weight is not None:
+        device = weight.device
     else:
-        if bias_a is None:
+        if bias is None:
             return
-        device = bias_a.device
-    os.wait_stream(comfy.model_management.current_stream(device))
+        device = bias.device
+    offload_stream.wait_stream(comfy.model_management.current_stream(device))
 
 
 class CastWeightBiasOp:

comfy_api/latest/_io.py

@@ -5,7 +5,7 @@ import inspect
 from abc import ABC, abstractmethod
 from collections import Counter
 from collections.abc import Iterable
-from dataclasses import asdict, dataclass, field
+from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import Any, Callable, Literal, TypedDict, TypeVar, TYPE_CHECKING
 from typing_extensions import NotRequired, final
@@ -166,7 +166,7 @@ class Input(_IO_V3):
     '''
     Base class for a V3 Input.
     '''
-    def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None, raw_link: bool=None):
+    def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None):
         super().__init__()
         self.id = id
         self.display_name = display_name
@@ -174,7 +174,6 @@ class Input(_IO_V3):
         self.tooltip = tooltip
         self.lazy = lazy
         self.extra_dict = extra_dict if extra_dict is not None else {}
-        self.rawLink = raw_link
 
     def as_dict(self):
         return prune_dict({
@@ -182,7 +181,6 @@ class Input(_IO_V3):
             "optional": self.optional,
             "tooltip": self.tooltip,
             "lazy": self.lazy,
-            "rawLink": self.rawLink,
         }) | prune_dict(self.extra_dict)
 
     def get_io_type(self):
@@ -197,8 +195,8 @@ class WidgetInput(Input):
     '''
     def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
                  default: Any=None,
-                 socketless: bool=None, widget_type: str=None, force_input: bool=None, extra_dict=None, raw_link: bool=None):
-        super().__init__(id, display_name, optional, tooltip, lazy, extra_dict, raw_link)
+                 socketless: bool=None, widget_type: str=None, force_input: bool=None, extra_dict=None):
+        super().__init__(id, display_name, optional, tooltip, lazy, extra_dict)
         self.default = default
         self.socketless = socketless
         self.widget_type = widget_type
@@ -254,8 +252,8 @@ class Boolean(ComfyTypeIO):
         '''Boolean input.'''
         def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
                     default: bool=None, label_on: str=None, label_off: str=None,
-                    socketless: bool=None, force_input: bool=None, extra_dict=None, raw_link: bool=None):
-            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, force_input, extra_dict, raw_link)
+                    socketless: bool=None, force_input: bool=None):
+            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, force_input)
             self.label_on = label_on
             self.label_off = label_off
             self.default: bool
@@ -274,8 +272,8 @@ class Int(ComfyTypeIO):
         '''Integer input.'''
         def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
                     default: int=None, min: int=None, max: int=None, step: int=None, control_after_generate: bool=None,
-                    display_mode: NumberDisplay=None, socketless: bool=None, force_input: bool=None, extra_dict=None, raw_link: bool=None):
-            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, force_input, extra_dict, raw_link)
+                    display_mode: NumberDisplay=None, socketless: bool=None, force_input: bool=None):
+            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, force_input)
             self.min = min
             self.max = max
             self.step = step
@@ -300,8 +298,8 @@ class Float(ComfyTypeIO):
         '''Float input.'''
         def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
                     default: float=None, min: float=None, max: float=None, step: float=None, round: float=None,
-                    display_mode: NumberDisplay=None, socketless: bool=None, force_input: bool=None, extra_dict=None, raw_link: bool=None):
-            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, force_input, extra_dict, raw_link)
+                    display_mode: NumberDisplay=None, socketless: bool=None, force_input: bool=None):
+            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, force_input)
             self.min = min
             self.max = max
             self.step = step
@@ -326,8 +324,8 @@ class String(ComfyTypeIO):
         '''String input.'''
         def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
                     multiline=False, placeholder: str=None, default: str=None, dynamic_prompts: bool=None,
-                    socketless: bool=None, force_input: bool=None, extra_dict=None, raw_link: bool=None):
-            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, force_input, extra_dict, raw_link)
+                    socketless: bool=None, force_input: bool=None):
+            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, force_input)
             self.multiline = multiline
             self.placeholder = placeholder
             self.dynamic_prompts = dynamic_prompts
@@ -360,14 +358,12 @@ class Combo(ComfyTypeIO):
             image_folder: FolderType=None,
             remote: RemoteOptions=None,
             socketless: bool=None,
-            extra_dict=None,
-            raw_link: bool=None,
         ):
             if isinstance(options, type) and issubclass(options, Enum):
                 options = [v.value for v in options]
             if isinstance(default, Enum):
                 default = default.value
-            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, None, extra_dict, raw_link)
+            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless)
             self.multiselect = False
             self.options = options
             self.control_after_generate = control_after_generate
@@ -399,8 +395,8 @@ class MultiCombo(ComfyTypeI):
     class Input(Combo.Input):
         def __init__(self, id: str, options: list[str], display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
                     default: list[str]=None, placeholder: str=None, chip: bool=None, control_after_generate: bool=None,
-                    socketless: bool=None, extra_dict=None, raw_link: bool=None):
-            super().__init__(id, options, display_name, optional, tooltip, lazy, default, control_after_generate, socketless=socketless, extra_dict=extra_dict, raw_link=raw_link)
+                    socketless: bool=None):
+            super().__init__(id, options, display_name, optional, tooltip, lazy, default, control_after_generate, socketless=socketless)
             self.multiselect = True
             self.placeholder = placeholder
             self.chip = chip
@@ -433,9 +429,9 @@ class Webcam(ComfyTypeIO):
         Type = str
         def __init__(
                 self, id: str, display_name: str=None, optional=False,
-                tooltip: str=None, lazy: bool=None, default: str=None, socketless: bool=None, extra_dict=None, raw_link: bool=None
+                tooltip: str=None, lazy: bool=None, default: str=None, socketless: bool=None
         ):
-            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless, None, None, extra_dict, raw_link)
+            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless)
 
 
 @comfytype(io_type="MASK")
@@ -779,7 +775,7 @@ class MultiType:
         '''
         Input that permits more than one input type; if `id` is an instance of `ComfyType.Input`, then that input will be used to create a widget (if applicable) with overridden values.
         '''
-        def __init__(self, id: str | Input, types: list[type[_ComfyType] | _ComfyType], display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None, raw_link: bool=None):
+        def __init__(self, id: str | Input, types: list[type[_ComfyType] | _ComfyType], display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None):
             # if id is an Input, then use that Input with overridden values
             self.input_override = None
             if isinstance(id, Input):
@@ -792,7 +788,7 @@ class MultiType:
                 # if is a widget input, make sure widget_type is set appropriately
                 if isinstance(self.input_override, WidgetInput):
                     self.input_override.widget_type = self.input_override.get_io_type()
-            super().__init__(id, display_name, optional, tooltip, lazy, extra_dict, raw_link)
+            super().__init__(id, display_name, optional, tooltip, lazy, extra_dict)
             self._io_types = types
 
         @property
@@ -846,8 +842,8 @@ class MatchType(ComfyTypeIO):
 
     class Input(Input):
         def __init__(self, id: str, template: MatchType.Template,
-                    display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None, raw_link: bool=None):
-            super().__init__(id, display_name, optional, tooltip, lazy, extra_dict, raw_link)
+                    display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None):
+            super().__init__(id, display_name, optional, tooltip, lazy, extra_dict)
             self.template = template
 
         def as_dict(self):
@@ -1214,9 +1210,9 @@ class Schema:
     """Display name of node."""
     category: str = "sd"
     """The category of the node, as per the "Add Node" menu."""
-    inputs: list[Input] = field(default_factory=list)
-    outputs: list[Output] = field(default_factory=list)
-    hidden: list[Hidden] = field(default_factory=list)
+    inputs: list[Input]=None
+    outputs: list[Output]=None
+    hidden: list[Hidden]=None
     description: str=""
     """Node description, shown as a tooltip when hovering over the node."""
     is_input_list: bool = False
@@ -1257,10 +1253,11 @@ class Schema:
         - verify ids on inputs and outputs are unique - both internally and in relation to each other
         '''
         nested_inputs: list[Input] = []
-        for input in self.inputs:
-            nested_inputs.extend(input.get_all())
-        input_ids = [i.id for i in nested_inputs]
-        output_ids = [o.id for o in self.outputs]
+        if self.inputs is not None:
+            for input in self.inputs:
+                nested_inputs.extend(input.get_all())
+        input_ids = [i.id for i in nested_inputs] if nested_inputs is not None else []
+        output_ids = [o.id for o in self.outputs] if self.outputs is not None else []
         input_set = set(input_ids)
         output_set = set(output_ids)
         issues = []
@@ -1276,36 +1273,36 @@ class Schema:
         if len(issues) > 0:
             raise ValueError("\n".join(issues))
         # validate inputs and outputs
-        for input in self.inputs:
-            input.validate()
-        for output in self.outputs:
-            output.validate()
+        if self.inputs is not None:
+            for input in self.inputs:
+                input.validate()
+        if self.outputs is not None:
+            for output in self.outputs:
+                output.validate()
 
     def finalize(self):
         """Add hidden based on selected schema options, and give outputs without ids default ids."""
-        # ensure inputs, outputs, and hidden are lists
-        if self.inputs is None:
-            self.inputs = []
-        if self.outputs is None:
-            self.outputs = []
-        if self.hidden is None:
-            self.hidden = []
         # if is an api_node, will need key-related hidden
         if self.is_api_node:
+            if self.hidden is None:
+                self.hidden = []
             if Hidden.auth_token_comfy_org not in self.hidden:
                 self.hidden.append(Hidden.auth_token_comfy_org)
             if Hidden.api_key_comfy_org not in self.hidden:
                 self.hidden.append(Hidden.api_key_comfy_org)
         # if is an output_node, will need prompt and extra_pnginfo
         if self.is_output_node:
+            if self.hidden is None:
+                self.hidden = []
             if Hidden.prompt not in self.hidden:
                 self.hidden.append(Hidden.prompt)
             if Hidden.extra_pnginfo not in self.hidden:
                 self.hidden.append(Hidden.extra_pnginfo)
         # give outputs without ids default ids
-        for i, output in enumerate(self.outputs):
-            if output.id is None:
-                output.id = f"_{i}_{output.io_type}_"
+        if self.outputs is not None:
+            for i, output in enumerate(self.outputs):
+                if output.id is None:
+                    output.id = f"_{i}_{output.io_type}_"
 
     def get_v1_info(self, cls, live_inputs: dict[str, Any]=None) -> NodeInfoV1:
         # NOTE: live_inputs will not be used anymore very soon and this will be done another way

comfy_api_nodes/apis/kling_api.py

@@ -46,41 +46,21 @@ class TaskStatusVideoResult(BaseModel):
     url: str | None = Field(None, description="URL for generated video")
 
 
-class TaskStatusImageResult(BaseModel):
-    index: int = Field(..., description="Image Number，0-9")
-    url: str = Field(..., description="URL for generated image")
-
-
-class OmniTaskStatusResults(BaseModel):
+class TaskStatusVideoResults(BaseModel):
     videos: list[TaskStatusVideoResult] | None = Field(None)
-    images: list[TaskStatusImageResult] | None = Field(None)
 
 
-class OmniTaskStatusResponseData(BaseModel):
+class TaskStatusVideoResponseData(BaseModel):
     created_at: int | None = Field(None, description="Task creation time")
     updated_at: int | None = Field(None, description="Task update time")
     task_status: str | None = None
     task_status_msg: str | None = Field(None, description="Additional failure reason. Only for polling endpoint.")
     task_id: str | None = Field(None, description="Task ID")
-    task_result: OmniTaskStatusResults | None = Field(None)
+    task_result: TaskStatusVideoResults | None = Field(None)
 
 
-class OmniTaskStatusResponse(BaseModel):
+class TaskStatusVideoResponse(BaseModel):
     code: int | None = Field(None, description="Error code")
     message: str | None = Field(None, description="Error message")
     request_id: str | None = Field(None, description="Request ID")
-    data: OmniTaskStatusResponseData | None = Field(None)
-
-
-class OmniImageParamImage(BaseModel):
-    image: str = Field(...)
-
-
-class OmniProImageRequest(BaseModel):
-    model_name: str = Field(..., description="kling-image-o1")
-    resolution: str = Field(..., description="'1k' or '2k'")
-    aspect_ratio: str | None = Field(...)
-    prompt: str = Field(...)
-    mode: str = Field("pro")
-    n: int | None = Field(1, le=9)
-    image_list: list[OmniImageParamImage] | None = Field(..., max_length=10)
+    data: TaskStatusVideoResponseData | None = Field(None)

comfy_api_nodes/nodes_kling.py

@@ -6,7 +6,6 @@ For source of truth on the allowed permutations of request fields, please refere
 
 import logging
 import math
-import re
 
 import torch
 from typing_extensions import override
@@ -50,14 +49,12 @@ from comfy_api_nodes.apis import (
     KlingSingleImageEffectModelName,
 )
 from comfy_api_nodes.apis.kling_api import (
-    OmniImageParamImage,
     OmniParamImage,
     OmniParamVideo,
     OmniProFirstLastFrameRequest,
-    OmniProImageRequest,
     OmniProReferences2VideoRequest,
     OmniProText2VideoRequest,
-    OmniTaskStatusResponse,
+    TaskStatusVideoResponse,
 )
 from comfy_api_nodes.util import (
     ApiEndpoint,
@@ -213,36 +210,7 @@ VOICES_CONFIG = {
 }
 
 
-def normalize_omni_prompt_references(prompt: str) -> str:
-    """
-    Rewrites Kling Omni-style placeholders used in the app, like:
-
-        @image, @image1, @image2, ... @imageN
-        @video, @video1, @video2, ... @videoN
-
-    into the API-compatible form:
-
-        <<<image_1>>>, <<<image_2>>>, ...
-        <<<video_1>>>, <<<video_2>>>, ...
-
-    This is a UX shim for ComfyUI so users can type the same syntax as in the Kling app.
-    """
-    if not prompt:
-        return prompt
-
-    def _image_repl(match):
-        return f"<<<image_{match.group('idx') or '1'}>>>"
-
-    def _video_repl(match):
-        return f"<<<video_{match.group('idx') or '1'}>>>"
-
-    # (?<!\w) avoids matching e.g. "test@image.com"
-    # (?!\w) makes sure we only match @image / @image<digits> and not @imageFoo
-    prompt = re.sub(r"(?<!\w)@image(?P<idx>\d*)(?!\w)", _image_repl, prompt)
-    return re.sub(r"(?<!\w)@video(?P<idx>\d*)(?!\w)", _video_repl, prompt)
-
-
-async def finish_omni_video_task(cls: type[IO.ComfyNode], response: OmniTaskStatusResponse) -> IO.NodeOutput:
+async def finish_omni_video_task(cls: type[IO.ComfyNode], response: TaskStatusVideoResponse) -> IO.NodeOutput:
     if response.code:
         raise RuntimeError(
             f"Kling request failed. Code: {response.code}, Message: {response.message}, Data: {response.data}"
@@ -250,9 +218,8 @@ async def finish_omni_video_task(cls: type[IO.ComfyNode], response: OmniTaskStat
     final_response = await poll_op(
         cls,
         ApiEndpoint(path=f"/proxy/kling/v1/videos/omni-video/{response.data.task_id}"),
-        response_model=OmniTaskStatusResponse,
+        response_model=TaskStatusVideoResponse,
         status_extractor=lambda r: (r.data.task_status if r.data else None),
-        max_poll_attempts=160,
     )
     return IO.NodeOutput(await download_url_to_video_output(final_response.data.task_result.videos[0].url))
 
@@ -834,7 +801,7 @@ class OmniProTextToVideoNode(IO.ComfyNode):
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/kling/v1/videos/omni-video", method="POST"),
-            response_model=OmniTaskStatusResponse,
+            response_model=TaskStatusVideoResponse,
             data=OmniProText2VideoRequest(
                 model_name=model_name,
                 prompt=prompt,
@@ -897,7 +864,6 @@ class OmniProFirstLastFrameNode(IO.ComfyNode):
         end_frame: Input.Image | None = None,
         reference_images: Input.Image | None = None,
     ) -> IO.NodeOutput:
-        prompt = normalize_omni_prompt_references(prompt)
         validate_string(prompt, min_length=1, max_length=2500)
         if end_frame is not None and reference_images is not None:
             raise ValueError("The 'end_frame' input cannot be used simultaneously with 'reference_images'.")
@@ -929,7 +895,7 @@ class OmniProFirstLastFrameNode(IO.ComfyNode):
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/kling/v1/videos/omni-video", method="POST"),
-            response_model=OmniTaskStatusResponse,
+            response_model=TaskStatusVideoResponse,
             data=OmniProFirstLastFrameRequest(
                 model_name=model_name,
                 prompt=prompt,
@@ -984,7 +950,6 @@ class OmniProImageToVideoNode(IO.ComfyNode):
         duration: int,
         reference_images: Input.Image,
     ) -> IO.NodeOutput:
-        prompt = normalize_omni_prompt_references(prompt)
         validate_string(prompt, min_length=1, max_length=2500)
         if get_number_of_images(reference_images) > 7:
             raise ValueError("The maximum number of reference images is 7.")
@@ -997,7 +962,7 @@ class OmniProImageToVideoNode(IO.ComfyNode):
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/kling/v1/videos/omni-video", method="POST"),
-            response_model=OmniTaskStatusResponse,
+            response_model=TaskStatusVideoResponse,
             data=OmniProReferences2VideoRequest(
                 model_name=model_name,
                 prompt=prompt,
@@ -1058,7 +1023,6 @@ class OmniProVideoToVideoNode(IO.ComfyNode):
         keep_original_sound: bool,
         reference_images: Input.Image | None = None,
     ) -> IO.NodeOutput:
-        prompt = normalize_omni_prompt_references(prompt)
         validate_string(prompt, min_length=1, max_length=2500)
         validate_video_duration(reference_video, min_duration=3.0, max_duration=10.05)
         validate_video_dimensions(reference_video, min_width=720, min_height=720, max_width=2160, max_height=2160)
@@ -1081,7 +1045,7 @@ class OmniProVideoToVideoNode(IO.ComfyNode):
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/kling/v1/videos/omni-video", method="POST"),
-            response_model=OmniTaskStatusResponse,
+            response_model=TaskStatusVideoResponse,
             data=OmniProReferences2VideoRequest(
                 model_name=model_name,
                 prompt=prompt,
@@ -1139,7 +1103,6 @@ class OmniProEditVideoNode(IO.ComfyNode):
         keep_original_sound: bool,
         reference_images: Input.Image | None = None,
     ) -> IO.NodeOutput:
-        prompt = normalize_omni_prompt_references(prompt)
         validate_string(prompt, min_length=1, max_length=2500)
         validate_video_duration(video, min_duration=3.0, max_duration=10.05)
         validate_video_dimensions(video, min_width=720, min_height=720, max_width=2160, max_height=2160)
@@ -1162,7 +1125,7 @@ class OmniProEditVideoNode(IO.ComfyNode):
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/kling/v1/videos/omni-video", method="POST"),
-            response_model=OmniTaskStatusResponse,
+            response_model=TaskStatusVideoResponse,
             data=OmniProReferences2VideoRequest(
                 model_name=model_name,
                 prompt=prompt,
@@ -1175,90 +1138,6 @@ class OmniProEditVideoNode(IO.ComfyNode):
         return await finish_omni_video_task(cls, response)
 
 
-class OmniProImageNode(IO.ComfyNode):
-
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="KlingOmniProImageNode",
-            display_name="Kling Omni Image (Pro)",
-            category="api node/image/Kling",
-            description="Create or edit images with the latest model from Kling.",
-            inputs=[
-                IO.Combo.Input("model_name", options=["kling-image-o1"]),
-                IO.String.Input(
-                    "prompt",
-                    multiline=True,
-                    tooltip="A text prompt describing the image content. "
-                    "This can include both positive and negative descriptions.",
-                ),
-                IO.Combo.Input("resolution", options=["1K", "2K"]),
-                IO.Combo.Input(
-                    "aspect_ratio",
-                    options=["16:9", "9:16", "1:1", "4:3", "3:4", "3:2", "2:3", "21:9"],
-                ),
-                IO.Image.Input(
-                    "reference_images",
-                    tooltip="Up to 10 additional reference images.",
-                    optional=True,
-                ),
-            ],
-            outputs=[
-                IO.Image.Output(),
-            ],
-            hidden=[
-                IO.Hidden.auth_token_comfy_org,
-                IO.Hidden.api_key_comfy_org,
-                IO.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        model_name: str,
-        prompt: str,
-        resolution: str,
-        aspect_ratio: str,
-        reference_images: Input.Image | None = None,
-    ) -> IO.NodeOutput:
-        prompt = normalize_omni_prompt_references(prompt)
-        validate_string(prompt, min_length=1, max_length=2500)
-        image_list: list[OmniImageParamImage] = []
-        if reference_images is not None:
-            if get_number_of_images(reference_images) > 10:
-                raise ValueError("The maximum number of reference images is 10.")
-            for i in reference_images:
-                validate_image_dimensions(i, min_width=300, min_height=300)
-                validate_image_aspect_ratio(i, (1, 2.5), (2.5, 1))
-            for i in await upload_images_to_comfyapi(cls, reference_images, wait_label="Uploading reference image"):
-                image_list.append(OmniImageParamImage(image=i))
-        response = await sync_op(
-            cls,
-            ApiEndpoint(path="/proxy/kling/v1/images/omni-image", method="POST"),
-            response_model=OmniTaskStatusResponse,
-            data=OmniProImageRequest(
-                model_name=model_name,
-                prompt=prompt,
-                resolution=resolution.lower(),
-                aspect_ratio=aspect_ratio,
-                image_list=image_list if image_list else None,
-            ),
-        )
-        if response.code:
-            raise RuntimeError(
-                f"Kling request failed. Code: {response.code}, Message: {response.message}, Data: {response.data}"
-            )
-        final_response = await poll_op(
-            cls,
-            ApiEndpoint(path=f"/proxy/kling/v1/images/omni-image/{response.data.task_id}"),
-            response_model=OmniTaskStatusResponse,
-            status_extractor=lambda r: (r.data.task_status if r.data else None),
-        )
-        return IO.NodeOutput(await download_url_to_image_tensor(final_response.data.task_result.images[0].url))
-
-
 class KlingCameraControlT2VNode(IO.ComfyNode):
     """
     Kling Text to Video Camera Control Node. This node is a text to video node, but it supports controlling the camera.
@@ -2056,7 +1935,6 @@ class KlingExtension(ComfyExtension):
             OmniProImageToVideoNode,
             OmniProVideoToVideoNode,
             OmniProEditVideoNode,
-            # OmniProImageNode,  # need support from backend
         ]