From 7d114d9dd7657d42ce17ce257b7b2851f2ae1bb5 Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Wed, 30 Apr 2025 04:31:37 -0500
Subject: [PATCH] Added pixverse_template support to Pixverse Text to Video
 node (#54)

---
 comfy_api_nodes/apis/pixverse_api.py | 15 +++++++++-
 comfy_api_nodes/nodes_pixverse.py    | 42 ++++++++++++++++++++++++++--
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/comfy_api_nodes/apis/pixverse_api.py b/comfy_api_nodes/apis/pixverse_api.py
index 174d50965..85eae2738 100644
--- a/comfy_api_nodes/apis/pixverse_api.py
+++ b/comfy_api_nodes/apis/pixverse_api.py
@@ -6,6 +6,19 @@ from typing import Optional
 from pydantic import BaseModel, Field
 
 
+pixverse_templates = {
+    "Microwave": 324641385496960,
+    "Suit Swagger": 328545151283968,
+    "Anything, Robot": 313358700761536,
+    "Subject 3 Fever": 327828816843648,
+    "kiss kiss": 315446315336768,
+}
+
+
+class PixverseIO:
+    TEMPLATE = "PIXVERSE_TEMPLATE"
+
+
 class PixverseStatus(int, Enum):
     successful = 1
     generating = 5
@@ -57,7 +70,7 @@ class PixverseDto_V2OpenAPIT2VReq(BaseModel):
     negative_prompt: Optional[str] = Field(None)
     seed: Optional[int] = Field(None)
     style: Optional[str] = Field(None)
-    template_id: Optional[str] = Field(None)
+    template_id: Optional[int] = Field(None)
     water_mark: Optional[bool] = Field(None)
 
 
diff --git a/comfy_api_nodes/nodes_pixverse.py b/comfy_api_nodes/nodes_pixverse.py
index b90601deb..70083c2e7 100644
--- a/comfy_api_nodes/nodes_pixverse.py
+++ b/comfy_api_nodes/nodes_pixverse.py
@@ -9,6 +9,8 @@ from comfy_api_nodes.apis.pixverse_api import (
     PixverseDuration,
     PixverseMotionMode,
     PixverseStatus,
+    PixverseIO,
+    pixverse_templates,
 )
 from comfy_api_nodes.apis.client import (
     ApiEndpoint,
@@ -24,6 +26,32 @@ import requests
 from io import BytesIO
 
 
+class PixverseTemplateNode:
+    """
+    Select template for Pixverse Video generation.
+    """
+
+    RETURN_TYPES = (PixverseIO.TEMPLATE,)
+    RETURN_NAMES = ("pixverse_template",)
+    FUNCTION = "create_template"
+    CATEGORY = "api node/video/Pixverse"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "template": (list(pixverse_templates.keys()), ),
+            }
+        }
+
+    def create_template(self, template: str):
+        template_id = pixverse_templates.get(template, None)
+        if template_id is None:
+            raise Exception(f"Template '{template}' is not recognized.")
+        # just return the integer
+        return (template_id,)
+
+
 class PixverseTextToVideoNode(ComfyNodeABC):
     """
     Generates videos synchronously based on prompt and output_size.
@@ -78,6 +106,12 @@ class PixverseTextToVideoNode(ComfyNodeABC):
                         "tooltip": "An optional text description of undesired elements on an image.",
                     },
                 ),
+                "pixverse_template": (
+                    PixverseIO.TEMPLATE,
+                    {
+                        "tooltip": "An optional template to influence style of generation, created by the Pixverse Template node."
+                    }
+                )
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
@@ -93,6 +127,7 @@ class PixverseTextToVideoNode(ComfyNodeABC):
         motion_mode: str,
         seed,
         negative_prompt: str=None,
+        pixverse_template: int=None,
         auth_token=None,
         **kwargs,
     ):
@@ -118,6 +153,7 @@ class PixverseTextToVideoNode(ComfyNodeABC):
                 duration=duration_seconds,
                 motion_mode=motion_mode,
                 negative_prompt=negative_prompt if negative_prompt else None,
+                template_id=pixverse_template,
                 seed=seed,
             ),
             auth_token=auth_token,
@@ -146,9 +182,11 @@ class PixverseTextToVideoNode(ComfyNodeABC):
 
 
 NODE_CLASS_MAPPINGS = {
-    "PixverseTextToVideoNode": PixverseTextToVideoNode
+    "PixverseTextToVideoNode": PixverseTextToVideoNode,
+    "PixverseTemplateNode": PixverseTemplateNode,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "PixverseTextToVideoNode": "Pixverse Text to Video"
+    "PixverseTextToVideoNode": "Pixverse Text to Video",
+    "PixverseTemplateNode": "Pixverse Template",
 }