From 275e55befeb1da31a01d38b71fdbf741562d5da0 Mon Sep 17 00:00:00 2001 From: David Kim <100145320+Groundbreaker-kim@users.noreply.github.com> Date: Tue, 29 Jul 2025 23:13:57 +0900 Subject: [PATCH] Refactoring the ComfyUI custom node codes to utilize the latest model, and adding Video Model and Preview functionalities (#1568) * Refactoring the code to utilize the latest model, and adding Video Model and Preview functionalities * Adding license headers --------- Co-authored-by: taebeom@google.com --- tools/comfyui_custom_nodes/README.md | 141 ++++++++++++++++-- tools/comfyui_custom_nodes/__init__.py | 10 +- .../src/vertexai_nodes/__init__.py | 78 ---------- .../src/vertexai_nodes/config.yaml | 48 ++++++ .../src/vertexai_nodes/modules/consts.py | 51 ++++--- .../src/vertexai_nodes/modules/utils.py | 1 - .../nodes/{flash.py => gemini.py} | 19 ++- .../src/vertexai_nodes/nodes/imagen.py | 5 +- .../vertexai_nodes/nodes/imagen_bg_swap.py | 1 + .../nodes/imagen_inpaint_insert.py | 1 + .../nodes/imagen_inpaint_remove.py | 1 + .../vertexai_nodes/nodes/imagen_outpaint.py | 1 + .../src/vertexai_nodes/nodes/veo.py | 41 ++--- .../src/vertexai_nodes/web/js/videoPreview.js | 1 + 14 files changed, 255 insertions(+), 144 deletions(-) create mode 100644 tools/comfyui_custom_nodes/src/vertexai_nodes/config.yaml rename tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/{flash.py => gemini.py} (94%) diff --git a/tools/comfyui_custom_nodes/README.md b/tools/comfyui_custom_nodes/README.md index f907cfb4449..c5b3b2d5926 100644 --- a/tools/comfyui_custom_nodes/README.md +++ b/tools/comfyui_custom_nodes/README.md @@ -1,33 +1,142 @@ +# Google Cloud Vertex AI Custom node for ComfyUI + +This repo is for developing a cloud-based visual AI workflow on Google Cloud Platform (GCP) by leveraging ComfyUI and Google Cloud's powerful foundation models. + ## ComfyUI Vertex AI Nodes -This custom node pack for ComfyUI provides seamless integration with Google Cloud's Vertex AI services, allowing you to leverage powerful AI models directly within your ComfyUI workflows. +This custom node pack for ComfyUI provides seamless integration with Google Cloud's Vertex AI services, allowing you to leverage powerful AI models directly within your ComfyUI workflows. This enables you to build sophisticated generative AI applications with the scalability and performance of Google Cloud. ### Nodes Included This pack currently includes the following nodes: -* **ImageGen 3:** Generates images using a Vertex AI image generation model -* **Gemini Flash:** Interacts with the Google Gemini Flash model for fast and efficient text generation. -* **Veo 2:** Generates video using a Vertex AI video model. +| Node Name | Description| +| :---------------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Gemini** | Interacts with Google's Gemini Flash, Pro model for fast and efficient text generation, ideal for real-time chat, summarization, and rapid content creation. | +| **Veo Video Generation** | Generates video content using the Google Veo 2,3 model, an advanced text-to-video diffusion model capable of producing high-quality, diverse video clips from text and image prompts. | +| **VideoPreviewNode** | Used to preview video stored in Google Cloud Storage (GCS) directly within your ComfyUI workflow, providing immediate visual feedback for video generation processes. | +| **Image to B64 Node** | Converts an input image into a Base64 encoded string, useful for embedding image data in text-based formats (e.g., JSON, HTML) or for API transmissions. | +| **Image Generation 3** | The primary node for generating high-quality images from text prompts using Google's Imagen 3 model, designed for diverse visual content creation. | +| **Inpaint Insert w Mask** | Inserts new content into a specified masked area of an image based on a text prompt, requiring an explicit mask to define the insertion region. | +| **Inpaint Insert w AutoMask** | Similar to "Inpaint Insert w Mask," but automatically generates the mask based on the prompt, streamlining the inpainting process. | +| **Inpaint Insert w SemanticMask** | Uses semantic understanding to intelligently generate a mask around specified objects (e.g., "the car") for precise content insertion based on a text prompt. | +| **Inpaint Remove w Mask** | Removes content from a specified masked area of an image and intelligently fills in the gap, requiring an explicit mask for removal. | +| **Inpaint Remove w AutoMask** | Similar to "Inpaint Remove w Mask," but automatically generates the mask for content removal based on the prompt. | +| **Inpaint Remove w SemanticMask** | Removes content based on semantic understanding, where the model automatically generates the mask around the specified object for removal (e.g., "remove the person"). | +| **Imagen Product Background Swap w Mask** | Replaces the background of a product image, requiring an explicit mask for the product to ensure clean separation from the original background. | +| **Imagen Product Background Swap w AutoMask** | Automatically detects the product and generates a mask for it, significantly speeding up the background swapping process for product images. | +| **Imagen Mask-Free Editing** | Enables image editing without the need for explicit masks. Users provide an image and a text prompt describing the desired change (e.g., "change the color"), and Imagen 3 intelligently applies the edits. | +| **Imagen Outpainting** | Extends the boundaries of an image beyond its original canvas, generating new content that seamlessly blends with the existing image based on a text prompt. | +# Pre-Requisite -### Installation -1. Navigate to your ComfyUI `custom_nodes` directory. -2. Add the contents of the `comfyui_custom_nodes` directory to the `custom_nodes` directory. -3. Set the values of PROJECT_ID and REGION in your environment variables or directly in the const.py file located at `src/vertexai_nodes/const.py`. These values are required for the nodes to function correctly. +--- +### Step 1: Create a Google Cloud Project + +Every Google Cloud resource belongs to a project. If you don't have one, create a new project: + +1. Go to the [Google Cloud Console](https://console.cloud.google.com/). +2. In the top navigation bar, click the **project dropdown** (usually displaying "My First Project" or your current project name). +3. Click **"New Project"**. +4. Enter a **Project name** (e.g., `my-api-access-project`). +5. (Optional) Choose an **Organization** and **Location** if applicable. +6. Click **"Create"**. +7. Once the project is created, ensure it is selected in the project dropdown in the Cloud Console. Note down your **Project ID** (e.g., `my-api-access-project-123456`) as you'll need it later. + +--- +### Step 2: Enable Required Google Cloud APIs + +For your application to use a specific Google Cloud service (e.g., Cloud Storage, Vertex AI, Gemini for Google Cloud API), you must enable its corresponding API in your project. + +1. In the Google Cloud Console, navigate to the **Navigation menu (☰)** on the top left. +2. Go to **APIs & Services > Dashboard**. +3. Click **"+ ENABLE APIS AND SERVICES"** at the top. +4. In the API Library, search for the API you need (e.g., "Vertex AI API", "Gemini for Google Cloud API", "Cloud Storage API"). +5. Click on the API you want to enable. +6. Click the **"Enable"** button. + +--- +### Step 3: Authenticate Your Local Environment + +Google Cloud APIs require authentication to verify your application's identity and permissions. There are two primary methods for local development: - ```python - PROJECT_ID = "your-project-id" - REGION = "your-region" # e.g., "us-central1" +#### Method 1: Google Cloud CLI (Recommended for Development) + +This method uses your personal Google account credentials to authenticate your `gcloud` CLI, which then automatically provides **Application Default Credentials (ADC)** to your local applications. This is the simplest and most secure method for individual development. + +1. **Ensure `gcloud` CLI is installed:** If not, follow the [installation guide](https://cloud.google.com/sdk/docs/install). +2. **Log in to your Google Account via `gcloud`:** Open your terminal or command prompt and run: + + ```bash + gcloud auth login ``` - Alternatively, you can set these environment variables in your terminal before starting ComfyUI: + + This command will open a browser window, prompting you to log in with your Google account. After successful login, you can close the browser. + +3. **Set your Google Cloud Project:** Tell `gcloud` which project to use for subsequent commands and for ADC: ```bash - export PROJECT_ID="your-project-id" - export REGION="your-region" # e.g., "us-central1" + gcloud config set project YOUR_PROJECT_ID ``` + Replace `YOUR_PROJECT_ID` with the actual ID of the project you created (e.g., `my-api-access-project-123456`). + +4. **Verify Authentication:** You can verify your active account and project: + + ```bash + gcloud auth list + gcloud config list project + ``` +Now, any Google Cloud client library or application running in your local environment that uses Application Default Credentials will automatically use the authenticated `gcloud` credentials. + +#### Method 2: Service Account Key File (Recommended for Production/CI/CD) + +For production environments, CI/CD pipelines, or situations where you don't want to use a personal account, **service accounts** are the secure and recommended method. A service account is a special type of Google account that represents your application rather than an individual user. + +**Caution:** Service account key files grant powerful access. Keep them secure and never commit them to version control. + +**Create a Service Account:** + +1. Go to the [Google Cloud Console](https://console.cloud.google.com/). +2. Navigate to the **Navigation menu (☰) > IAM & Admin > Service Accounts**. +3. Click **"+ CREATE SERVICE ACCOUNT"**. +4. Enter a **Service account name** (e.g., `my-local-api-sa`) and an optional description. +5. Click **"CREATE AND CONTINUE"**. + +**Grant Permissions (Roles):** + +1. In the "Grant this service account access to project" section, select the roles necessary for your application to interact with the required APIs. Grant the **principle of least privilege**: only give the permissions absolutely needed. + * For example, if you're using Cloud Storage, you might grant `Storage Object Admin` or `Storage Object Viewer`. If using Vertex AI, you might grant `Vertex AI User` or `Vertex AI Administrator`. +2. Click **"CONTINUE"**. + +**Create a Key (JSON File):** + +1. In the "Grant users access to this service account" section (optional, skip if not needed), click **"DONE"**. +2. On the Service Accounts page, find your newly created service account. +3. Click on the **three dots (⋮)** under the "Actions" column for your service account. +4. Select **"Manage keys"**. +5. Click **"ADD KEY" > "Create new key"**. +6. Choose **"JSON"** as the key type. +7. Click **"Create"**. +8. Your browser will download a JSON key file (e.g., `your-service-account-name-xxxx.json`). **Save this file in a secure location on your local machine.** + +**Set the `GOOGLE_APPLICATION_CREDENTIALS` Environment Variable:** + +Most Google Cloud client libraries automatically detect credentials if the `GOOGLE_APPLICATION_CREDENTIALS` environment variable is set to the path of your service account JSON key file. + +* **Linux/macOS:** + + ```bash + export GOOGLE_APPLICATION_CREDENTIALS="/path/to/your-service-account-key.json" + ``` +**Important:** This environment variable is usually session-specific. For persistent access, you might add this line to your shell's profile file (e.g., `.bashrc`, `.zshrc`, `~/.profile` for Linux/macOS, or configure it in your IDE/system environment variables). However, for development, setting it per-session or through your IDE's run configurations is often safer. + + +### Installation +1. Navigate to your ComfyUI `custom_nodes` directory. +2. Add the contents of this repository to the `custom_nodes` directory. +3. Set configurations in `config.yaml` file. You can find these values in your Google Cloud Console. 4. Install the dependencies present in the `requirements.txt` file. You can do this by running the following command in your terminal: ```bash @@ -35,6 +144,10 @@ This pack currently includes the following nodes: ``` 5. **Restart the ComfyUI server.** +### Installation from git +1. Navigate to your ComfyUI `custom_nodes` directory. +2. Clone the git repo under `custom_nodes` directory. +3. **Restart the ComfyUI server.** ### Usage After installing and restarting ComfyUI, you will find the "VertexAI" category added to the node menu. diff --git a/tools/comfyui_custom_nodes/__init__.py b/tools/comfyui_custom_nodes/__init__.py index c7c3367ecb1..f030fa586a0 100644 --- a/tools/comfyui_custom_nodes/__init__.py +++ b/tools/comfyui_custom_nodes/__init__.py @@ -15,8 +15,8 @@ """Init nodes""" -from .src.vertexai_nodes.nodes.flash import ( - GeminiFlashNode +from .src.vertexai_nodes.nodes.gemini import ( + GeminiNode ) from .src.vertexai_nodes.nodes.imagen import ( @@ -49,7 +49,7 @@ ) from .src.vertexai_nodes.nodes.veo import ( - Veo2Node, + VeoNode, VideoPreviewNode, ImageToBase64Node ) @@ -60,8 +60,8 @@ # that will be displayed in the UI. NODE_CLASS_MAPPINGS = { - "Gemini Flash 2.0": GeminiFlashNode, - "Veo 2 Video Generation": Veo2Node, + "Gemini": GeminiNode, + "Veo Video Generation": VeoNode, "VideoPreviewNode": VideoPreviewNode, "Image to B64 Node": ImageToBase64Node, diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/__init__.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/__init__.py index f2e07ac2182..c2c64ee64c0 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/__init__.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/__init__.py @@ -11,81 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - - -from .nodes.flash import ( - GeminiFlashNode -) - -from .nodes.imagen_bg_swap import ( - ProductBGSwapMaskNode, - ProductBGSwapAutoMaskNode -) - -from .nodes.imagen_custom_endpoint import ( - ImagenCustomEndpointNode -) - -from .nodes.imagen_inpaint_insert import ( - InpaintInsertMaskNode, - InpaintInsertAutoMaskNode, - InpaintInsertSemanticMaskNode -) - -from .nodes.imagen_inpaint_remove import ( - InpaintRemoveMaskNode, - InpaintRemoveAutoMaskNode, - InpaintRemoveSemanticMaskNode -) - -from .nodes.imagen_maskfree_editing import ( - MaskFreeEditNode -) - -from .nodes.imagen_outpaint import ( - OutpaintingNode -) - -from .nodes.imagen import ( - Imagen3Node -) - -from .nodes.veo import ( - Veo2Node, - VideoPreviewNode, - ImageToBase64Node -) - - -WEB_DIRECTORY = "./web" - - -NODE_CLASS_MAPPINGS = { - - "Gemini Flash 2.0": GeminiFlashNode, - "Veo 2 Video Generation": Veo2Node, - "VideoPreviewNode": VideoPreviewNode, - "Image to B64 Node": ImageToBase64Node, - - # Imagen3 nodes - "Image Generation 3": Imagen3Node, - - "Inpaint Insert w Mask": InpaintInsertMaskNode, - "Inpaint Insert w AutoMask": InpaintInsertAutoMaskNode, - "Inpaint Insert w SemanticMask": InpaintInsertSemanticMaskNode, - - "Inpaint Remove w Mask": InpaintRemoveMaskNode, - "Inpaint Remove w AutoMask": InpaintRemoveAutoMaskNode, - "Inpaint Remove w SemanticMask": InpaintRemoveSemanticMaskNode, - - "Imagen Product Background Swap w Mask": ProductBGSwapMaskNode, - "Imagen Product Background Swap w AutoMask": ProductBGSwapAutoMaskNode, - - "Imagen Mask-Free Editing": MaskFreeEditNode, - - "Imagen Outpainting": OutpaintingNode, - - "Imagen Custom Endpoint": ImagenCustomEndpointNode, -} - -__all__ = ['NODE_CLASS_MAPPINGS', 'WEB_DIRECTORY'] diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/config.yaml b/tools/comfyui_custom_nodes/src/vertexai_nodes/config.yaml new file mode 100644 index 00000000000..604429d7e9b --- /dev/null +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/config.yaml @@ -0,0 +1,48 @@ +# Copyright 2025 Google LLC + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +PROJECT_ID: kr-ais-demo +REGION: us-central1 +FLASH_MODELS: + - gemini-2.0-flash-001 + - gemini-2.5-pro-preview-05-06 + - gemini-2.5-flash-preview-04-17 +IMAGEN3_EDIT_MODELS: + - imagen-3.0-capability-001 +SAFETY_FILTER_LEVELS: + - BLOCK_LOW_AND_ABOVE + - BLOCK_MEDIUM_AND_ABOVE + - BLOCK_ONLY_HIGH + - BLOCK_NONE +PERSON_GENERATION_MODES: + - ALLOW_ADULT + - ALLOW_ALL + - DONT_ALLOW +AUTO_MASK_MODES: + - MASK_MODE_FOREGROUND + - MASK_MODE_BACKGROUND +IMAGEN3_GENERATION_MODELS: + - imagen-3.0-generate-002 +VEO_MODELS: + - veo-2.0-generate-001 + - veo-3.0-generate-preview +VEO_PERSON_GENERATION_MODES: + - allow_adult + - disallow +VEO_ASPECT_RATIOS: + - "16:9" + - "9:16" +IMAGEN_MODELS: + - imagen-3.0-generate-002 \ No newline at end of file diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/modules/consts.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/modules/consts.py index 00eba01c837..7fff54c50c6 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/modules/consts.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/modules/consts.py @@ -16,24 +16,41 @@ "Consts used in custom nodes" import os +from ast import literal_eval +# This will be replaced by Cloud Run service.yaml +from ..modules.utils import load_config_and_set_env_vars -PROJECT_ID = os.environ.get("PROJECT_ID") -REGION = os.environ.get("REGION", "us-central1") - -FLASH_MODELS = ['gemini-2.0-flash-001'] -IMAGEN_MODELS = ['imagen-3.0-generate-002'] -IMAGEN3_EDIT_MODELS = ['imagen-3.0-capability-001'] -IMAGEN3_GENERATION_MODELS = ['imagen-3.0-generate-002'] -VEO_MODELS = ['veo-2.0-generate-001'] - -AUTO_MASK_MODES = ['MASK_MODE_FOREGROUND', 'MASK_MODE_BACKGROUND'] -VEO_ASPECT_RATIOS = ["16:9", "9:16"] - -PERSON_GENERATION_MODES = ['ALLOW_ADULT', 'ALLOW_ALL', 'DONT_ALLOW'] -VEO_PERSON_GENERATION_MODES = ['allow_adult', 'disallow'] - -SAFETY_FILTER_LEVELS = ['BLOCK_LOW_AND_ABOVE', 'BLOCK_MEDIUM_AND_ABOVE', - 'BLOCK_ONLY_HIGH', 'BLOCK_NONE'] +# Check if PROJECT_ID or REGION are set as environment variables. +# If not, load the config.yaml file. +if ( + not os.environ.get("PROJECT_ID") + or not os.environ.get("REGION") + or not os.environ.get("FLASH_MODELS") +): + config_path = os.path.join(os.path.dirname(__file__), '../config.yaml') + load_config_and_set_env_vars(config_path=config_path) +PROJECT_ID = os.environ.get("PROJECT_ID") +REGION = os.environ.get("REGION") +FLASH_MODELS = literal_eval(os.environ.get("FLASH_MODELS")) +IMAGEN3_EDIT_MODELS = literal_eval(os.environ.get("IMAGEN3_EDIT_MODELS")) +SAFETY_FILTER_LEVELS = literal_eval(os.environ.get("SAFETY_FILTER_LEVELS")) +PERSON_GENERATION_MODES = literal_eval( + os.environ.get("PERSON_GENERATION_MODES") +) +AUTO_MASK_MODES = literal_eval(os.environ.get("AUTO_MASK_MODES")) +IMAGEN3_GENERATION_MODELS = literal_eval( + os.environ.get("IMAGEN3_GENERATION_MODELS") +) +VEO_MODELS = literal_eval(os.environ.get("VEO_MODELS")) +VEO_PERSON_GENERATION_MODES = literal_eval( + os.environ.get("VEO_PERSON_GENERATION_MODES") +) +VEO_ASPECT_RATIOS = literal_eval( + os.environ.get("VEO_ASPECT_RATIOS") +) +IMAGEN_MODELS = literal_eval( + os.environ.get("IMAGEN_MODELS") +) INPAINT_CATEGORY = "VertexAI/Imagen3/Inpainting" diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/modules/utils.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/modules/utils.py index 11b7e2c7c6c..26594aecf59 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/modules/utils.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/modules/utils.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - "Utils used in custom nodes" import logging diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/flash.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/gemini.py similarity index 94% rename from tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/flash.py rename to tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/gemini.py index ee3f160830a..c1b3edd3b95 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/flash.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/gemini.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + """Custom ComfyUI Node for Google Gemini Flash via Vertex AI""" from typing import Tuple, Dict, Any, Optional, List @@ -34,10 +35,11 @@ ) -class GeminiFlashNode: +class GeminiNode: """ - A ComfyUI node to interact with Google's Gemini Flash model via Vertex AI. - Takes a text prompt and generation parameters, returns the generated text. + A ComfyUI node to interact with Gemini Flash, Pro model via Vertex AI. + Takes a text, image, video prompt and generation parameters, + returns the generated text. """ CATEGORY = "VertexAI" @@ -153,7 +155,7 @@ def execute_generation(self, error_message = "Gemini Model not initialized. \ Check errors during ComfyUI startup/console logs." print(f"Error in execute_generation: {error_message}") - return (error_message,) + raise RuntimeError(error_message) gen_config_dict: Dict[str, Any] = { "temperature": temperature, @@ -236,13 +238,15 @@ def execute_generation(self, print( f"Warning: Could not extract text directly. {error_info}.\ Full Response: {response}") - generated_text = f"Error: {error_info}" + error_message = f"Error: {error_info}" + raise RuntimeError(error_message) except Exception as e: # pylint: disable=W0718 # Catch other unexpected errors during text extraction print( f"Error processing response: {e}. \ Full Response: {response}") - generated_text = f"Error: Failed to process response - {e}" + error_message = f"Error: Failed to process response - {e}" + raise RuntimeError(error_message) return (generated_text,) @@ -265,4 +269,5 @@ def execute_generation(self, error_message += " - The model may be temporarily overloaded \ or you've hit a rate limit. Try again later." - return (error_message,) + print(f"Critical Error: {error_message}") + raise RuntimeError(error_message) diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen.py index 65035641feb..8592332286d 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen.py @@ -105,5 +105,6 @@ def generate_image(self, return (torch.stack(image_tensors),) except Exception as e: # pylint: disable=W0718 - print(f"Error generating image: {e}") - return (torch.empty(1, 512, 512, 3),) + error_message = f"Error generating image: {e}" + print(error_message) + raise RuntimeError(error_message) diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_bg_swap.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_bg_swap.py index 60fa311acc2..ea0de6e8e0e 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_bg_swap.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_bg_swap.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + """Custom ComfyUI Nodes for Vertex AI Imagen3 Product Background Swap""" from google.genai.types import ( diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_inpaint_insert.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_inpaint_insert.py index 42b3486049a..f2d98cbf239 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_inpaint_insert.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_inpaint_insert.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + """ Custom ComfyUI Nodes for Vertex AI Imagen3 Image Inpainting. """ diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_inpaint_remove.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_inpaint_remove.py index f7087185386..f12e8ddeed8 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_inpaint_remove.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_inpaint_remove.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + """Custom ComfyUI Nodes for Imagen3 Image Removal (Inpainting based)""" from google.genai.types import ( diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_outpaint.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_outpaint.py index 5cda2bd2c01..21467f28fd9 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_outpaint.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/imagen_outpaint.py @@ -13,6 +13,7 @@ # limitations under the License. + from PIL import Image as PIL_Image, ImageDraw as PIL_ImageDraw from google.genai.types import ( diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/veo.py b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/veo.py index aedc8698c6f..8904597aa3a 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/veo.py +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/nodes/veo.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + """Custom ComfyUI Node for Google Veo 2 via Google GenAI.""" import time @@ -40,7 +41,7 @@ ) -class Veo2Node: +class VeoNode: """ ComfyUI node to generate videos using Google's Veo 2 model. It initiates a long-running operation, polls for completion, and returns @@ -144,13 +145,13 @@ def generate_video(self, error_message = "Google GenAI client is not initialized. \ Ensure the GOOGLE_API_KEY environment variable is set" print(f"[Veo2Node] Error: {error_message}") - return (error_message,) + raise RuntimeError(error_message) if not is_valid_gcs_uri(output_gcs_uri): error_message = f"Invalid 'output_gcs_uri': '{output_gcs_uri}'. \ Must start with 'gs://', specify a bucket/path, \ and not end with '/'. Example: gs://bucket/path/prefix" - return (error_message,) + raise RuntimeError(error_message) config = GenerateVideosConfig( number_of_videos=number_of_videos, @@ -213,7 +214,7 @@ def generate_video(self, error_message = f"Operation {operation_name} not found \ during polling. It might have expired or been deleted." print(f"[Veo2Node] Error: {error_message}") - return (error_message,) + raise RuntimeError(error_message) except Exception as poll_error: # pylint: disable=W0718 # Log non-critical polling errors but continue polling print( @@ -226,7 +227,7 @@ def generate_video(self, error_message = f"Video generation failed. \ Error: {operation.error}" print(f"[Veo2Node] Error: {error_message}") - return (error_message,) + raise RuntimeError(error_message) print(operation.response) @@ -250,38 +251,35 @@ def generate_video(self, print( f"[Veo2Node] Warning: {error_message}. \ Full response: {operation.response}") - return (error_message,) + raise RuntimeError(error_message) except api_exceptions.PermissionDenied as e: error_message = f"Permission Denied: Check API key permissions \ and GCS write access to '{output_gcs_uri}'. Details: {e}" print(f"[Veo2Node] Error: {error_message}") - return (error_message,) + raise RuntimeError(error_message) except api_exceptions.InvalidArgument as e: error_message = f"InvalidArgument: Model name('{model_name}'), \ parameters, prompt, or GCS path format. Details: {e}" print(f"[Veo2Node] Error: {error_message}") - return (error_message,) + raise RuntimeError(error_message) except api_exceptions.ResourceExhausted as e: error_message = f"Quota Exceeded: Check Google Cloud API \ quotas for the GenAI service. Details: {e}" print(f"[Veo2Node] Error: {error_message}") - return (error_message,) + raise RuntimeError(error_message) except api_exceptions.NotFound as e: error_message = f"Not Found: Ensure the model '{model_name}' \ is available or the specified GCS path exists. Details: {e}" print(f"[Veo2Node] Error: {error_message}") - return (error_message,) + raise RuntimeError(error_message) except Exception as e: # pylint: disable=W0718 error_message = f"An unexpected error occurred: \ {type(e).__name__} - {e}" if operation and hasattr(operation, 'name'): error_message += f" (Related Operation: {operation.name})" print(f"[Veo2Node] Error: {error_message}") - - # import traceback - # print(traceback.format_exc()) - return (error_message,) + raise RuntimeError(error_message) class VideoPreviewNode: @@ -347,8 +345,9 @@ def download_video_from_gcs_to_tempfile(self, gcs_urls): return file_paths except Exception as e: # pylint: disable=W0718 - print(f"Error downloading video from GCS: {e}") - return None + error_message = f"Error downloading video from GCS: {e}" + print(error_message) + raise RuntimeError(error_message) def load_video(self, gcs_urls): "Load video from temp file" @@ -397,7 +396,8 @@ def encode_image_to_base64(self, image: torch.Tensor) -> Tuple[str]: or an error message. """ if not isinstance(image, torch.Tensor): - return ("Error: Input 'image' is not a valid tensor.",) + error_message = "Error: Input 'image' must be a torch.Tensor." + raise RuntimeError(error_message) if image.dim() == 4 and image.shape[0] > 0: # Select the first image from the batch @@ -406,8 +406,9 @@ def encode_image_to_base64(self, image: torch.Tensor) -> Tuple[str]: # Assume it's a single image without batch dimension img_tensor = image else: - return (f"Error: Input image tensor has unexpected\ - dimensions: {image.shape}",) + error_message = f"Error: Input image tensor has unexpected \ + dimensions: {image.shape}" + raise RuntimeError(error_message) try: # Convert tensor values from [0.0, 1.0] to [0, 255] @@ -434,4 +435,4 @@ def encode_image_to_base64(self, image: torch.Tensor) -> Tuple[str]: except Exception as e: # pylint: disable=W0718 error_message = f"Error encoding image to base64: {e}" print(f"[ImageToBase64Node] {error_message}") - return (error_message,) + raise RuntimeError(error_message) diff --git a/tools/comfyui_custom_nodes/src/vertexai_nodes/web/js/videoPreview.js b/tools/comfyui_custom_nodes/src/vertexai_nodes/web/js/videoPreview.js index d770d8f7641..4321f79511e 100644 --- a/tools/comfyui_custom_nodes/src/vertexai_nodes/web/js/videoPreview.js +++ b/tools/comfyui_custom_nodes/src/vertexai_nodes/web/js/videoPreview.js @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. + import { app } from "../../../scripts/app.js"; import { api } from '../../../scripts/api.js'