[Partner Nodes] added 4K resolution for Veo models; added Veo 3 Lite model

comfy_api_nodes/nodes_veo2.py

@@ -24,8 +24,9 @@
 AVERAGE_DURATION_VIDEO_GEN = 32
 MODELS_MAP = {
     "veo-2.0-generate-001": "veo-2.0-generate-001",
-    "veo-3.1-generate": "veo-3.1-generate-preview",
-    "veo-3.1-fast-generate": "veo-3.1-fast-generate-preview",
+    "veo-3.1-generate": "veo-3.1-generate-001",
+    "veo-3.1-fast-generate": "veo-3.1-fast-generate-001",
+    "veo-3.1-lite": "veo-3.1-lite-generate-001",
     "veo-3.0-generate-001": "veo-3.0-generate-001",
     "veo-3.0-fast-generate-001": "veo-3.0-fast-generate-001",
 }
@@ -247,17 +248,8 @@ def status_extractor(response):
         raise Exception("Video generation completed but no video was returned")
 
 
-class Veo3VideoGenerationNode(VeoVideoGenerationNode):
-    """
-    Generates videos from text prompts using Google's Veo 3 API.
-
-    Supported models:
-    - veo-3.0-generate-001
-    - veo-3.0-fast-generate-001
-
-    This node extends the base Veo node with Veo 3 specific features including
-    audio generation and fixed 8-second duration.
-    """
+class Veo3VideoGenerationNode(IO.ComfyNode):
+    """Generates videos from text prompts using Google's Veo 3 API."""
 
     @classmethod
     def define_schema(cls):
@@ -279,6 +271,13 @@ def define_schema(cls):
                     default="16:9",
                     tooltip="Aspect ratio of the output video",
                 ),
+                IO.Combo.Input(
+                    "resolution",
+                    options=["720p", "1080p", "4k"],
+                    default="720p",
+                    tooltip="Output video resolution. 4K is not available for veo-3.1-lite and veo-3.0 models.",
+                    optional=True,
+                ),
                 IO.String.Input(
                     "negative_prompt",
                     multiline=True,
@@ -289,11 +288,11 @@ def define_schema(cls):
                 IO.Int.Input(
                     "duration_seconds",
                     default=8,
-                    min=8,
+                    min=4,
                     max=8,
-                    step=1,
+                    step=2,
                     display_mode=IO.NumberDisplay.number,
-                    tooltip="Duration of the output video in seconds (Veo 3 only supports 8 seconds)",
+                    tooltip="Duration of the output video in seconds",
                     optional=True,
                 ),
                 IO.Boolean.Input(
@@ -332,10 +331,10 @@ def define_schema(cls):
                     options=[
                         "veo-3.1-generate",
                         "veo-3.1-fast-generate",
+                        "veo-3.1-lite",
                         "veo-3.0-generate-001",
                         "veo-3.0-fast-generate-001",
                     ],
-                    default="veo-3.0-generate-001",
                     tooltip="Veo 3 model to use for video generation",
                     optional=True,
                 ),
@@ -356,21 +355,111 @@ def define_schema(cls):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model", "generate_audio"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["model", "generate_audio", "resolution", "duration_seconds"]),
                 expr="""
                 (
                   $m := widgets.model;
+                  $r := widgets.resolution;
                   $a := widgets.generate_audio;
-                  ($contains($m,"veo-3.0-fast-generate-001") or $contains($m,"veo-3.1-fast-generate"))
-                    ? {"type":"usd","usd": ($a ? 1.2 : 0.8)}
-                    : ($contains($m,"veo-3.0-generate-001") or $contains($m,"veo-3.1-generate"))
-                      ? {"type":"usd","usd": ($a ? 3.2 : 1.6)}
-                      : {"type":"range_usd","min_usd":0.8,"max_usd":3.2}
+                  $seconds := widgets.duration_seconds;
+                  $pps :=
+                    $contains($m, "lite")
+                      ? ($r = "1080p" ? ($a ? 0.08 : 0.05) : ($a ? 0.05 : 0.03))
+                    : $contains($m, "3.1-fast")
+                      ? ($r = "4k" ? ($a ? 0.30 : 0.25) : $r = "1080p" ? ($a ? 0.12 : 0.10) : ($a ? 0.10 : 0.08))
+                    : $contains($m, "3.1-generate")
+                      ? ($r = "4k" ? ($a ? 0.60 : 0.40) : ($a ? 0.40 : 0.20))
+                    : $contains($m, "3.0-fast")
+                      ? ($a ? 0.15 : 0.10)
+                    : ($a ? 0.40 : 0.20);
+                  {"type":"usd","usd": $pps * $seconds}
                 )
                 """,
             ),
         )
 
