refactor: cleanup/remove code for testing video stuff

add basically all inputs to one node, venice parameters etc, use dynamic combo for hiding unsupported capabilities
fix encode tensor for vision
add text model spec get to utils/catalog

.gitignore

@@ -4,6 +4,7 @@ js/test.js
 veniceai_config.json
 data/*
 testing_video/*
+nodes/test/*
 
 # Ignore Mac system files
 .DS_Store

__init__.py

@@ -1,41 +1,31 @@
-import importlib
-import logging
 import os
-from pathlib import Path
 
-from .pyserver import routes  # noqa: F401
+from comfy_api.latest import ComfyExtension, io
 
-NODE_PACKAGE = Path(__file__).with_name("nodes")
-node_list = sorted(
-    module.stem for module in NODE_PACKAGE.glob("*.py") if module.is_file() and not module.name.startswith("_")
-)
+from .nodes.gen_image_node import GenerateImage
+from .nodes.gen_speech_node import GenerateSpeech
+from .nodes.gen_text_node import GenerateTextAdvanced
+from .nodes.gen_video_from_text_node import GenerateVideoFromText
+from .nodes.i2i_enhance_upscale import I2IEnhanceUpscale
+from .nodes.test_node import DCTestNode
 
-NODE_CLASS_MAPPINGS = {}
-NODE_DISPLAY_NAME_MAPPINGS = {}
-NODE_MODULE_PREFIX = f"{__name__}.nodes"
 
-for module_name in node_list:
-    try:
-        imported_module = importlib.import_module(f"{NODE_MODULE_PREFIX}.{module_name}")
-    except ImportError as e:
-        logging.warning(f"Could not import module '{module_name}': {e}")
-        continue
+class VeniceExtension(ComfyExtension):
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            DCTestNode,
+            GenerateImage,
+            GenerateSpeech,
+            GenerateTextAdvanced,
+            GenerateVideoFromText,
+            I2IEnhanceUpscale,
+        ]
 
-    has_classes = hasattr(imported_module, "NODE_CLASS_MAPPINGS")
-    has_display = hasattr(imported_module, "NODE_DISPLAY_NAME_MAPPINGS")
 
-    if has_classes and has_display:
-        NODE_CLASS_MAPPINGS.update(imported_module.NODE_CLASS_MAPPINGS)
-        NODE_DISPLAY_NAME_MAPPINGS.update(imported_module.NODE_DISPLAY_NAME_MAPPINGS)
-    elif has_classes or has_display:
-        logging.warning(
-            "Module '%s' defines '%s' but is missing '%s'; module skipped",
-            module_name,
-            "NODE_CLASS_MAPPINGS" if has_classes else "NODE_DISPLAY_NAME_MAPPINGS",
-            "NODE_DISPLAY_NAME_MAPPINGS" if has_classes else "NODE_CLASS_MAPPINGS",
-        )
-    # modules without either mapping are ignored silently
+async def comfy_entrypoint() -> VeniceExtension:
+    return VeniceExtension()
+
 
 WEB_DIRECTORY = os.path.join(os.path.dirname(__file__), "js")
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
+__all__ = ["WEB_DIRECTORY"]

nodes/catalog_utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Iterable, Sequence, Tuple
+from typing import Any, Dict, Iterable, Sequence, Tuple
 
 from ..venice_catalog import get_characters, get_models, get_styles
 
@@ -46,10 +46,29 @@ def image_model_choices() -> Tuple[str, ...]:
 def text2video_model_choices() -> Tuple[str, ...]:
     return _safe_values(lambda: get_models(), "text2video_models")
 
+
 def image2video_model_choices() -> Tuple[str, ...]:
     return _safe_values(lambda: get_models(), "image2video_models")
 
 
+def video_model_specs() -> Dict[str, Dict[str, Any]]:
+    """Return a by-id mapping of video model specs with constraints for DynamicCombo use."""
+    models = get_models()
+    return models.get("video_models_by_id", {})
+
+
+def image_model_specs() -> Dict[str, Dict[str, Any]]:
+    """Return a by-id mapping of image model specs and constraints for UI validation."""
+    models = get_models()
+    return models.get("image_models_by_id", {})
+
+
+def text_model_specs() -> Dict[str, Dict[str, Any]]:
+    """Return a by-id mapping of text model specs and constraints for UI validation."""
+    models = get_models()
+    return models.get("text_models_by_id", {})
+
+
 def text_model_choices() -> Tuple[str, ...]:
     return _safe_values(lambda: get_models(), "text_models")
 

nodes/gen_image_base.py

@@ -2,13 +2,10 @@ import base64
 import io
 import logging
 
-import numpy as np
-import torch  # type: ignore
-import torchvision.transforms as transforms  # type: ignore
+import torchvision.transforms as transforms
 from PIL import Image
 
 
-
 class GenerateImageBase:
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "generate"
@@ -34,10 +31,3 @@ class GenerateImageBase:
 
         except Exception as e:
             raise Exception(f"Error processing image result: {str(e)}") from e
-
-    # unused
-    def create_blank_image(self):
-        blank_img = Image.new("RGB", (64, 64), color="black")
-        img_array = np.array(blank_img).astype(np.float32) / 255.0
-        img_tensor = torch.from_numpy(img_array)[None,]
-        return (img_tensor,)

nodes/gen_image_node.py

@@ -1,208 +1,324 @@
 import logging
-import re
+from typing import Any, Dict
 
 import torch
 
+from comfy_api.latest import io
+
 from ..globals import API_ENDPOINTS
-from ..nodes.catalog_utils import image_model_choices, style_choices
+from ..nodes.catalog_utils import image_model_specs, style_choices
 from ..nodes.gen_image_base import GenerateImageBase
 from ..nodes.utils import ensure_multiple_of, ensure_prompt_length
 from ..venice_client import client
 
+LOG = logging.getLogger(__name__)
+
+
+class GenerateImage(io.ComfyNode):
+    _processor = GenerateImageBase()
 
-class GenerateImage(GenerateImageBase):
     @classmethod
-    def INPUT_TYPES(cls):
-        model_choices = image_model_choices()
-        style_preset_options = style_choices()
+    def _image_specs(cls, require: bool = False) -> Dict[str, Dict[str, Any]]:
+        specs = image_model_specs() or {}
+        if require and not specs:
+            raise ValueError(
+                "No Venice image model specs available; refresh the catalog in VeniceAI settings and retry."
+            )
+        return specs
 
-        return {
-            "required": {
-                "model": (
-                    model_choices,
-                    {
-                        "default": model_choices[0],
-                        "tooltip": "Model to use for image generation",
-                    },
-                ),
-                "prompt": (
-                    "STRING",
-                    {
-                        "default": "A flying cat made of lettuce",
-                        "multiline": True,
-                        "tooltip": "The text prompt to guide the image generation",
-                        "placeholder": "Positive Prompt. Example: A flying cat made of lettuce",
-                    },
-                ),
-                "neg_prompt": (
-                    "STRING",
-                    {
-                        "placeholder": "Negative Prompt. (Ignored for Flux based models.)\n Example: bad composition, rating_explicit, bad quality,",
-                        "multiline": True,
-                        "tooltip": "Negative prompt. This is ignored when using flux-dev or flux-dev-uncensored or similar models that do not support CFG (Classifier-Free-Guidance)",
-                    },
-                ),
-                "width": (
-                    "INT",
-                    {
-                        "default": 1024,
-                        "min": 0,
-                        "max": 2048,  # limit is 1280 but i dont want to restrict this in case of future updates, https://docs.venice.ai/api-reference/endpoint/image/generate#body-height
-                        "step": 16,
-                        "tooltip": "Must be a multiple of 32. Maximum allowed by venice.ai at time of writing is 1280",
-                    },
-                ),
-                "height": (
-                    "INT",
-                    {
-                        "default": 1024,
-                        "min": 0,
-                        "max": 2048,
-                        "step": 16,
-                        "tooltip": "Must be a multiple of 32. Maximum allowed by venice.ai at time of writing is 1280",
-                    },
-                ),
-                "batch_size": (
-                    "INT",
-                    {
-                        "default": 1,
-                        "min": 1,
-                        "max": 4,
-                        "tooltip": "Number of images to generate in a single batch. IMPORTANT: Doesn't do actual batches like ComfyUI would, just sends batches amount of different requests to Venice.",
-                    },
-                ),
-                "steps": (
-                    "INT",
-                    {
-                        "default": 20,
-                        "min": 1,
-                        "max": 50,
-                        "tooltip": (
-                            "Number of inference steps. The following models have reduced max steps from "
-                            "the global max: venice-sd35: 30 max steps, hidream: 50 max steps, fluently-xl: 50 max steps, "
-                            "flux-dev: 30 max steps, flux-dev-uncensored: 30 max steps, getphat-flux: 50 max steps, "
-                            "lustify-sdxl: 50 max steps, pony-realism: 50 max steps, stable-diffusion-3.5: 30 max steps, "
-                            "juggernaut-xi: 50 max steps."
-                        ),
-                    },
-                ),
-                "guidance": (
-                    "FLOAT",
-                    {
-                        "default": 3.0,
-                        "min": 0.0,
-                        "max": 20.0,
-                        "step": 0.05,
-                        "tooltip": "CFG scale parameter or 'Guidance' for Flux",
-                    },
-                ),
-                # "lora_strength": ("INT", {"default": 50, "min": 0, "max": 100}), # check docs idk how to work this yet
-                "style_preset": (
-                    style_preset_options,
-                    {
-                        "default": style_preset_options[0],
-                    },
-                ),
-                "hide_watermark": (
-                    "BOOLEAN",
-                    {
-                        "default": True,
-                        "tooltip": "Whether to hide the Venice watermark. Venice may ignore this parameter for certain generated content (mainl NSFW seems like).",
-                    },
-                ),
-                "safe_mode": (
-                    "BOOLEAN",
-                    {
-                        "default": False,
-                        "tooltip": "Whether to use safe mode. If enabled, this will blur images that are classified as having adult content.",
-                    },
-                ),
-                # "format": (["png", "jpeg", "webp"], {"default": "png",}),
-            },
-            "optional": {
-                "seed": ("INT", {"default": -1, "min": -0x3B9AC9FF, "max": 0x3B9AC9FF})
-            },  # 0xffffffffffffffff is 64 bit integer limit, current hex is 999999999, venice max
-        }
+    @staticmethod
+    def _style_options() -> tuple[str, ...]:
+        options = list(style_choices())
+        if not options:
+            return ("none_available",)
+        return tuple(options)
+
+    @staticmethod
+    def _option_input_id(model_id: str, field: str) -> str:
+        sanitized = "".join(ch if ch.isalnum() or ch in {"_", "-"} else "_" for ch in model_id)
+        return f"{sanitized}__{field}"
 
-    # todo: add variants for batch size
-    # todo: implement lora and lora_strength
-    # todo: see if aspect_ratio needs to be added (some models incl nano banana pro use this)
-    # https://docs.venice.ai/api-reference/endpoint/image/generate
-    # todo: see if resolution needs to be added (some models incl nano banana pro use this)
-    # todo: add enable_web_search, mention it charges extra credits
-    # todo: change up default limits, prompt length max is 7500 now
     @classmethod
-    def VALIDATE_INPUTS(cls, input_types, **kwargs):
-        width = input_types.get("width", kwargs.get("width"))
-        height = input_types.get("height", kwargs.get("height"))
-        prompt = input_types.get("prompt", kwargs.get("prompt"))
-        neg_prompt = input_types.get("neg_prompt", kwargs.get("neg_prompt"))
+    def _get_option_value(cls, model_payload: Dict[str, Any], model_id: str, field: str) -> Any:
+        candidates = (
+            cls._option_input_id(model_id, field),
+            field,
+            f"{field}__{model_id}",
+            f"{model_id}__{field}",
+        )
+        for key in candidates:
+            if key in model_payload:
+                return model_payload.get(key)
+        return None
 
-        ensure_multiple_of(width, height)
-        ensure_prompt_length(prompt, 1500, "Prompt")
-        ensure_prompt_length(neg_prompt, 1500, "Negative Prompt", allow_empty=True)
+    @classmethod
+    def _resolve_option_value(cls, model_payload: Dict[str, Any], model_id: str, field: str, default: Any) -> Any:
+        value = cls._get_option_value(model_payload, model_id, field)
+        return default if value is None else value
 
-        return True
+    @staticmethod
+    def _coerce_positive_int(value: Any) -> int | None:
+        try:
+            candidate = int(value)
+        except (TypeError, ValueError):
+            return None
+        return candidate if candidate > 0 else None
 
-    def generate(
-        self,
+    @staticmethod
+    def _width_height_divisor(constraints: Dict[str, Any]) -> int:
+        divisor = constraints.get("widthHeightDivisor")
+        if isinstance(divisor, int) and divisor > 0:
+            return divisor
+        return 16
+
+    @classmethod
+    def _steps_limits(cls, constraints: Dict[str, Any]) -> tuple[int, int]:
+        steps = constraints.get("steps") or {}
+        default = cls._coerce_positive_int(steps.get("default"))
+        max_value = cls._coerce_positive_int(steps.get("max"))
+        default = default if default is not None else 20
+        max_steps = max_value if max_value is not None else 50
+        if default > max_steps:
+            max_steps = default
+        return default, max_steps
+
+    @classmethod
+    def _model_option_inputs(
+        cls,
+        model_id: str,
+        width_divisor: int,
+        steps_default: int,
+        steps_max: int,
+    ) -> list[io.Input]:
+        return [
+            io.Int.Input(
+                cls._option_input_id(model_id, "width"),
+                display_name="width",
+                default=1024,
+                min=0,
+                max=2048,
+                step=width_divisor,
+                tooltip="Target width for the generated image; stepping is tied to the model's `widthHeightDivisor`. Defaults to 16",
+            ),
+            io.Int.Input(
+                cls._option_input_id(model_id, "height"),
+                display_name="height",
+                default=1024,
+                min=0,
+                max=2048,
+                step=width_divisor,
+                tooltip="Target height for the generated image; stepping is tied to the model's `widthHeightDivisor`. Defaults to 16",
+            ),
+            io.Int.Input(
+                cls._option_input_id(model_id, "steps"),
+                display_name="steps",
+                default=steps_default,
+                min=1,
+                max=steps_max,
+                tooltip="Number of inference steps. Model constraints can reduce the range and have different defaults.",
+            ),
+        ]
+
+    @classmethod
+    def _build_model_options(cls) -> list[io.DynamicCombo.Option]:
+        specs = cls._image_specs(require=False)
+        options: list[io.DynamicCombo.Option] = []
+
+        def _sorted_model_items() -> list[tuple[str, Dict[str, Any]]]:
+            return sorted(specs.items())
+
+        for model_id, spec in _sorted_model_items():
+            if (
+                model_id == "nano-banana"
+            ):  # todo: implement ui for nano-banana, might be able to use code from video node
+                continue
+            constraints = spec.get("constraints") or {}
+            width_divisor = cls._width_height_divisor(constraints)
+            steps_default, steps_max = cls._steps_limits(constraints)
+            option_inputs = cls._model_option_inputs(
+                model_id,
+                width_divisor,
+                steps_default,
+                steps_max,
+            )
+            options.append(io.DynamicCombo.Option(model_id, option_inputs))
+
+        if not options:
+            option_inputs = cls._model_option_inputs(
+                "none_available",
+                16,
+                20,
+                50,
+            )
+            options.append(io.DynamicCombo.Option("none_available", option_inputs))
+
+        return options
+
+    @staticmethod
+    def _prompt_limit_from_spec(spec: Dict[str, Any] | None, default: int = 1500) -> int:
+        if not spec:
+            return default
+        constraints = spec.get("constraints") or {}
+        limit = constraints.get("promptCharacterLimit")
+        if isinstance(limit, int) and limit > 0:
+            return limit
+        try:
+            normalized = int(limit)
+        except (TypeError, ValueError):
+            return default
+        return normalized if normalized > 0 else default
+
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        model_options = cls._build_model_options()
+        style_options = cls._style_options()
+
+        return io.Schema(
+            node_id="GenerateImage_VENICE",
+            display_name="Generate Image (Venice)",
+            category="venice.ai",
+            inputs=[
+                io.String.Input(
+                    "prompt",
+                    default="A flying cat made of lettuce",
+                    multiline=True,
+                    placeholder="Positive Prompt. Example: A flying cat made of lettuce",
+                    tooltip="The text prompt to guide the image generation. Character limit depends on model (usually around 1500-7500 characters).",
+                ),
+                io.String.Input(
+                    "neg_prompt",
+                    default="",
+                    multiline=True,
+                    placeholder="Negative Prompt. Example: low quality, vacant scene",
+                    tooltip=(
+                        "Negative prompt (ignored for models that do not support CFG - z-image-turbo, flux-dev, etc.). "
+                        "Character limit depends on model (usually around 1500-7500 characters)."
+                    ),
+                ),
+                io.DynamicCombo.Input(
+                    "model",
+                    options=model_options,
+                    tooltip="Select a Venice image model to auto-populate valid parameters",
+                ),
+                io.Int.Input(
+                    "batch_size",
+                    default=1,
+                    min=1,
+                    max=4,
+                    tooltip="Number of images to generate in a single batch (sequential requests, does not use variants api (yet?)).",
+                ),
+                io.Float.Input(
+                    "guidance",
+                    default=6.0,
+                    min=0.0,
+                    max=20.0,
+                    step=0.05,
+                    tooltip=(
+                        "CFG scale (SDXL based models work well with 6.0, most newer ones work with 3-4. "
+                        "Closed Source models may ignore this setting and distilled models too, such as z-image-turbo or flux-dev and similar)."
+                    ),
+                ),
+                io.Combo.Input(
+                    "style_preset",
+                    options=list(style_options),
+                    default=style_options[0],
+                    tooltip="Venice.ai style preset to apply to the generated image.",
+                ),
+                io.Boolean.Input(
+                    "hide_watermark",
+                    default=True,
+                    tooltip="Hide the Venice watermark when possible.",
+                ),
+                io.Boolean.Input(
+                    "safe_mode",
+                    default=False,
+                    tooltip="Enable safe mode (blurs NSFW content).",
+                ),
+                io.Int.Input(
+                    "seed",
+                    optional=True,
+                    default=42,
+                    min=-0x3B9AC9FF,
+                    max=0x3B9AC9FF,
+                    tooltip="Seed for reproducibility.",
+                ),
+            ],
+            outputs=[io.Image.Output(id="image", display_name="Image")],
+        )
+
+    @classmethod
+    def execute(
+        cls,
         model,
         prompt,
         neg_prompt,
-        width,
-        height,
-        batch_size,
-        steps,
         guidance,
-        # lora_strength,
+        batch_size,
         style_preset,
         hide_watermark,
         safe_mode,
-        # format,
         seed=-1,
-    ):
-        if re.match(r"^flux.*", model):
-            logging.info(f"VeniceAPI INFO: Ignoring negative prompt for {model}.")
-            neg_prompt = ""
+    ) -> io.NodeOutput:
+        if not isinstance(model, dict) or "model" not in model:
+            raise ValueError("Model selection is required")
 
-        images_tensor = ()  # empty tuple for tensors
+        model_id = model.get("model")
+        specs = cls._image_specs(require=True)
+        spec = specs.get(model_id)
+        if not spec:
+            raise ValueError("Selected model is missing from the Venice catalog; refresh the catalog and try again.")
+
+        constraints = spec.get("constraints") or {}
+        prompt_limit = cls._prompt_limit_from_spec(spec)
+        width_height_divisor = cls._width_height_divisor(constraints)
+        steps_default, _steps_max = cls._steps_limits(constraints)
+
+        width = int(cls._resolve_option_value(model, model_id, "width", 1024))
+        height = int(cls._resolve_option_value(model, model_id, "height", 1024))
+        steps = int(cls._resolve_option_value(model, model_id, "steps", steps_default))
+        guidance = float(guidance)
+        batch_size = int(batch_size)
+        style_options = cls._style_options()
+        if style_preset not in style_options:
+            style_preset = style_options[0]
+        hide_watermark = bool(hide_watermark)
+        safe_mode = bool(safe_mode)
+        seed = seed
+
+        ensure_multiple_of(width, height, multiple=width_height_divisor)
+        ensure_prompt_length(prompt, prompt_limit, "Prompt")
+        ensure_prompt_length(neg_prompt, prompt_limit, "Negative Prompt", allow_empty=True)
+
+        seed_value = -1 if seed is None else int(seed)
+        images_tensor = ()
 
         try:
             payload = {
-                "model": model,
+                "model": model_id,
                 "prompt": prompt,
                 "negative_prompt": neg_prompt,
-                # "lora_strength": lora_strength,
                 "style_preset": style_preset,
                 "height": height,
                 "width": width,
                 "steps": steps,
                 "cfg_scale": guidance,
-                "seed": seed,
+                "seed": seed_value,
                 "return_binary": False,
                 "hide_watermark": hide_watermark,
                 "safe_mode": safe_mode,
-                "format": "png",  # hardcoded because, change to format var and uncomment related stuff above if want dynamic
-                "embed_exif_metadata": True,  # this might not work and be overriden by comfyui on image save
+                "format": "png",
+                "embed_exif_metadata": True,
             }
-            if style_preset == "none":
-                del payload["style_preset"]
+            if style_preset in ("none", "none_available"):
+                payload.pop("style_preset", None)
 
             for i in range(batch_size):
-                payload["seed"] = seed + i
+                payload["seed"] = seed_value + i
                 response_json = client.post_json(API_ENDPOINTS["image_generate"], payload)
-                images_tensor += self.process_result(response_json)
+                images_tensor += cls._processor.process_result(response_json)
 
             merged = torch.cat(images_tensor, dim=0)
-            return (merged,)
+            return io.NodeOutput(merged)
 
-        except Exception as e:
-            raise Exception(f"Error processing image result: {str(e)}") from e
-
-
-NODE_CLASS_MAPPINGS = {
-    "GenerateImage_VENICE": GenerateImage,
-}
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "GenerateImage_VENICE": "Generate Image (Venice)",
-}
+        except Exception as exc:
+            raise Exception(f"Error processing image result: {str(exc)}") from exc

nodes/gen_speech_node.py

@@ -1,176 +1,224 @@
+import logging
 import os
+import subprocess
 import tempfile
 
 import requests
 import torch
-import torchaudio
+
+from comfy_api.latest import io
+
+try:
+    import imageio_ffmpeg
+except ImportError:  # pragma: no cover
+    imageio_ffmpeg = None
 
 from ..globals import API_ENDPOINTS
 from ..nodes.catalog_utils import tts_model_choices, tts_voice_choices
 from ..nodes.utils import ensure_prompt_length
 from ..venice_client import client
 
+try:
+    from torchaudio import sox_io_backend
+except ImportError:  # pragma: no cover - default backend may not be available everywhere
+    sox_io_backend = None
 
-class GenerateSpeech:
+try:
+    from torchaudio import soundfile_backend
+except ImportError:  # pragma: no cover
+    soundfile_backend = None
+
+
+LOG = logging.getLogger(__name__)
+
+
+class GenerateSpeech(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(cls):
-        model_options = tts_model_choices()
-        voice_options = tts_voice_choices()
+    def _model_options(cls) -> list[str]:
+        options = list(tts_model_choices())
+        return options or ["none_available"]
 
-        return {
-            "required": {
-                "model": (
-                    model_options,
-                    {
-                        "default": model_options[0],
-                    },
+    @classmethod
+    def _voice_options(cls) -> list[str]:
+        options = list(tts_voice_choices())
+        return options or ["none_available"]
+
+    @staticmethod
+    def _ffmpeg_decode(temp_path: str) -> tuple[torch.Tensor, int]:
+        if imageio_ffmpeg is None:
+            LOG.error("imageio-ffmpeg is not installed; ffmpeg fallback unavailable")
+            raise RuntimeError("imageio-ffmpeg is not installed; install it to enable ffmpeg fallback.")
+
+        try:
+            ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
+        except Exception as exc:
+            LOG.error("Unable to download ffmpeg via imageio-ffmpeg: %s", exc)
+            raise RuntimeError("Failed to download ffmpeg via imageio-ffmpeg") from exc
+
+        LOG.info("Decoding Venice audio via ffmpeg executable %s", ffmpeg)
+        try:
+            process = subprocess.run(
+                [
+                    ffmpeg,
+                    "-hide_banner",
+                    "-loglevel",
+                    "error",
+                    "-i",
+                    temp_path,
+                    "-acodec",
+                    "pcm_f32le",
+                    "-f",
+                    "f32le",
+                    "-ac",
+                    "1",
+                    "-ar",
+                    "44100",
+                    "-",
+                ],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                check=True,
+            )
+        except subprocess.CalledProcessError as exc:
+            LOG.error("ffmpeg failed to decode %s: %s", temp_path, exc.stderr.decode(errors="ignore"))
+            raise RuntimeError(f"ffmpeg failed to decode {temp_path}: {exc.stderr.decode(errors='ignore')}") from exc
+
+        audio_data = torch.frombuffer(process.stdout, dtype=torch.float32).clone()
+        audio_data = audio_data.reshape(-1, 1).transpose(0, 1)
+        return audio_data, 44100
+
+    @staticmethod
+    def _load_with_torchaudio_backends(temp_path: str, response_format: str) -> tuple[torch.Tensor, int]:
+        backends = [sox_io_backend, soundfile_backend]
+        errors: list[str] = []
+        for backend in backends:
+            if backend is None:
+                continue
+            backend_name = getattr(backend, "__name__", "torchaudio_backend")
+            try:
+                return backend.load(temp_path)
+            except Exception as exc:
+                error = str(exc)
+                errors.append(error)
+                LOG.warning("torchaudio backend %s failed to load %s: %s", backend_name, response_format, error)
+        if errors:
+            raise RuntimeError(f"Failed to load audio format '{response_format}' with all methods: {', '.join(errors)}")
+        raise RuntimeError(
+            f"Failed to load audio format '{response_format}' because no torchaudio backend is available."
+        )
+
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        response_formats = ["mp3", "opus", "aac", "flac", "wav", "pcm"]
+        model_options = cls._model_options()
+        voice_options = cls._voice_options()
+
+        return io.Schema(
+            node_id="GenerateSpeech_VENICE",
+            display_name="Generate Speech (Venice)",
+            category="venice.ai",
+            inputs=[
+                io.Combo.Input(
+                    "model",
+                    options=model_options,
+                    default=model_options[0],
+                    tooltip="Model to use for text-to-speech",
                 ),
-                "input": (
-                    "STRING",
-                    {
-                        "default": "",
-                        "multiline": True,
-                        "tooltip": ("The text to generate audio for. The maximum length is 4096 characters."),
-                    },
+                io.String.Input(
+                    "input",
+                    default="",
+                    multiline=True,
+                    placeholder="Text to speak",
+                    tooltip="The text prompt used for speech generation (max 4096 chars)",
                 ),
-                "response_format": (
-                    [
-                        # todo: some dont work because idk implementing would be ass
-                        "mp3",
-                        # "opus",
-                        # "aac",
-                        # "flac",
-                        "wav",
-                        "pcm",
-                    ],
-                    {
-                        "default": "mp3",
-                        "tooltip": (
-                            "mp3: widely supported, lossy; "
-                            # "opus: very good quality at low bitrate; "
-                            # "aac: lossy, good for streaming; "
-                            # "flac: lossless compressed audio; "
-                            "wav: lossless raw audio; "
-                            "pcm: uncompressed raw audio."
-                        ),
-                    },
+                io.Combo.Input(
+                    "response_format",
+                    options=response_formats,
+                    default=response_formats[0],
+                    tooltip="Audio format to request from the Venice TTS API",
                 ),
-                "speed": (
-                    "FLOAT",
-                    {
-                        "default": 1.0,
-                        "min": 0.25,
-                        "max": 4,
-                        "step": 0.01,
-                        "tooltip": (
-                            "The text to image style to apply during prompt enhancement. "
-                            "Does best with short descriptive prompts, like gold, marble or angry, menacing."
-                        ),
-                    },
+                io.Float.Input(
+                    "speed",
+                    default=1.0,
+                    min=0.25,
+                    max=4.0,
+                    step=0.01,
+                    tooltip="Playback speed multiplier (1.0 = normal speed)",
                 ),
-                # "streaming": (
-                #     "BOOLEAN",
-                #     {
-                #         "default": False,
-                #         "tooltip": (
-                #             "Should the content stream back sentence by sentence or be processed and returned as a complete audio file."
-                #
-                # ),
-                #     },
-                # ),
-                "voice": (
-                    voice_options,
-                    {
-                        "default": voice_options[0],
-                    },
+                io.Combo.Input(
+                    "voice",
+                    options=voice_options,
+                    default=voice_options[0],
+                    tooltip="Voice preset to use for the TTS model",
                 ),
-            }
-        }
+            ],
+            outputs=[io.Audio.Output(id="audio", display_name="audio")],
+        )
 
-    RETURN_TYPES = ("AUDIO",)
-    RETURN_NAMES = ("audio",)
-    FUNCTION = "gen_speech"
-    CATEGORY = "venice.ai"
+    @classmethod
+    def execute(cls, model, input, response_format, speed, voice) -> io.NodeOutput:
+        ensure_prompt_length(input, 4096, label="TTS input")
 
-    EXPERIMENTAL = True
-
-    def gen_speech(self, model, input, response_format, speed, voice):
-        ensure_prompt_length(input, 4096, label="Speech input")
-
-        url = API_ENDPOINTS["speech_generate"]
-
-        # remove everything from voice string after and including the hyphen " - blabla"
-        voice = voice.split(" - ")[0] if " - " in voice else voice
-
-        # Prepare JSON payload
+        # todo: currently models' voices show up as "model-name - voice_name"
+        # todo: this can be in dynamiccombo so the single combo dropdown is not cluttered with all the voices of all selectable models
+        normalized_voice = voice.split(" - ")[-1].strip() if " - " in voice else voice
         payload = {
             "model": model,
             "input": input,
             "speed": speed,
-            "voice": voice,
+            "voice": normalized_voice,
             "response_format": response_format,
             "streaming": False,
         }
 
-        # Send request
         try:
             response = client.request(
                 "POST",
-                url,
+                API_ENDPOINTS["speech_generate"],
                 json=payload,
                 headers={"Content-Type": "application/json"},
             )
-        except requests.exceptions.RequestException as e:
-            raise RuntimeError(f"Generate Speech (Venice) API request failed: {str(e)}")
+        except requests.exceptions.RequestException as exc:
+            raise RuntimeError(f"Generate Speech (Venice) API request failed: {str(exc)}")
 
-        # Convert the audio response to a ComfyUI-compatible tensor and return it
-        if not response.content or len(response.content) == 0:
+        if not response.content:
             raise RuntimeError("No audio data received from Venice API.")
 
-        # Save to temporary file and load with torchaudio for better format support
-        with tempfile.NamedTemporaryFile(suffix=f".{response_format}", delete=False) as temp_file:
-            temp_file.write(response.content)
-            temp_file_path = temp_file.name
-
+        temp_path = None
+        waveform = None
+        sample_rate = None
         try:
-            # Load audio using torchaudio from the temporary file
-            waveform, sample_rate = torchaudio.load(temp_file_path)
-        except Exception as e:
-            # Fallback: try different approaches for problematic formats
+            # Save to temp file and load with torchaudio for better format support
+            with tempfile.NamedTemporaryFile(suffix=f".{response_format}", delete=False) as temp_file:
+                temp_file.write(response.content)
+                temp_path = temp_file.name
+
             try:
-                if response_format == "pcm":
-                    # For PCM, we need to handle it as raw audio data
-                    # Assume 16-bit PCM, mono, 16kHz (adjust as needed based on API response)
+                waveform, sample_rate = cls._load_with_torchaudio_backends(temp_path, response_format)
+            except RuntimeError as audio_exc:
+                LOG.warning("torchaudio decoding failed for %s: %s", response_format, audio_exc)
+                if response_format == "pcm" and response.content:
+                    LOG.info("Falling back to raw PCM interpretation for %s", response_format)
                     audio_data = torch.frombuffer(response.content, dtype=torch.int16).float() / 32768.0
-                    waveform = audio_data.unsqueeze(0)  # Add channel dimension
-                    sample_rate = 16000  # Default sample rate, adjust if needed
+                    waveform = audio_data.unsqueeze(0)
+                    sample_rate = 16000  # Default sample rate, might need change?
                 else:
-                    # For other formats, try loading without specifying format
-                    waveform, sample_rate = torchaudio.load(temp_file_path, format=None)
-            except Exception as e2:
-                raise RuntimeError(
-                    f"Failed to load audio format '{response_format}' with all methods. Errors: {str(e)}, {str(e2)}"
-                )
+                    LOG.info("Attempting ffmpeg fallback for %s audio", response_format)
+                    waveform, sample_rate = cls._ffmpeg_decode(temp_path)
         finally:
-            # Clean up temporary file
-            try:
-                os.unlink(temp_file_path)
-            except:
-                pass
+            if temp_path:
+                try:
+                    os.unlink(temp_path)
+                except OSError:
+                    pass
+
+        if waveform is None or sample_rate is None:
+            raise RuntimeError("Unable to decode Venice speech response.")
 
-        # Ensure shape is [B, C, T] (batch size 1)
         if waveform.dim() == 2:
             waveform = waveform.unsqueeze(0)  # [1, C, T]
         elif waveform.dim() == 1:
             waveform = waveform.unsqueeze(0).unsqueeze(0)  # [1, 1, T]
 
-        return ({"waveform": waveform, "sample_rate": sample_rate},)
-
-
-NODE_CLASS_MAPPINGS = {
-    "GenerateSpeech_VENICE": GenerateSpeech,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "GenerateSpeech_VENICE": "Generate Speech [BETA] (Venice)",
-}
+        audio_value: io.Audio.Type = {"waveform": waveform, "sample_rate": sample_rate}
+        return io.NodeOutput(audio_value)

nodes/gen_text_advanced_node.py

@@ -1,287 +0,0 @@
-from ..globals import API_ENDPOINTS
-from ..nodes.catalog_utils import text_model_choices
-from ..nodes.utils import encode_tensor_for_vision, ensure_prompt_length
-from ..venice_client import client
-
-
-class GenerateTextAdvanced:
-    @classmethod
-    def INPUT_TYPES(cls):
-        model_options = text_model_choices()
-        return {
-            "required": {
-                "model": (
-                    model_options,
-                    {
-                        "default": model_options[0],
-                        "tooltip": ("The model to use for text generation."),
-                    },
-                ),
-                "prompt": (
-                    "STRING",
-                    {
-                        "default": "",
-                        "multiline": True,
-                        "tooltip": ("The prompt to generate text from. Ask, command or chat with the model."),
-                    },
-                ),
-                "system_prompt": (
-                    "STRING",
-                    {
-                        "default": "",
-                        "multiline": True,
-                        "tooltip": ("Optional system prompt to guide the model's behavior."),
-                    },
-                ),
-                "enable_system_prompt": (
-                    "BOOLEAN",
-                    {
-                        "default": True,
-                        "tooltip": ("Enable or disable system prompt being passed on."),
-                    },
-                ),
-                "frequency_penalty": (
-                    "FLOAT",
-                    {
-                        "default": 0.0,
-                        "min": -2.0,
-                        "max": 2.0,
-                        "step": 0.05,
-                        "tooltip": (
-                            "Positive values penalize new tokens based on their existing frequency in the text so far, "
-                            "decreasing the model's likelihood to repeat the same line verbatim."
-                        ),
-                    },
-                ),
-                "presence_penalty": (
-                    "FLOAT",
-                    {
-                        "default": 0.0,
-                        "min": -2.0,
-                        "max": 2.0,
-                        "step": 0.05,
-                        "tooltip": (
-                            "Positive values penalize new tokens based on whether they appear in the text so far, "
-                            "increasing the model's likelihood to talk about new topics."
-                        ),
-                    },
-                ),
-                "repetition_penalty": (
-                    "FLOAT",
-                    {
-                        "default": 1.2,
-                        "min": 0.0,
-                        "max": 2.0,
-                        "step": 0.05,
-                        "tooltip": ("1.0 means no penalty. Values > 1.0 discourage repetition."),
-                    },
-                ),
-                "max_temp": (
-                    "FLOAT",
-                    {
-                        "default": 1.5,
-                        "min": 0.0,
-                        "max": 2.0,
-                        "step": 0.05,
-                        "tooltip": ("Maximum temperature value for dynamic temperature scaling."),
-                    },
-                ),
-                "min_temp": (
-                    "FLOAT",
-                    {
-                        "default": 0.1,
-                        "min": 0.0,
-                        "max": 2.0,
-                        "step": 0.05,
-                        "tooltip": ("Minimum temperature value for dynamic temperature scaling."),
-                    },
-                ),
-                "max_completion_tokens": (
-                    "INT",
-                    {
-                        "default": 420,
-                        "min": 1,
-                        "max": 131072,
-                        "step": 1,
-                        "tooltip": (
-                            "An upper bound for the number of tokens that can be generated for "
-                            "a completion, including visible output tokens and reasoning tokens."
-                        ),
-                    },
-                ),
-                "temperature": (
-                    "FLOAT",
-                    {
-                        "default": 0.5,
-                        "min": 0.0,
-                        "max": 2.0,
-                        "step": 0.05,
-                        "tooltip": (
-                            "Higher values like 0.8 will make the output more random, "
-                            "while lower values like 0.2 will make it more focused and deterministic. "
-                            "We generally recommend altering this or top_p but not both."
-                        ),
-                    },
-                ),
-                "top_k": (
-                    "INT",
-                    {
-                        "default": 40,
-                        "min": 0,
-                        "tooltip": ("The number of highest probability vocabulary tokens to keep for top-k-filtering."),
-                    },
-                ),
-                "top_p": (
-                    "FLOAT",
-                    {
-                        "default": 0.8,
-                        "min": 0.0,
-                        "max": 2.0,
-                        "step": 0.01,
-                        "tooltip": (
-                            "An alternative to sampling with temperature, called nucleus sampling, "
-                            "where the model considers the results of the tokens with top_p probability mass. "
-                            "So 0.1 means only the tokens comprising the top 10% probability mass are considered."
-                        ),
-                    },
-                ),
-                "min_p": (
-                    "FLOAT",
-                    {
-                        "default": 0.05,
-                        "min": 0.0,
-                        "max": 1.0,
-                        "step": 0.01,
-                        "tooltip": (
-                            "Sets a minimum probability threshold for token selection. "
-                            "Tokens with probabilities below this value are filtered out."
-                        ),
-                    },
-                ),
-                # "stop": ("STRING", {"default": "", "tooltip": "Up to 4 sequences where the API will stop generating further tokens. Defaults to null.", "placeholder": "stop: [\"\\n\"]"}),
-                # "stop_token_ids": ("STRING", {"default": "", "tooltip": "Array of token IDs where the API will stop generating further tokens. Example: [151643, 151645]", "placeholder": "151643, 151645, ..."}),
-                "enable_vision": (
-                    "BOOLEAN",
-                    {
-                        "default": False,
-                        "tooltip": (
-                            "Enable or disable vision tasks. "
-                            "Requires image_for_vision input to be populated and "
-                            "for the LLM to actually support vision tasks to process."
-                        ),
-                    },
-                ),
-            },
-            "optional": {
-                "venice_parameters": (
-                    "STRING",
-                    {
-                        "forceInput": True,
-                        "tooltip": (
-                            "Optional input. "
-                            "Use the Textgen Parameters (Venice) node to use "
-                            "extra, venice specific parameters for text generation."
-                        ),
-                    },
-                ),
-                "image_for_vision": (
-                    "IMAGE",
-                    {
-                        "tooltip": (
-                            "Optional input. "
-                            "Add an image for vision-supported LLMs to process. "
-                            "Will only be processed if 'enable_vision' is 'True' and "
-                            "if the LLM actually supports vision tasks."
-                        ),
-                    },
-                ),
-            },
-        }
-
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("response",)
-    FUNCTION = "generate_text"
-    CATEGORY = "venice.ai"
-    DESCRIPTION = (
-        "Text Generation node that makes use of Venice.AI's text generation API. "
-        "Use Textgen Parameters (Venice) node to pass on extra Venice.AI specific parameters."
-        "Does not have chat history context. "
-    )
-
-    def generate_text(
-        # region params
-        self,
-        model,
-        prompt,
-        system_prompt,
-        enable_system_prompt,
-        frequency_penalty,
-        presence_penalty,
-        repetition_penalty,
-        max_temp,
-        min_temp,
-        max_completion_tokens,
-        temperature,
-        top_k,
-        top_p,
-        min_p,
-        enable_vision,
-        **kwargs,
-        # endregion
-    ):
-
-        ensure_prompt_length(prompt, 1500, label="Prompt")
-        user_content = []
-        venice_parameters = kwargs.get("venice_parameters")
-        image_for_vision = kwargs.get("image_for_vision")
-
-        if image_for_vision is not None and enable_vision:
-            encoded_image = encode_tensor_for_vision(image_for_vision[0])
-            user_content.extend(
-                [
-                    {"type": "text", "text": prompt},
-                    {"type": "image_url", "image_url": {"url": encoded_image}},
-                ]
-            )
-        else:
-            user_content.append({"type": "text", "text": prompt})
-
-        if not enable_system_prompt:
-            system_prompt = ""
-
-        messages = [{"role": "system", "content": system_prompt}]
-        messages.append({"role": "user", "content": user_content})
-
-        payload = {
-            "model": model,
-            "messages": messages,
-            "frequency_penalty": frequency_penalty,
-            "presence_penalty": presence_penalty,
-            "repetition_penalty": repetition_penalty,
-            "max_temp": max_temp,
-            "min_temp": min_temp,
-            "max_completion_tokens": max_completion_tokens,
-            "temperature": temperature,
-            "top_k": top_k,
-            "top_p": top_p,
-            "min_p": min_p,
-        }
-        if venice_parameters is not None:
-            payload["venice_parameters"] = venice_parameters
-
-        json_response = client.post_json(API_ENDPOINTS["text_generate"], payload)
-        try:
-            content = json_response["choices"][0]["message"]["content"]
-        except (KeyError, IndexError, TypeError) as e:
-            raise ValueError(f"Unexpected API response format: {json_response}") from e
-
-        return (content,)
-
-
-NODE_CLASS_MAPPINGS = {
-    "GenerateTextAdvanced_VENICE": GenerateTextAdvanced,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "GenerateTextAdvanced_VENICE": "Generate Text Advanced BETA (Venice)",
-}

nodes/gen_text_node.py

@@ -1,58 +1,438 @@
+import json
+import logging
+from typing import Any, Dict, Iterable
+
+from comfy_api.latest import io
+
 from ..globals import API_ENDPOINTS
-from ..nodes.catalog_utils import text_model_choices
-from ..nodes.utils import encode_tensor_for_vision, ensure_prompt_length
+from ..nodes.catalog_utils import character_choices, text_model_specs
+from ..nodes.utils import encode_tensor_for_vision
 from ..venice_client import client
 
+logger = logging.getLogger(__name__)
 
-class GenerateText:
+
+class GenerateTextAdvanced(io.ComfyNode):
     @classmethod
-    def INPUT_TYPES(cls):
-        model_choices = text_model_choices()
+    def _option_input_id(cls, model_id: str, field: str) -> str:
+        return f"{model_id}__{field}"
 
-        return {
-            "required": {
-                "model": (
-                    model_choices,
-                    {
-                        "default": model_choices[0],
-                    },
+    @staticmethod
+    def _constraint_default(value: Any, fallback: float) -> float:
+        if isinstance(value, dict):
+            default = value.get("default")
+        else:
+            default = value
+        if isinstance(default, (int, float)):
+            return float(default)
+        return fallback
+
+    @classmethod
+    def _get_option_value(cls, model_payload: Dict[str, Any], model_id: str, field: str) -> Any:
+        candidates = (
+            cls._option_input_id(model_id, field),
+            field,
+            f"{field}__{model_id}",
+            f"{model_id}__{field}",
+        )
+        for key in candidates:
+            if key in model_payload:
+                return model_payload.get(key)
+        return None
+
+    @classmethod
+    def _text_specs(cls) -> Dict[str, Dict[str, Any]]:
+        specs = text_model_specs() or {}
+        if not specs:
+            raise ValueError("No Venice text model specs available")
+        return specs
+
+    @staticmethod
+    def _normalize_stop_tokens(value: str | Iterable[str] | None) -> list[str]:
+        tokens: list[str] = []
+        if not value:
+            return tokens
+        segments: Iterable[str] = value.splitlines() if isinstance(value, str) else value
+        for segment in segments:
+            for raw_token in str(segment).split(","):
+                trimmed = raw_token.strip()
+                if trimmed:
+                    tokens.append(trimmed)
+        return tokens
+
+    @classmethod
+    def _build_model_options(cls) -> list[io.DynamicCombo.Option]:
+        specs = cls._text_specs()
+        options: list[io.DynamicCombo.Option] = []
+        for model_id, spec in sorted(specs.items(), key=lambda item: item[0]):
+            constraints = spec.get("constraints") or {}
+            capabilities = spec.get("capabilities") or {}
+            temperature_default = cls._constraint_default(constraints.get("temperature"), 0.5)
+            top_p_default = cls._constraint_default(constraints.get("top_p"), 0.8)
+
+            option_inputs: list[io.Input] = [
+                io.Float.Input(
+                    cls._option_input_id(model_id, "temperature"),
+                    display_name="temperature",
+                    default=temperature_default,
+                    min=0.0,
+                    max=2.0,
+                    step=0.01,
+                    tooltip="Sampling temperature (per-model default taken from the catalog).",
                 ),
-                "system_prompt": ("STRING", {"default": "", "multiline": True}),
-                "prompt": ("STRING", {"default": "", "multiline": True}),
-                "frequency_penalty": ("FLOAT", {"default": 1.5, "min": 0.0, "max": 2.0, "step": 0.1}),
-                "presence_penalty": ("FLOAT", {"default": 1.5, "min": 0.0, "max": 2.0, "step": 0.1}),
-                "temperature": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.1}),
-                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.1}),
-                "enable_vision": ("BOOLEAN", {"default": False}),
-            },
-            "optional": {
-                "image_for_vision": ("IMAGE",),
-            },
-        }
+                io.Float.Input(
+                    cls._option_input_id(model_id, "top_p"),
+                    display_name="top_p",
+                    default=top_p_default,
+                    min=0.0,
+                    max=1.0,
+                    step=0.01,
+                    tooltip="Nucleus sampling probability (per-model default taken from the catalog).",
+                ),
+            ]
 
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("response",)
-    FUNCTION = "generate_text"
-    CATEGORY = "venice.ai"
+            if capabilities.get("supportsVision"):
+                option_inputs.append(
+                    io.Boolean.Input(
+                        cls._option_input_id(model_id, "enable_vision"),
+                        display_name="enable_vision",
+                        default=False,
+                        tooltip="Enable vision inputs when the model supports vision.",
+                    )
+                )
 
-    def generate_text(
-        self,
+            if capabilities.get("supportsReasoning"):
+                option_inputs.append(
+                    io.Boolean.Input(
+                        cls._option_input_id(model_id, "reasoning"),
+                        display_name="reasoning",
+                        default=True,
+                        tooltip="Toggle reasoning capabilities for this model.",
+                    )
+                )
+
+            option_inputs.extend(cls._venice_parameter_inputs(model_id, capabilities))
+
+            options.append(io.DynamicCombo.Option(model_id, option_inputs))
+
+        return options
+
+    @classmethod
+    def _venice_parameter_inputs(cls, model_id: str, capabilities: Dict[str, Any]) -> list[io.Input]:
+        character_options = ["none", *character_choices()]
+        inputs: list[io.Input] = [
+            io.Combo.Input(
+                cls._option_input_id(model_id, "vp_character_slug"),
+                display_name="vp_character_slug",
+                options=character_options,
+                default=character_options[0],
+                optional=True,
+                tooltip="Select a Venice character slug (public ID) from the catalog.",
+            ),
+        ]
+
+        if capabilities.get("supportsReasoning"):
+            inputs.extend(
+                [
+                    io.Boolean.Input(
+                        cls._option_input_id(model_id, "vp_strip_thinking_response"),
+                        display_name="vp_strip_thinking_response",
+                        default=False,
+                        tooltip="Strip thinking blocks from the response on reasoning models.",
+                    ),
+                    io.Boolean.Input(
+                        cls._option_input_id(model_id, "vp_disable_thinking"),
+                        display_name="vp_disable_thinking",
+                        default=False,
+                        tooltip="Disable thinking blocks for supported reasoning models.",
+                    ),
+                ]
+            )
+
+        if capabilities.get("supportsWebSearch"):
+            inputs.extend(
+                [
+                    io.Combo.Input(
+                        cls._option_input_id(model_id, "vp_enable_web_search"),
+                        display_name="vp_enable_web_search",
+                        options=["auto", "off", "on"],
+                        default="off",
+                        optional=True,
+                        tooltip="Set to auto/off/on to control Venice web search for this request.",
+                    ),
+                    io.Boolean.Input(
+                        cls._option_input_id(model_id, "vp_enable_web_scraping"),
+                        display_name="vp_enable_web_scraping",
+                        default=False,
+                        tooltip="Enable Venice web scraping for URLs found in the latest user message.",
+                    ),
+                    io.Boolean.Input(
+                        cls._option_input_id(model_id, "vp_enable_web_citations"),
+                        display_name="vp_enable_web_citations",
+                        default=False,
+                        tooltip="Request citations when web search returns sources.",
+                    ),
+                ]
+            )
+
+        inputs.append(
+            io.Boolean.Input(
+                cls._option_input_id(model_id, "vp_include_venice_system_prompt"),
+                display_name="vp_include_venice_system_prompt",
+                default=False,
+                tooltip="Include Venice-supplied system prompts alongside your own.",
+            )
+        )
+
+        return inputs
+
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        model_options = cls._build_model_options()
+
+        return io.Schema(
+            node_id="GenerateTextAdvanced_VENICE",
+            display_name="Generate Text Advanced (Venice)",
+            category="venice.ai",
+            inputs=[
+                io.String.Input(
+                    "prompt",
+                    default="",
+                    multiline=True,
+                    tooltip="The prompt to generate text from. Ask, command or chat with the model.",
+                ),
+                io.String.Input(
+                    "system_prompt",
+                    default="",
+                    multiline=True,
+                    tooltip="Optional system prompt to guide the model's behavior.",
+                ),
+                io.Image.Input(
+                    "vision_image",
+                    display_name="vision_image",
+                    optional=True,
+                    tooltip="Optional image for vision-capable models. Enable the vision toggle to send it.",
+                ),
+                io.DynamicCombo.Input(
+                    "model",
+                    options=model_options,
+                    tooltip="The model to use for text generation.",
+                ),
+                io.Float.Input(
+                    "frequency_penalty",
+                    default=0.0,
+                    min=-2.0,
+                    max=2.0,
+                    step=0.05,
+                    tooltip=(
+                        "Positive values penalize new tokens based on their existing frequency in the text so far, "
+                        "decreasing the model's likelihood to repeat the same line verbatim."
+                    ),
+                ),
+                io.Float.Input(
+                    "presence_penalty",
+                    default=0.0,
+                    min=-2.0,
+                    max=2.0,
+                    step=0.05,
+                    tooltip=(
+                        "Positive values penalize new tokens based on whether they appear in the text so far, "
+                        "increasing the model's likelihood to talk about new topics."
+                    ),
+                ),
+                io.Float.Input(
+                    "repetition_penalty",
+                    default=1.2,
+                    min=0.0,
+                    max=2.0,
+                    step=0.05,
+                    tooltip="1.0 means no penalty. Values > 1.0 discourage repetition.",
+                ),
+                io.Float.Input(
+                    "max_temp",
+                    default=1.5,
+                    min=0.0,
+                    max=2.0,
+                    step=0.05,
+                    tooltip="Maximum temperature value for dynamic temperature scaling.",
+                ),
+                io.Float.Input(
+                    "min_temp",
+                    default=0.1,
+                    min=0.0,
+                    max=2.0,
+                    step=0.05,
+                    tooltip="Minimum temperature value for dynamic temperature scaling.",
+                ),
+                io.Int.Input(
+                    "max_completion_tokens",
+                    default=420,
+                    min=1,
+                    max=131072,
+                    step=1,
+                    tooltip=(
+                        "An upper bound for the number of tokens that can be generated for "
+                        "a completion, including visible output tokens and reasoning tokens."
+                    ),
+                ),
+                io.Int.Input(
+                    "top_k",
+                    default=40,
+                    min=0,
+                    tooltip="The number of highest probability vocabulary tokens to keep for top-k-filtering.",
+                ),
+                io.Float.Input(
+                    "min_p",
+                    default=0.05,
+                    min=0.0,
+                    max=1.0,
+                    step=0.01,
+                    tooltip=(
+                        "Sets a minimum probability threshold for token selection. "
+                        "Tokens with probabilities below this value are filtered out."
+                    ),
+                ),
+                io.String.Input(
+                    "stop_tokens",
+                    default="",
+                    tooltip="Optional comma- or newline-separated tokens to stop generation on (requires at least one).",
+                ),
+                    io.Int.Input(
+                        "seed",
+                        default=42,
+                        min=1,
+                        tooltip="Seed for Venice randomness; must be 1 or greater.",
+                    ),
+                io.Boolean.Input(
+                    "enable_system_prompt",
+                    default=True,
+                    tooltip="Enable or disable system prompt being passed on.",
+                ),
+                io.Boolean.Input(
+                    "debug_append_response",
+                    default=False,
+                    tooltip="Append the raw Venice response after three newlines for debugging.",
+                ),
+            ],
+            outputs=[io.String.Output(id="response", display_name="response")],
+        )
+
+    @classmethod
+    def execute(
+        cls,
         model,
-        system_prompt,
         prompt,
+        system_prompt,
+        vision_image,
         frequency_penalty,
         presence_penalty,
-        temperature,
-        top_p,
-        enable_vision,
-        **kwargs,
-    ):
-        ensure_prompt_length(prompt, 1500, label="Prompt")
-        user_content = []
-        image_for_vision = kwargs.get("image_for_vision")
+        repetition_penalty,
+        max_temp,
+        min_temp,
+        max_completion_tokens,
+        top_k,
+        min_p,
+        stop_tokens,
+        seed,
+        enable_system_prompt,
+        debug_append_response,
+    ) -> io.NodeOutput:
+        if isinstance(model, str):
+            model = {"model": model}
+        if not isinstance(model, dict) or "model" not in model:
+            raise ValueError("Model selection is required")
 
-        if image_for_vision is not None and enable_vision:
-            encoded_image = encode_tensor_for_vision(image_for_vision[0])
+        model_id = model.get("model")
+        specs = cls._text_specs()
+        spec = specs.get(model_id)
+        if not spec:
+            raise ValueError("Selected model is missing from the Venice catalog; refresh the catalog and try again.")
+
+        constraints = spec.get("constraints") or {}
+        capabilities = spec.get("capabilities") or {}
+
+        temperature_value = cls._get_option_value(model, model_id, "temperature")
+        if temperature_value is None:
+            temperature_value = cls._constraint_default(constraints.get("temperature"), 0.5)
+        try:
+            temperature_value = float(temperature_value)
+        except (TypeError, ValueError):
+            temperature_value = 0.5
+
+        top_p_value = cls._get_option_value(model, model_id, "top_p")
+        if top_p_value is None:
+            top_p_value = cls._constraint_default(constraints.get("top_p"), 0.8)
+        try:
+            top_p_value = float(top_p_value)
+        except (TypeError, ValueError):
+            top_p_value = 0.8
+
+        reasoning_value = cls._get_option_value(model, model_id, "reasoning")
+        reasoning_enabled = (
+            bool(reasoning_value) if reasoning_value is not None else bool(capabilities.get("supportsReasoning"))
+        )
+        reasoning_effort_value = "medium"
+        vision_enabled = (
+            bool(cls._get_option_value(model, model_id, "enable_vision"))
+            if capabilities.get("supportsVision")
+            else False
+        )
+        vision_tensor = None
+        if vision_image is not None:
+            candidates = vision_image if isinstance(vision_image, (list, tuple)) else (vision_image,)
+            for candidate in candidates:
+                if candidate is not None:
+                    vision_tensor = candidate
+                    break
+        normalized_stop_tokens = cls._normalize_stop_tokens(stop_tokens)
+        try:
+            seed_value = int(seed)
+        except (TypeError, ValueError):
+            seed_value = 42
+        if seed_value < 1:
+            seed_value = 1
+
+        venice_parameters: Dict[str, Any] = {}
+
+        def _set_bool(field: str, key: str) -> None:
+            value = cls._get_option_value(model, model_id, field)
+            if value is not None:
+                venice_parameters[key] = bool(value)
+
+        slug_value = cls._get_option_value(model, model_id, "vp_character_slug")
+        if isinstance(slug_value, str):
+            trimmed = slug_value.strip()
+            if trimmed and trimmed.lower() not in {"", "none"}:
+                venice_parameters["character_slug"] = trimmed
+
+        if capabilities.get("supportsReasoning"):
+            _set_bool("vp_strip_thinking_response", "strip_thinking_response")
+            _set_bool("vp_disable_thinking", "disable_thinking")
+
+        if capabilities.get("supportsWebSearch"):
+            web_search = cls._get_option_value(model, model_id, "vp_enable_web_search")
+            if isinstance(web_search, str):
+                trimmed = web_search.strip()
+                if trimmed:
+                    venice_parameters["enable_web_search"] = trimmed
+            elif web_search is not None:
+                venice_parameters["enable_web_search"] = str(web_search)
+            _set_bool("vp_enable_web_scraping", "enable_web_scraping")
+            _set_bool("vp_enable_web_citations", "enable_web_citations")
+
+        _set_bool("vp_include_venice_system_prompt", "include_venice_system_prompt")
+
+        if vision_tensor is not None and (not capabilities.get("supportsVision") or not vision_enabled):
+            logger.warning(
+                "Vision image provided but model %s does not support vision, or enable_vision is disabled",
+                model_id,
+            )
+        if vision_enabled and vision_tensor is None:
+            raise ValueError("Vision input is enabled but no image was provided")
+
+        user_content = []
+        if vision_enabled and vision_tensor is not None:
+            encoded_image = encode_tensor_for_vision(vision_tensor)
             user_content.extend(
                 [
                     {"type": "text", "text": prompt},
@@ -62,28 +442,57 @@ class GenerateText:
         else:
             user_content.append({"type": "text", "text": prompt})
 
+        if not enable_system_prompt:
+            system_prompt = ""
+
         messages = [{"role": "system", "content": system_prompt}]
         messages.append({"role": "user", "content": user_content})
 
         payload = {
-            "model": model,
+            "model": model_id,
             "messages": messages,
             "frequency_penalty": frequency_penalty,
+            "logprobs": False,  # unused, not supported by all models
+            "top_logprobs": 0,  # x >= 0
+            "max_completion_tokens": max_completion_tokens,
+            "max_temp": max_temp,
+            "min_p": min_p,
+            "min_temp": min_temp,
+            "n": 1,  # basically batch size
             "presence_penalty": presence_penalty,
-            "temperature": temperature,
-            "top_p": top_p,
+            "repetition_penalty": repetition_penalty,
+            "seed": seed_value,
+            "stream": False,
+            "temperature": temperature_value,
+            "top_k": top_k,
+            "top_p": top_p_value,
+            "parallel_tool_calls": True,
         }
+        if reasoning_enabled:
+            payload["reasoning"] = {"mode": reasoning_effort_value}
+            payload["reasoning_effort"] = reasoning_effort_value
+        if normalized_stop_tokens:
+            payload["stop"] = normalized_stop_tokens
+        if venice_parameters:
+            payload["venice_parameters"] = venice_parameters
 
         json_response = client.post_json(API_ENDPOINTS["text_generate"], payload)
-        content = json_response["choices"][0]["message"]["content"]
-        # print(content)
-        return (content,)
+        try:
+            choice = json_response["choices"][0]
+            logger.info(
+                "Venice LLM finish_reason=%s native_finish_reason=%s stop_reason=%s",
+                choice.get("finish_reason"),
+                choice.get("native_finish_reason"),
+                choice.get("stop_reason"),
+            )
+            content = choice["message"]["content"]
+        except (KeyError, IndexError, TypeError) as exc:
+            raise ValueError(f"Unexpected API response format: {json_response}") from exc
+        if debug_append_response:
+            try:
+                raw_dump = json.dumps(json_response, indent=2)
+            except (TypeError, ValueError):
+                raw_dump = str(json_response)
+            content = f"{content}\n\n\n{raw_dump}"
 
-
-NODE_CLASS_MAPPINGS = {
-    "GenerateText_VENICE": GenerateText,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "GenerateText_VENICE": "Generate Text (Venice)",
-}
+        return io.NodeOutput(content)

nodes/gen_text_venice_params_node.py

@@ -1,123 +0,0 @@
-from ..nodes.catalog_utils import character_choices
-
-
-class GenerateTextVeniceParameters:
-    @classmethod
-    def INPUT_TYPES(cls):
-        character_options = character_choices()
-
-        return {
-            "required": {
-                "character_slug": (
-                    character_options,
-                    {
-                        "default": character_options[0],
-                        "tooltip": ("The character slug of a public Venice character."),
-                    },
-                ),
-                "enable_character": (
-                    "BOOLEAN",
-                    {
-                        "default": False,
-                        "tooltip": ("Enable or disable character parameter being passed on."),
-                    },
-                ),
-                "strip_thinking_response": (
-                    "BOOLEAN",
-                    {
-                        "default": False,
-                        "tooltip": (
-                            "Strip <think></think> blocks from the response. "
-                            "Applicable only to reasoning / thinking models. "
-                            "Also available to use as a model feature suffix. "
-                            "Defaults to false."
-                        ),
-                    },
-                ),
-                "disable_thinking": (
-                    "BOOLEAN",
-                    {
-                        "default": False,
-                        "tooltip": (
-                            "On supported reasoning models, will disable thinking and "
-                            "strip the <think></think> blocks from the response. "
-                            "Defaults to false."
-                        ),
-                    },
-                ),
-                "web_search": (
-                    ["auto", "on", "off"],
-                    {
-                        "default": "auto",
-                        "tooltip": (
-                            "Auto will enable it based on the model's discretion. "
-                            "On will force web search on the request. "
-                            "Citations will be returned either in the first chunk of a "
-                            "streaming result, or in the non streaming response."
-                            "Defaults to off. "
-                        ),
-                    },
-                ),
-                "enable_web_citations": (
-                    "BOOLEAN",
-                    {
-                        "default": False,
-                        "tooltip": (
-                            "When web search is enabled, this will request that the LLM cite "
-                            "its sources using a [REF]0[/REF] format. "
-                            "Defaults to false."
-                        ),
-                    },
-                ),  # NOTE: include_search_results_in_stream not implemented
-                "use_venice_system_prompt": (
-                    "BOOLEAN",
-                    {
-                        "default": False,
-                        "tooltip": (
-                            "Whether to include the Venice supplied system prompts "
-                            "alongside specified system prompts. "
-                            "Defaults to true."
-                        ),
-                    },
-                ),
-            },
-        }
-
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("venice_parameters",)
-    FUNCTION = "pass_on_params"
-    CATEGORY = "venice.ai"
-    DESCRIPTION = (
-        "Passes on parameters unique parameters to Venice's API implementation for Venice text generation nodes."
-    )
-
-    def pass_on_params(
-        self,
-        character_slug,
-        enable_character,
-        strip_thinking_response,
-        disable_thinking,
-        web_search,
-        enable_web_citations,
-        use_venice_system_prompt,
-    ):
-        venice_params = {
-            "strip_thinking_response": strip_thinking_response,
-            "disable_thinking": disable_thinking,
-            "enable_web_search": web_search,
-            "enable_web_citations": enable_web_citations,
-            "include_venice_system_prompt": use_venice_system_prompt,
-        }
-        if enable_character:
-            venice_params["character_slug"] = character_slug
-
-        return (venice_params,)
-
-
-NODE_CLASS_MAPPINGS = {
-    "GenerateTextVeniceParameters_VENICE": GenerateTextVeniceParameters,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "GenerateTextVeniceParameters_VENICE": "Textgen Parameters (Venice)",
-}

nodes/gen_video_from_text_node.py

@@ -1,151 +1,241 @@
+import logging
+from typing import Any, Dict, Iterable
+
 from comfy_api.latest import InputImpl, io
 
-from ..nodes.catalog_utils import image2video_model_choices, text2video_model_choices
-from ..nodes.utils import ensure_prompt_length
-from ..nodes.video_utils import (
-    get_testing_video_path,
-    list_testing_videos,
-    poll_video_until_ready,
-    queue_video_job,
-)
+from ..nodes.catalog_utils import video_model_specs
+from ..nodes.utils import encode_tensor_for_vision, ensure_prompt_length
+from ..nodes.video_utils import poll_video_until_ready, queue_video_job
+
+LOG = logging.getLogger(__name__)
 
 
-class GenerateVideoFromText:
+class GenerateVideoFromText(io.ComfyNode):
+
+    @staticmethod
+    def _option_input_id(model_id: str, field: str) -> str:
+        sanitized = "".join(ch if ch.isalnum() or ch in {"_", "-"} else "_" for ch in model_id)
+        return f"{sanitized}__{field}"
+
+    @staticmethod
+    def _constraint_values(value: Iterable | None) -> list[str]:
+        if not value:
+            return []
+        if isinstance(value, (str, bytes)):
+            normalized = str(value).strip()
+            return [normalized] if normalized else []
+        return [str(item) for item in value if item]
+
     @classmethod
-    def INPUT_TYPES(cls):
-        i2v_models = image2video_model_choices()
-        t2v_models = text2video_model_choices()
-        model_choices = i2v_models + t2v_models
-        # todo: make model choices show model name instead of id for readability and prepend i2v/t2v
+    def _video_specs(cls) -> Dict[str, Dict[str, Any]]:
+        specs = video_model_specs() or {}
+        if not specs:
+            raise ValueError(
+                "No Venice video model specs available; refresh the catalog in VeniceAI settings and retry."
+            )
+        return specs
 
-        video_choices = list_testing_videos()
-        existing_default = video_choices[0] if video_choices else "none_available"
-        choices_for_combo = video_choices or ["none_available"]
+    @classmethod
+    def _get_option_value(cls, model_payload: Dict[str, Any], model_id: str, field: str) -> Any:
+        candidates = (
+            cls._option_input_id(model_id, field),
+            field,
+            f"{field}__{model_id}",
+            f"{model_id}__{field}",
+        )
+        for key in candidates:
+            if key in model_payload:
+                return model_payload.get(key)
+        return None
 
-        return {
-            "required": {
-                "model": (
-                    model_choices,
-                    {
-                        # "default": model_choices[0],
-                        "default": "longcat-distilled-text-to-video",
-                        "tooltip": "Model to use for text-to-video generation",
-                    },
-                ),
-                "prompt": (
-                    "STRING",
-                    {
-                        "default": "A cat made of lettuce flying through space",
-                        "placeholder": "Positive Prompt. Example: A cat made of lettuce flying through space",
-                        "tooltip": "Text prompt to generate the video from",
-                        "multiline": True,
-                    },
-                ),
-                "negative_prompt": (
-                    "STRING",
-                    {
-                        "default": "low resolution, error, worst quality, low quality, defects",
-                        "placeholder": "Negative Prompt",
-                        "multiline": True,
-                        "tooltip": "Negative prompt to avoid elements in the video",
-                    },
-                ),
-                "duration": (
-                    ["4s", "5s", "6s", "8s", "10s", "12s", "14s", "15s", "16s", "18s", "20s"],
-                    {
-                        "default": "5s",
-                        "tooltip": "Duration of the generated video",
-                    },
-                ),
-                "aspect_ratio": (
-                    ["16:9", "9:16", "1:1"],
-                    {
-                        "default": "16:9",
-                        "tooltip": "Aspect ratio for the video",
-                    },
-                ),
-                "resolution": (
-                    ["1080p", "720p", "480p"],
-                    {
-                        "default": "720p",
-                        "tooltip": "Resolution of the generated video",
-                    },
-                ),
-                "audio": (
-                    "BOOLEAN",
-                    {
-                        "default": True,
-                        "tooltip": "Generate audio if the model supports it",
-                    },
-                ),
-                "use_existing_video": (
-                    "BOOLEAN",
-                    {
-                        "default": True,  # NOTE: IMPORTANT DEFAULT TO TRUE FOR TESTING PURPOSES THE WHOLE TIME DO NOT REMOVE UNTIL DEPLOYMENT
-                        "tooltip": "Use a cached video from testing_video instead of calling the Venice API",
-                    },
-                ),
-                "existing_video": (
-                    choices_for_combo,
-                    {
-                        "default": existing_default,
-                        "tooltip": "Select the cached video file that should be emitted when bypassing the API",
-                    },
-                ),
-            }
-        }
+    @classmethod
+    def _build_model_options(cls) -> list[io.DynamicCombo.Option]:
+        specs = cls._video_specs()
+        options: list[io.DynamicCombo.Option] = []
 
-    RETURN_TYPES = ("VIDEO",)
-    RETURN_NAMES = ("video",)
-    FUNCTION = "execute"
-    CATEGORY = "venice.ai"
+        def _sorted_models_by_group(group: str) -> list[tuple[str, Dict[str, Any]]]:
+            return sorted(
+                (
+                    (model_id, spec)
+                    for model_id, spec in specs.items()
+                    if spec.get("constraints", {}).get("model_type") == group
+                ),
+                key=lambda item: item[0],
+            )
 
+        ordered_specs = [
+            # didnt know about this, is same as [] + []
+            *_sorted_models_by_group("text-to-video"),
+            *_sorted_models_by_group("image-to-video"),
+        ]
+
+        for model_id, spec in ordered_specs:
+            constraints = spec.get("constraints") or {}
+            aspect_ratios = cls._constraint_values(constraints.get("aspect_ratios"))
+            resolutions = cls._constraint_values(constraints.get("resolutions"))
+            durations = cls._constraint_values(constraints.get("durations"))
+            audio_default = bool(constraints.get("audio")) if constraints.get("audio") is not None else False
+
+            option_inputs: list[io.Input] = []
+            if constraints.get("model_type") == "image-to-video":
+                option_inputs.append(
+                    io.Image.Input(
+                        cls._option_input_id(model_id, "image"),
+                        display_name="image",
+                        tooltip="Source image for image-to-video models",
+                    )
+                )
+
+            if durations:
+                option_inputs.append(
+                    io.Combo.Input(
+                        id=cls._option_input_id(model_id, "duration"),
+                        display_name="duration",
+                        options=durations,
+                        default=durations[0],
+                        tooltip="Duration allowed by the selected model",
+                    )
+                )
+            if aspect_ratios:
+                option_inputs.append(
+                    io.Combo.Input(
+                        cls._option_input_id(model_id, "aspect_ratio"),
+                        display_name="aspect_ratio",
+                        options=aspect_ratios,
+                        default=aspect_ratios[0],
+                        tooltip="Aspect ratios allowed by the selected model",
+                    )
+                )
+            if resolutions:
+                option_inputs.append(
+                    io.Combo.Input(
+                        cls._option_input_id(model_id, "resolution"),
+                        display_name="resolution",
+                        options=resolutions,
+                        default=resolutions[0],
+                        tooltip="Resolutions allowed by the selected model",
+                    )
+                )
+
+            if constraints.get("audio_configurable"):
+                option_inputs.append(
+                    io.Boolean.Input(
+                        cls._option_input_id(model_id, "audio"),
+                        display_name="audio",
+                        default=audio_default,
+                        tooltip="Generate audio (only when the model allows toggling)",
+                    )
+                )
+
+            options.append(io.DynamicCombo.Option(model_id, option_inputs))
+
+        return options
+
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        model_options = cls._build_model_options()
+
+        return io.Schema(
+            node_id="TextToVideo_VENICE",
+            display_name="Generate Video from Text (Venice)",
+            category="venice.ai",
+            inputs=[
+                io.DynamicCombo.Input(
+                    "model",
+                    options=model_options,
+                    tooltip="Select a Venice video model to auto-populate valid parameters",
+                ),
+                io.String.Input(
+                    "prompt",
+                    default="A cat made of lettuce flying through space",
+                    placeholder="Positive Prompt. Example: A cat made of lettuce flying through space",
+                    tooltip="Text prompt to generate the video from",
+                    multiline=True,
+                ),
+                io.String.Input(
+                    "negative_prompt",
+                    default="low resolution, error, worst quality, low quality, defects",
+                    placeholder="Negative Prompt",
+                    tooltip="Negative prompt to avoid elements in the video",
+                    multiline=True,
+                ),
+            ],
+            outputs=[
+                io.Video.Output(id="video", display_name="Video"),
+            ],
+        )
+
+    @classmethod
     def execute(
-        self,
+        cls,
         model,
         prompt,
         negative_prompt,
-        duration,
-        aspect_ratio,
-        resolution,
-        audio,
-        use_existing_video,
-        existing_video,
-    ):
+    ) -> io.NodeOutput:
         ensure_prompt_length(prompt, 2500, label="Prompt")
         ensure_prompt_length(negative_prompt, 2500, label="Negative Prompt", allow_empty=True)
 
+        if not isinstance(model, dict) or "model" not in model:
+            raise ValueError("Model selection is required")
+
+        model_id = model.get("model")
+        specs = cls._video_specs()
+        spec = specs.get(model_id)
+        if not spec:
+            raise ValueError("Selected model is missing from the Venice catalog; refresh the catalog and try again.")
+        constraints = spec.get("constraints") or {}
+
+        durations = cls._constraint_values(constraints.get("durations"))
+        aspect_ratios = cls._constraint_values(constraints.get("aspect_ratios"))
+        resolutions = cls._constraint_values(constraints.get("resolutions"))
+
+        duration = cls._get_option_value(model, model_id, "duration")
+        if durations:
+            if duration is None:
+                raise ValueError(f"Model {model_id} requires a duration selection")
+            if duration not in durations:
+                raise ValueError(f"Duration '{duration}' is not supported by model {model_id}")
+
+        aspect_ratio = cls._get_option_value(model, model_id, "aspect_ratio")
+        if aspect_ratios:
+            if aspect_ratio is None:
+                raise ValueError(f"Model {model_id} requires an aspect ratio selection")
+            if aspect_ratio not in aspect_ratios:
+                raise ValueError(f"Aspect ratio '{aspect_ratio}' is not supported by model {model_id}")
+
+        resolution = cls._get_option_value(model, model_id, "resolution")
+        if resolutions:
+            if resolution is None:
+                raise ValueError(f"Model {model_id} requires a resolution selection")
+            if resolution not in resolutions:
+                raise ValueError(f"Resolution '{resolution}' is not supported by model {model_id}")
+
+        audio_configurable = bool(constraints.get("audio_configurable"))
+        audio_default = bool(constraints.get("audio")) if constraints.get("audio") is not None else False
+        audio_value = cls._get_option_value(model, model_id, "audio") if audio_configurable else None
+        audio = audio_value if audio_value is not None else audio_default
+
         payload = {
-            "model": model,
+            "model": model_id,
             "prompt": prompt,
             "negative_prompt": negative_prompt,
-            "duration": duration,
-            "aspect_ratio": aspect_ratio,
-            "resolution": resolution,
-            # "audio": audio, # todo: this will error with bad request if model without audio support is used, fix with node schema v3 rewrite
         }
+        if durations:
+            payload["duration"] = duration
+        if aspect_ratios:
+            payload["aspect_ratio"] = aspect_ratio
+        if resolutions:
+            payload["resolution"] = resolution
+        if audio is not None:
+            payload["audio"] = audio
 
-        if use_existing_video:
-            if not existing_video:
-                raise ValueError("No cached video selected")
-            cached_files = list_testing_videos()
-            if existing_video not in cached_files:
-                raise ValueError("Selected cached video does not exist anymore")
-            video_path = get_testing_video_path(existing_video)
-            if not video_path.exists():
-                raise ValueError("Cached video file disappeared")
-            return io.NodeOutput(InputImpl.VideoFromFile(video_path))
+        if constraints.get("model_type") == "image-to-video":
+            image = cls._get_option_value(model, model_id, "image")
+            if image is None:
+                raise ValueError(f"Model {model_id} requires an input image")
+            payload["image_url"] = encode_tensor_for_vision(image)
 
-        model_id, queue_id = queue_video_job(payload)
-        video_path, _ = poll_video_until_ready(model=model_id, queue_id=queue_id)
+        model_id_resp, queue_id = queue_video_job(payload)
+        video_path, _ = poll_video_until_ready(model=model_id_resp, queue_id=queue_id)
 
         return io.NodeOutput(InputImpl.VideoFromFile(video_path))
-
-
-NODE_CLASS_MAPPINGS = {
-    "TextToVideo_VENICE": GenerateVideoFromText,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "TextToVideo_VENICE": "Generate Video from Text (Venice)",
-}

nodes/i2i_enhance_upscale.py

@@ -0,0 +1,155 @@
+import base64
+import io as python_io
+import logging
+
+import requests
+from PIL import Image
+from torchvision.transforms import ToPILImage, ToTensor  # type: ignore
+
+from comfy_api.latest import io
+
+from ..globals import API_ENDPOINTS
+from ..nodes.utils import ensure_prompt_length
+from ..venice_client import client
+
+LOG = logging.getLogger(__name__)
+
+
+class I2IEnhanceUpscale(io.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> io.Schema:
+        return io.Schema(
+            node_id="I2IEnhanceUpscale_VENICE",
+            display_name="Img2Img Enhance + Upscale (Venice)",
+            category="venice.ai",
+            inputs=[
+                io.Image.Input(
+                    "image",
+                    tooltip="Image tensor to enhance or upscale",
+                ),
+                io.Float.Input(
+                    "scale",
+                    default=2.0,
+                    min=1.0,
+                    max=4.0,
+                    step=0.01,
+                    tooltip=(
+                        "Scale factor for upscaling the image. Valid values are 1, 2, 3, or 4.\n"
+                        "If set to 1, the image will not be upscaled but enhanced, 'enhance' must be set to 'True'."
+                    ),
+                ),
+                io.Boolean.Input(
+                    "enhance",
+                    default=False,
+                    tooltip=(
+                        "Whether to enhance the image using Venice's image engine during upscaling.\n"
+                        "Must be set to 'True' if scale is set to 1."
+                    ),
+                ),
+                io.Float.Input(
+                    "enhance_creativity",
+                    default=0.5,
+                    min=0.0,
+                    max=1.0,
+                    step=0.01,
+                    tooltip=(
+                        "Higher values let the enhancement AI change the image more. "
+                        "Setting this to 1 effectively creates an entirely new image."
+                    ),
+                ),
+                io.String.Input(
+                    "enhance_prompt",
+                    default="",
+                    multiline=True,
+                    placeholder="Prompt for enhance. Example: gold, graffiti, minimalistic",
+                    tooltip=(
+                        "The text to image style to apply during prompt enhancement. "
+                        "Does best with short descriptive prompts, like gold, marble or angry, menacing."
+                    ),
+                ),
+                io.Float.Input(
+                    "replication",
+                    default=0.1,
+                    min=0.0,
+                    max=1.0,
+                    step=0.01,
+                    tooltip=(
+                        "How strongly lines and noise in the base image are preserved. "
+                        "Higher values are noisier but less plastic/AI 'generated'/hallucinated"
+                    ),
+                ),
+            ],
+            outputs=[io.Image.Output(id="image", display_name="Image")],
+        )
+
+    @classmethod
+    def execute(
+        cls,
+        image,
+        scale,
+        enhance,
+        enhance_creativity,
+        enhance_prompt,
+        replication,
+    ) -> io.NodeOutput:
+        ensure_prompt_length(enhance_prompt, 1500, "Enhance prompt", allow_empty=True)
+        if scale == 1 and not enhance:
+            raise ValueError("Upscale Image (Venice) 'enhance' must be set to 'True' if scale is 1.")
+        if scale == 4:
+            LOG.info(
+                (
+                    "Upscale Image (Venice) A scale of 4 with large images will result "
+                    "in the scale being dynamically set (by venice) to ensure the "
+                    "final image stays within the maximum size limits."
+                )
+            )
+
+        # Convert tensor to PIL Image
+        try:
+            # Get first image from batch
+            img_tensor = image[0].detach().cpu()  # Shape: (H, W, C)
+
+            # Ensure RGB format by taking first 3 channels
+            if img_tensor.shape[-1] > 3:
+                img_tensor = img_tensor[:, :, :3]
+
+            # Convert to CHW format and create PIL Image
+            pil_image = ToPILImage()(img_tensor.permute(2, 0, 1))
+        except Exception as exc:
+            raise ValueError(f"Upscale Image (Venice) Failed to convert tensor to PIL image: {str(exc)}")
+
+        # Convert image to base64
+        byte_io = python_io.BytesIO()
+        pil_image.save(byte_io, format="PNG")
+        byte_io.seek(0)
+        image_base64 = base64.b64encode(byte_io.read()).decode("utf-8")
+
+        payload = {
+            "image": image_base64,
+            "scale": scale,
+            "enhance": enhance,
+            "enhanceCreativity": enhance_creativity,
+            "enhancePrompt": enhance_prompt,
+            "replication": replication,
+        }
+
+        response = None
+        try:
+            response = client.request(
+                "POST",
+                API_ENDPOINTS["upscale_image"],
+                json=payload,
+                headers={"Content-Type": "application/json"},
+            )
+        except requests.exceptions.RequestException as exc:
+            raise RuntimeError(f"Upscale Image (Venice) API request failed: {str(exc)}")
+
+        try:
+            upscaled_image = Image.open(python_io.BytesIO(response.content))
+            tensor = ToTensor()(upscaled_image)  # Converts to (C, H, W)
+            tensor = tensor.permute(1, 2, 0)  # Convert to (H, W, C)
+            tensor = tensor.unsqueeze(0)  # Add batch dimension (1, H, W, C)
+        except Exception as exc:
+            raise ValueError(f"Upscale Image (Venice) Failed to process response image: {str(exc)}")
+
+        return io.NodeOutput(tensor)

nodes/i2i_enhance_upscale_node.py

@@ -1,166 +0,0 @@
-import base64
-import io
-import logging
-
-import requests
-from PIL import Image
-from torchvision.transforms import ToPILImage, ToTensor  # type: ignore
-
-from ..globals import API_ENDPOINTS
-from ..nodes.utils import ensure_prompt_length
-from ..venice_client import client
-
-
-class I2IEnhanceUpscale:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "image": ("IMAGE",),
-                "scale": (
-                    "FLOAT",
-                    {
-                        "default": 2,
-                        "min": 1,
-                        "max": 4,
-                        "step": 0.01,
-                        "tooltip": (
-                            "Scale factor for upscaling the image. Valid values are 1, 2, 3, or 4.\n"
-                            "If set to 1, the image will not be upscaled but enhanced, 'enhanced setting must be set to 'True'."
-                        ),
-                    },
-                ),
-                "enhance": (
-                    "BOOLEAN",
-                    {
-                        "default": False,
-                        "tooltip": (
-                            "Whether to enhance the image using Venice's image engine during upscaling.\n"
-                            "Must be set to 'True' if scale is set to 1."
-                        ),
-                    },
-                ),
-                "enhance_creativity": (
-                    "FLOAT",
-                    {
-                        "default": 0.5,
-                        "min": 0.0,
-                        "max": 1.0,
-                        "step": 0.01,
-                        "tooltip": (
-                            "Higher values let the enhancement AI change the image more. "
-                            "Setting this to 1 effectively creates an entirely new image."
-                        ),
-                    },
-                ),
-                "enhance_prompt": (
-                    "STRING",
-                    {
-                        "default": "",
-                        "placeholder": "gold, graffiti, minimalistic",
-                        "multiline": True,
-                        "tooltip": (
-                            "The text to image style to apply during prompt enhancement. "
-                            "Does best with short descriptive prompts, like gold, marble or angry, menacing."
-                        ),
-                    },
-                ),
-                "replication": (
-                    "FLOAT",
-                    {
-                        "default": 0.1,
-                        "min": 0.0,
-                        "max": 1.0,
-                        "step": 0.01,
-                        "tooltip": (
-                            "How strongly lines and noise in the base image are preserved. "
-                            "Higher values are noisier but less plastic/AI 'generated'/hallucinated"
-                        ),
-                    },
-                ),
-            }
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("image",)
-    FUNCTION = "i2i_enhance_upscale"
-    CATEGORY = "venice.ai"
-
-    def i2i_enhance_upscale(self, image, scale, enhance, enhance_creativity, enhance_prompt, replication):
-        response = None
-
-        ensure_prompt_length(enhance_prompt, 1500, label="Enhance prompt", allow_empty=True)
-        if scale == 1:
-            raise ValueError("Upscale Image (Venice) 'enhance' must be set to 'True' if scale is 1.")
-        if scale == 4:
-            logging.info(
-                (
-                    "Upscale Image (Venice) A scale of 4 with large images will result "
-                    "in the scale being dynamically set (by venice) to ensure the "
-                    "final image stays within the maximum size limits."
-                )
-            )
-
-        # if not api_key:
-        #     raise ValueError("VENICEAI_API_KEY environment variable not set")
-
-        # Convert tensor to PIL Image
-        try:
-            # Get first image from batch
-            img_tensor = image[0].detach().cpu()  # Shape: (H, W, C)
-
-            # Ensure RGB format by taking first 3 channels
-            if img_tensor.shape[-1] > 3:
-                img_tensor = img_tensor[:, :, :3]
-
-            # Convert to CHW format and create PIL Image
-            pil_image = ToPILImage()(img_tensor.permute(2, 0, 1))
-        except Exception as e:
-            raise ValueError(f"Upscale Image (Venice) Failed to convert tensor to PIL image: {str(e)}")
-
-        # Convert image to base64
-        byte_io = io.BytesIO()
-        pil_image.save(byte_io, format="PNG")
-        byte_io.seek(0)
-        image_base64 = base64.b64encode(byte_io.read()).decode("utf-8")
-
-        # Prepare JSON payload
-        payload = {
-            "image": image_base64,
-            "scale": scale,
-            "enhance": enhance,
-            "enhanceCreativity": enhance_creativity,
-            "enhancePrompt": enhance_prompt,
-            "replication": replication,
-        }
-
-        # Send request
-        try:
-            response = client.request(
-                "POST",
-                API_ENDPOINTS["upscale_image"],
-                json=payload,
-                headers={"Content-Type": "application/json"},
-            )
-        except requests.exceptions.RequestException as e:
-            raise RuntimeError(f"Upscale Image (Venice) API request failed: {str(e)}")
-
-        # Convert response to tensor
-        try:
-            upscaled_image = Image.open(io.BytesIO(response.content))
-            tensor = ToTensor()(upscaled_image)  # Converts to (C, H, W)
-            tensor = tensor.permute(1, 2, 0)  # Convert to (H, W, C)
-            tensor = tensor.unsqueeze(0)  # Add batch dimension (1, H, W, C)
-        except Exception as e:
-            raise ValueError(f"Upscale Image (Venice) Failed to process response image: {str(e)}")
-
-        return (tensor,)
-
-
-NODE_CLASS_MAPPINGS = {
-    "I2IEnhanceUpscale_VENICE": I2IEnhanceUpscale,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "I2IEnhanceUpscale_VENICE": "Img2Img Enhance + Upscale (Venice)",
-}

nodes/util_nodes.py

@@ -1,27 +0,0 @@
-class CharCountTextBox:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "input_text": ("STRING", {"default": "", "multiline": True}),
-            }
-        }
-
-    CATEGORY = "venice.ai"
-
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("text",)
-    FUNCTION = "return_same_text"
-
-    def return_same_text(self, input_text):
-
-        return {"ui": {"text": input_text}, "result": (input_text,)}
-
-
-NODE_CLASS_MAPPINGS = {
-    "CharCountTextBox": CharCountTextBox,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "CharCountTextBox": "Textbox w/ char count",
-}

nodes/utils.py

@@ -17,7 +17,7 @@ def _round_down(value: int, multiple: int) -> int:
     return (value // multiple) * multiple
 
 
-def ensure_prompt_length(text: str, maximum: int, label: str, allow_empty: bool = False) -> None:
+def ensure_prompt_length(text: str, maximum: int, label: str, *, allow_empty: bool = False) -> None:
     if not text:
         if allow_empty:
             return
@@ -48,6 +48,10 @@ def encode_tensor_for_vision(image_tensor: torch.Tensor, *, max_encoded_bytes: i
         If the tensor does not have 3 dimensions or the last dimension is not 3 or 4.
     """
     tensor = image_tensor.detach().cpu()
+    if tensor.ndim == 4 and tensor.shape[0] == 1:
+        tensor = tensor[0]
+    if tensor.ndim == 3 and tensor.shape[-1] not in {3, 4} and tensor.shape[0] in {3, 4}:
+        tensor = tensor.permute(1, 2, 0)
     if tensor.ndim != 3 or tensor.shape[-1] not in {3, 4}:
         raise ValueError("Vision images must be height×width×(3 or 4 channels)")
 
@@ -115,6 +119,7 @@ def ensure_multiple_of(width: int, height: int, *, multiple: int = 32) -> None:
         dimensions = " and ".join(bad_dimensions)
         raise ValueError(f"{dimensions} must be multiples of {multiple}")
 
+
 # unused right now, might be useful, or not
 def ensure_api_key_valid() -> None:
     key = venice_config.apikey.strip()

nodes/video_utils.py

@@ -1,10 +1,7 @@
 import logging
-import os
 import tempfile
 import time
-from pathlib import Path
-from typing import List, Optional, Tuple
-
+from typing import Optional, Tuple
 
 from comfy.utils import ProgressBar  # type: ignore
 
@@ -17,10 +14,6 @@ LOG = logging.getLogger(__name__)
 POLL_INTERVAL_SECONDS = 5
 MAX_POLLS = 100  # MAX_POLLS * POLL_INTERVAL_SECONDS = x minutes total wait time
 PROGRESS_BAR_TOTAL = 100
-PROJECT_ROOT = Path(__file__).resolve().parents[1]
-VIDEO_OUTPUT_DIR = PROJECT_ROOT / "testing_video"
-VIDEO_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-DEBUG_SAVE_API_VIDEOS = os.environ.get("VENICE_CLIENT_DEBUG", "").lower() in {"1", "true"}
 
 
 def queue_video_job(payload: dict) -> Tuple[str, str]:
@@ -47,13 +40,6 @@ def _guess_suffix(content_type: str) -> str:
     return ".bin"
 
 
-def _build_video_path(queue_id: str, suffix: str) -> Path:
-    clean_id = "".join(ch if ch.isalnum() or ch in "-_" else "_" for ch in queue_id)
-    timestamp = int(time.time())
-    filename = f"{clean_id or 'queue'}_{timestamp}{suffix}"
-    return VIDEO_OUTPUT_DIR / filename
-
-
 def poll_video_until_ready(
     *,
     model: str,
@@ -123,14 +109,9 @@ def poll_video_until_ready(
 
         # Got binary content (video)
         suffix = _guess_suffix(ctype)
-        if DEBUG_SAVE_API_VIDEOS:
-            video_path = _build_video_path(queue_id, suffix)
-            with open(video_path, "wb") as fp:
-                fp.write(resp.content)
-        else:
-            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
-                tmp_file.write(resp.content)
-                video_path = Path(tmp_file.name)
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
+            tmp_file.write(resp.content)
+            video_path = tmp_file.name
 
         LOG.debug(
             "Venice video ready: queue_id=%s saved_to=%s content_type=%s",
@@ -147,18 +128,3 @@ def poll_video_until_ready(
     raise VeniceAPIError(
         f"Timed out waiting for Venice video. queue_id={queue_id} after {max_polls * poll_interval:.0f}s"
     )
-
-
-def list_testing_videos() -> List[str]:
-    """Return cached video filenames that can be selected inside a node."""
-    if not VIDEO_OUTPUT_DIR.exists():
-        return []
-    return sorted(entry.name for entry in VIDEO_OUTPUT_DIR.iterdir() if entry.is_file())
-
-
-def get_testing_video_path(filename: str) -> Path:
-    """Resolve the video file inside testing_video and ensure it exists."""
-    video_path = VIDEO_OUTPUT_DIR / filename
-    if not video_path.exists() or not video_path.is_file():
-        raise FileNotFoundError(f"Test video not found: {filename}")
-    return video_path

requirements.txt

@@ -1,4 +1,7 @@
 requests>=2.31.0
 Pillow>=10.0.0
 numpy>=1.24.0
-configparser>=5.3.0
+configparser>=5.3.0
+torchaudio
+torchvision
+imageio-ffmpeg

venice_catalog.py

@@ -78,6 +78,106 @@ def _enrich_models(raw: Dict[str, Any]) -> Dict[str, Any]:
     return raw.copy()
 
 
+def _extract_video_models(payload: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Extract Venice video models and expose a by-id mapping of their constraints.
+
+    The mapping is intentionally lightweight (id, name, model_type, constraints) so downstream
+    callers (nodes) can build DynamicCombo inputs without re-parsing the raw payload.
+    """
+    data = payload.get("data", []) or []
+    by_id: Dict[str, Any] = {}
+
+    for model in data:
+        if model.get("type") != "video":
+            continue
+        model_id = model.get("id")
+        if not model_id:
+            continue
+
+        model_spec = model.get("model_spec") or {}
+        constraints = model_spec.get("constraints") or {}
+        model_type = constraints.get("model_type")
+        if not isinstance(constraints, dict):
+            constraints = {}
+
+        entry = {
+            "id": model_id,
+            "name": model_spec.get("name"),
+            "model_type": model_type,
+            "constraints": {
+                "aspect_ratios": constraints.get("aspect_ratios") or [],
+                "resolutions": constraints.get("resolutions") or [],
+                "durations": constraints.get("durations") or [],
+                "audio": constraints.get("audio"),
+                "audio_configurable": constraints.get("audio_configurable"),
+                "model_type": model_type,
+            },
+            "raw": model,
+        }
+        by_id[model_id] = entry
+
+    return by_id
+
+
+def _extract_image_models(payload: Dict[str, Any]) -> Dict[str, Any]:
+    data = payload.get("data", []) or []
+    by_id: Dict[str, Any] = {}
+
+    for model in data:
+        if model.get("type") != "image":
+            continue
+        model_id = model.get("id")
+        if not model_id:
+            continue
+
+        model_spec = model.get("model_spec") or {}
+        constraints = model_spec.get("constraints") or {}
+        if not isinstance(constraints, dict):
+            constraints = {}
+
+        entry = {
+            "id": model_id,
+            "name": model_spec.get("name"),
+            "constraints": dict(constraints),
+            "raw": model,
+        }
+        by_id[model_id] = entry
+
+    return by_id
+
+
+def _extract_text_models(payload: Dict[str, Any]) -> Dict[str, Any]:
+    data = payload.get("data", []) or []
+    by_id: Dict[str, Any] = {}
+
+    for model in data:
+        if model.get("type") != "text":
+            continue
+        model_id = model.get("id")
+        if not model_id:
+            continue
+
+        model_spec = model.get("model_spec") or {}
+        constraints = model_spec.get("constraints") or {}
+        capabilities = model_spec.get("capabilities") or {}
+        if not isinstance(constraints, dict):
+            constraints = {}
+        if not isinstance(capabilities, dict):
+            capabilities = {}
+
+        entry = {
+            "id": model_id,
+            "name": model_spec.get("name"),
+            "constraints": dict(constraints),
+            "capabilities": dict(capabilities),
+            "raw": model,
+        }
+        by_id[model_id] = entry
+
+    return by_id
+
+
 def _should_refresh(last_refresh: float) -> bool:
     if _CACHE_TTL <= 0:
         return False
@@ -184,9 +284,16 @@ def get_models(model_type: Optional[str] = None, *, force_refresh: bool = False)
     if model_type:
         return {"models": _model_store.filter_by_type(payload, model_type)}
 
+    video_models_by_id = _extract_video_models(payload)
+    image_models_by_id = _extract_image_models(payload)
+    text_models_by_id = _extract_text_models(payload)
+
+    # todo: maybe dataclass is better for this
     filtered = {
+        # todo: the *_by_id might be enough so this stuff below can be removed
         "image_models": sorted([m.get("id") for m in payload.get("data", []) if m.get("type") == "image"]),
         "text_models": sorted([m.get("id") for m in payload.get("data", []) if m.get("type") == "text"]),
+        # todo: voices should be linked to models like in models json
         "tts_models": sorted([m.get("id") for m in payload.get("data", []) if m.get("type") == "tts"]),
         "tts_voices": sorted(
             [
@@ -198,25 +305,22 @@ def get_models(model_type: Optional[str] = None, *, force_refresh: bool = False)
         ),
         "text2video_models": sorted(
             [
-                m.get("id")
-                for m in payload.get("data", [])
-                if m.get("type") == "video"
-                and isinstance(m.get("model_spec"), dict)
-                and isinstance(m["model_spec"].get("constraints"), dict)
-                and m["model_spec"]["constraints"].get("model_type") == "text-to-video"
+                model_id
+                for model_id, spec in video_models_by_id.items()
+                if spec.get("constraints", {}).get("model_type") == "text-to-video"
             ]
         ),
         "image2video_models": sorted(
             [
-                m.get("id")
-                for m in payload.get("data", [])
-                if m.get("type") == "video"
-                and isinstance(m.get("model_spec"), dict)
-                and isinstance(m["model_spec"].get("constraints"), dict)
-                and m["model_spec"]["constraints"].get("model_type") == "image-to-video"
+                model_id
+                for model_id, spec in video_models_by_id.items()
+                if spec.get("constraints", {}).get("model_type") == "image-to-video"
             ]
         ),
+        "video_models_by_id": video_models_by_id,
+        "image_models_by_id": image_models_by_id,
         "model_list_json": payload,
+        "text_models_by_id": text_models_by_id,
     }
     return filtered