+    @classmethod
+    async def execute(
+        cls,
+        prompt,
+        aspect_ratio="16:9",
+        resolution="720p",
+        negative_prompt="",
+        duration_seconds=8,
+        enhance_prompt=True,
+        person_generation="ALLOW",
+        seed=0,
+        image=None,
+        model="veo-3.0-generate-001",
+        generate_audio=False,
+    ):
+        if "lite" in model and resolution == "4k":
+            raise Exception("4K resolution is not supported by the veo-3.1-lite model.")
+
+        model = MODELS_MAP[model]
+
+        instances = [{"prompt": prompt}]
+        if image is not None:
+            image_base64 = tensor_to_base64_string(image)
+            if image_base64:
+                instances[0]["image"] = {"bytesBase64Encoded": image_base64, "mimeType": "image/png"}
+
+        parameters = {
+            "aspectRatio": aspect_ratio,
+            "personGeneration": person_generation,
+            "durationSeconds": duration_seconds,
+            "enhancePrompt": True,
+            "generateAudio": generate_audio,
+        }
+        if negative_prompt:
+            parameters["negativePrompt"] = negative_prompt
+        if seed > 0:
+            parameters["seed"] = seed
+        if "veo-3.1" in model:
+            parameters["resolution"] = resolution
+
+        initial_response = await sync_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/veo/{model}/generate", method="POST"),
+            response_model=VeoGenVidResponse,
+            data=VeoGenVidRequest(
+                instances=instances,
+                parameters=parameters,
+            ),
+        )
+
+        poll_response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/veo/{model}/poll", method="POST"),
+            response_model=VeoGenVidPollResponse,
+            status_extractor=lambda r: "completed" if r.done else "pending",
+            data=VeoGenVidPollRequest(operationName=initial_response.name),
+            poll_interval=9.0,
+            estimated_duration=AVERAGE_DURATION_VIDEO_GEN,
+        )
+
+        if poll_response.error:
+            raise Exception(f"Veo API error: {poll_response.error.message} (code: {poll_response.error.code})")
+
+        response = poll_response.response
+        filtered_count = response.raiMediaFilteredCount
+        if filtered_count:
+            reasons = response.raiMediaFilteredReasons or []
+            reason_part = f": {reasons[0]}" if reasons else ""
+            raise Exception(
+                f"Content blocked by Google's Responsible AI filters{reason_part} "
+                f"({filtered_count} video{'s' if filtered_count != 1 else ''} filtered)."
+            )
+
+        if response.videos:
+            video = response.videos[0]
+            if video.bytesBase64Encoded:
+                return IO.NodeOutput(InputImpl.VideoFromFile(BytesIO(base64.b64decode(video.bytesBase64Encoded))))
+            if video.gcsUri:
+                return IO.NodeOutput(await download_url_to_video_output(video.gcsUri))
+            raise Exception("Video returned but no data or URL was provided")
+        raise Exception("Video generation completed but no video was returned")
+
 
 class Veo3FirstLastFrameNode(IO.ComfyNode):
 
@@ -394,7 +483,7 @@ def define_schema(cls):
                     default="",
                     tooltip="Negative text prompt to guide what to avoid in the video",
                 ),
-                IO.Combo.Input("resolution", options=["720p", "1080p"]),
+                IO.Combo.Input("resolution", options=["720p", "1080p", "4k"]),
                 IO.Combo.Input(
                     "aspect_ratio",
                     options=["16:9", "9:16"],
@@ -424,8 +513,7 @@ def define_schema(cls):
                 IO.Image.Input("last_frame", tooltip="End frame"),
                 IO.Combo.Input(
                     "model",
-                    options=["veo-3.1-generate", "veo-3.1-fast-generate"],
-                    default="veo-3.1-fast-generate",
+                    options=["veo-3.1-generate", "veo-3.1-fast-generate", "veo-3.1-lite"],
                 ),
                 IO.Boolean.Input(
                     "generate_audio",
@@ -443,26 +531,20 @@ def define_schema(cls):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model", "generate_audio", "duration"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["model", "generate_audio", "duration", "resolution"]),
                 expr="""
                 (
-                  $prices := {
-                    "veo-3.1-fast-generate": { "audio": 0.15, "no_audio": 0.10 },
-                    "veo-3.1-generate":      { "audio": 0.40, "no_audio": 0.20 }
-                  };
                   $m := widgets.model;
-                  $ga := (widgets.generate_audio = "true");
+                  $r := widgets.resolution;
+                  $ga := widgets.generate_audio;
                   $seconds := widgets.duration;
-                  $modelKey :=
-                    $contains($m, "veo-3.1-fast-generate") ? "veo-3.1-fast-generate" :
-                    $contains($m, "veo-3.1-generate")      ? "veo-3.1-generate" :
-                    "";
-                  $audioKey := $ga ? "audio" : "no_audio";
-                  $modelPrices := $lookup($prices, $modelKey);
-                  $pps := $lookup($modelPrices, $audioKey);
-                  ($pps != null)
-                    ? {"type":"usd","usd": $pps * $seconds}
-                    : {"type":"range_usd","min_usd": 0.4, "max_usd": 3.2}
+                  $pps :=
+                    $contains($m, "lite")
+                      ? ($r = "1080p" ? ($ga ? 0.08 : 0.05) : ($ga ? 0.05 : 0.03))
+                    : $contains($m, "fast")
+                      ? ($r = "4k" ? ($ga ? 0.30 : 0.25) : $r = "1080p" ? ($ga ? 0.12 : 0.10) : ($ga ? 0.10 : 0.08))
+                    : ($r = "4k" ? ($ga ? 0.60 : 0.40) : ($ga ? 0.40 : 0.20));
+                  {"type":"usd","usd": $pps * $seconds}
                 )
                 """,
             ),
@@ -482,6 +564,9 @@ async def execute(
         model: str,
         generate_audio: bool,
     ):
+        if "lite" in model and resolution == "4k":
+            raise Exception("4K resolution is not supported by the veo-3.1-lite model.")
+
         model = MODELS_MAP[model]
         initial_response = await sync_op(
             cls,
@@ -519,7 +604,7 @@ async def execute(
             data=VeoGenVidPollRequest(
                 operationName=initial_response.name,
             ),
-            poll_interval=5.0,
+            poll_interval=9.0,
             estimated_duration=AVERAGE_DURATION_VIDEO_GEN,
         )