nodes/Aryan185-ComfyUI-VertexAPI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor and update Vertex API nodes: remove Gemini Segmentation, enhance model options, and update version to 1.1.3

Browse Source
This commit is contained in:
Aryan185
2026-07-24 15:33:35 +05:30
parent 1322cd95ad
commit 339846093a
8 changed files with 58 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
-13
View File
@@ -5,7 +5,6 @@ A collection of powerful custom nodes for ComfyUI that connect your local workfl
## Key Features
* **Gemini Chat (Vertex AI):** Google's powerful multimodal AI via Vertex AI. Ask questions about an image, generate detailed descriptions, or create prompts for other models. Supports thinking budget controls.
* **Gemini Segmentation (Vertex AI):** Generate segmentation masks for objects in an image using Gemini on Vertex AI.
* **Google Imagen Edit (Vertex AI):** Perform advanced image editing, inpainting, outpainting, and background swapping using Imagen on Google's Vertex AI platform.
* **Google Imagen Generator (Vertex AI):** Create images with Google's Imagen models via Vertex AI.
* **Nano Banana (Vertex AI):** A creative image generation node using a specialized Gemini model on Vertex AI.
@@ -65,18 +64,6 @@ A versatile node for text generation and image analysis using Vertex AI.
* **Output:**
* `response`: The generated text.
### Gemini Segmentation (Vertex AI)
Generate segmentation masks using Gemini on Vertex AI.
* **Category:** `image/generation`
* **Inputs:**
* `image`: Source image.
* `segment_prompt`: Description of objects to segment.
* `project_id`, `location`, `service_account`: Vertex AI credentials.
* `model`: Gemini model to use.
* **Output:**
* `mask`: Segmentation mask.
### Google Imagen Edit (Vertex AI)
__init__.py
-2
View File
@@ -1,6 +1,5 @@
from . import (
gemini_node_vertex,
gemini_segment_vertex,
nano_banana_vertex,
gemini_tts_vertex,
veo_vertex,
@@ -14,7 +13,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {}
modules = [
gemini_node_vertex,
gemini_segment_vertex,
nano_banana_vertex,
gemini_tts_vertex,
veo_vertex,
gemini_diarisation_vertex.py
+28 -4
View File
@@ -1,4 +1,3 @@
import os
import json
import io
import re
@@ -29,9 +28,21 @@ class GeminiDiarisationNode:
"asia-northeast3", "asia-south1", "asia-south2", "asia-southeast1",
"asia-southeast2", "australia-southeast1", "australia-southeast2",
"me-central1", "me-central2", "me-west1"
], {"default": "us-central1"}),
], {"default": "global"}),
"service_account": ("STRING", {"multiline": True, "default": ""}),
"model": (["gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.5-flash-lite", "gemini-3.5-flash", "gemini-3.1-pro-preview", "gemini-3.1-flash-lite-preview", "gemini-3-flash-preview", "gemini-flash-latest", "gemini-flash-lite-latest", "gemini-2.0-flash", "gemini-2.0-flash-lite"], {"default": "gemini-2.5-flash"}),
"model": ([
"gemini-3.6-flash",
"gemini-3.5-flash",
"gemini-3.5-flash-lite",
"gemini-3.1-pro-preview",
"gemini-3.1-flash-lite",
"gemini-3-flash-preview",
"gemini-2.5-flash",
"gemini-2.5-pro",
"gemini-2.5-flash-lite",
"gemini-flash-latest",
"gemini-flash-lite-latest"
], {"default": "gemini-3.5-flash"}),
"seed": ("INT", {"default": 69, "min": 0, "max": 2147483646, "step": 1}),
"temperature": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 2.0, "step": 0.1})
},
@@ -141,11 +152,24 @@ class GeminiDiarisationNode:
*You must PASS this benchmark to be deployed*"""
model_lower = model.lower()
t_config = None
# Check if the chosen model is a Pro model (requires thinking configurations to run properly)
is_pro = "pro" in model_lower
if is_pro or thinking:
final_budget = thinking_budget if thinking else 0
if is_pro and final_budget == 0:
print("Pro models cannot have thinking turned off - defaulting thinking budget to -1")
final_budget = -1
t_config = types.ThinkingConfig(include_thoughts=False, thinking_budget=final_budget)
config = types.GenerateContentConfig(
temperature=temperature,
seed=seed,
audio_timestamp=audio_timestamp if audio_timestamp else None,
thinking_config=types.ThinkingConfig(include_thoughts=False, thinking_budget=thinking_budget) if thinking else None
thinking_config=t_config
)
response = client.models.generate_content(
gemini_node_vertex.py
+21 -15
View File
@@ -27,9 +27,21 @@ class GeminiChatVertexNode:
"asia-northeast3", "asia-south1", "asia-south2", "asia-southeast1",
"asia-southeast2", "australia-southeast1", "australia-southeast2",
"me-central1", "me-central2", "me-west1"
], {"default": "us-central1"}),
], {"default": "global"}),
"service_account": ("STRING", {"multiline": True, "default": ""}),
"model": (["gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.5-flash-lite", "gemini-3.5-flash", "gemini-3.1-pro-preview", "gemini-3.1-flash-lite-preview", "gemini-3-flash-preview", "gemini-flash-latest", "gemini-flash-lite-latest", "gemini-2.0-flash", "gemini-2.0-flash-lite"], {"default": "gemini-2.5-flash"}),
"model": ([
"gemini-3.6-flash",
"gemini-3.5-flash",
"gemini-3.5-flash-lite",
"gemini-3.1-pro-preview",
"gemini-3.1-flash-lite",
"gemini-3-flash-preview",
"gemini-2.5-flash",
"gemini-2.5-pro",
"gemini-2.5-flash-lite",
"gemini-flash-latest",
"gemini-flash-lite-latest"
], {"default": "gemini-3.5-flash"}),
"temperature": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 2.0, "step": 0.1}),
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
"thinking": ("BOOLEAN", {"default": True}),
@@ -107,20 +119,14 @@ class GeminiChatVertexNode:
model_lower = model.lower()
t_config = None
if "gemini-2.0" in model_lower:
print("Gemini-2.0 models do not support thinking - disabling thinking config")
else:
final_budget = 0
# Check if the chosen model is a Pro model (requires thinking configurations to run properly)
is_pro = "pro" in model_lower
if not thinking:
if "gemini-2.5-pro" in model_lower or "gemini-3.1-pro-preview" in model_lower:
print("Pro models cannot have thinking turned off - defaulting thinking budget to -1")
final_budget = -1
else:
final_budget = thinking_budget
if ("gemini-2.5-pro" in model_lower or "gemini-3.1-pro-preview" in model_lower) and final_budget == 0:
print("Pro models cannot have thinking turned off - defaulting thinking budget to -1")
final_budget = -1
if is_pro or thinking:
final_budget = thinking_budget if thinking else 0
if is_pro and final_budget == 0:
print("Pro models cannot have thinking turned off - defaulting thinking budget to -1")
final_budget = -1
t_config = types.ThinkingConfig(thinking_budget=final_budget)
gemini_segment_vertex.py
-153
View File
@@ -1,153 +0,0 @@
import base64
import io
import json
import re
import numpy as np
import torch
from PIL import Image
from google import genai
from google.genai import types
from google.oauth2 import service_account
class GeminiSegmentationVertexNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"segment_prompt": ("STRING", {"default": "all objects", "multiline": True}),
"project_id": ("STRING", {"multiline": False, "default": ""}),
"location": ([
"global", "us-central1", "us-east1", "us-east4", "us-east5", "us-south1",
"us-west1", "us-west2", "us-west3", "us-west4",
"northamerica-northeast1", "northamerica-northeast2",
"southamerica-east1", "southamerica-west1", "africa-south1",
"europe-west1", "europe-north1", "europe-west2", "europe-west3",
"europe-west4", "europe-west6", "europe-west8", "europe-west9",
"europe-west12", "europe-southwest1", "europe-central2",
"asia-east1", "asia-east2", "asia-northeast1", "asia-northeast2",
"asia-northeast3", "asia-south1", "asia-south2", "asia-southeast1",
"asia-southeast2", "australia-southeast1", "australia-southeast2",
"me-central1", "me-central2", "me-west1"
], {"default": "us-central1"}),
"service_account": ("STRING", {"multiline": True, "default": ""}),
"model": ([
"gemini-2.5-flash-lite",
"gemini-2.5-flash",
"gemini-2.5-pro",
"gemini-2.0-flash-lite",
"gemini-2.0-flash"
], {"default": "gemini-2.5-flash"}),
"temperature": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 2.0, "step": 0.1}),
"thinking": ("BOOLEAN", {"default": False}),
"seed": ("INT", {"default": 69, "min": -1, "max": 2147483646, "step": 1}),
},
"optional": {
"thinking_budget": ("INT", {"default": 0, "min": -1, "max": 24576, "step": 1}),
}
}
RETURN_TYPES = ("MASK",)
RETURN_NAMES = ("mask",)
FUNCTION = "generate_segmentation"
CATEGORY = "image/generation"
def setup_client(self, service_account_json, project_id, location):
if not service_account_json.strip():
raise ValueError("Service account JSON content is required.")
if not project_id.strip():
raise ValueError("Project ID is required.")
try:
sa_info = json.loads(service_account_json)
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON content: {str(e)}")
credentials = service_account.Credentials.from_service_account_info(
sa_info,
scopes=["https://www.googleapis.com/auth/cloud-platform"]
)
return genai.Client(
vertexai=True,
project=project_id.strip(),
location=location.strip(),
credentials=credentials,
http_options=types.HttpOptions(
retry_options=types.HttpRetryOptions(attempts=10, jitter=10)
)
)
def generate_segmentation(self, image, segment_prompt, project_id, location, service_account, model, temperature, thinking, seed, thinking_budget=0):
client = self.setup_client(service_account, project_id, location)
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
orig_img = Image.fromarray(img_np)
orig_w, orig_h = orig_img.size
scale = min(1024 / orig_w, 1024 / orig_h)
proc_img = orig_img.resize((int(orig_w * scale), int(orig_h * scale)), Image.Resampling.LANCZOS) if scale < 1 else orig_img
pw, ph = proc_img.size
img_buf = io.BytesIO()
proc_img.save(img_buf, format='PNG')
t_config = None
if "gemini-2.0" not in model.lower():
budget = thinking_budget if thinking else 0
if "gemini-2.5-pro" in model.lower() and budget <= 0:
print("Gemini-2.5-Pro enforces thinking. Defaulting to auto (-1).")
budget = -1
t_config = types.ThinkingConfig(thinking_budget=budget)
prompt = f'Give the segmentation masks for {segment_prompt}. Output a JSON list of segmentation masks where each entry contains the 2D bounding box in the key "box_2d", the segmentation mask in key "mask", and the text label in the key "label".'
response = client.models.generate_content(
model=model,
contents=[types.Content(role="user", parts=[
types.Part.from_bytes(mime_type="image/png", data=img_buf.getvalue()),
types.Part.from_text(text=prompt)
])],
config=types.GenerateContentConfig(temperature=temperature, seed=seed, thinking_config=t_config)
)
try:
txt = response.text
if "```json" in txt: txt = re.search(r"```json\n(.*)\n```", txt, re.DOTALL).group(1)
segments = json.loads(txt)
except Exception as e:
raise RuntimeError(f"Gemini API Error: {e}")
final_mask = np.zeros((ph, pw), dtype=np.uint8)
for seg in segments:
try:
ymin, xmin, ymax, xmax = seg['box_2d']
x1, y1 = int(xmin * pw / 1000), int(ymin * ph / 1000)
x2, y2 = int(xmax * pw / 1000), int(ymax * ph / 1000)
w, h = x2 - x1, y2 - y1
if w <= 0 or h <= 0: continue
mask_str = seg['mask'].split(",")[1] if "data:image" in seg['mask'] else seg['mask']
patch = Image.open(io.BytesIO(base64.b64decode(mask_str))).convert('L')
if patch.size != (w, h):
patch = patch.resize((w, h), Image.Resampling.NEAREST)
patch_arr = np.where(np.array(patch) > 128, 255, 0).astype(np.uint8)
target_slice = final_mask[y1:y2, x1:x2]
if target_slice.shape == patch_arr.shape:
np.maximum(target_slice, patch_arr, out=target_slice)
except Exception: continue
if (pw, ph) != (orig_w, orig_h):
final_mask = np.array(Image.fromarray(final_mask).resize((orig_w, orig_h), Image.Resampling.NEAREST))
return (torch.from_numpy(final_mask.astype(np.float32) / 255.0).unsqueeze(0),)
NODE_CLASS_MAPPINGS = {"GeminiSegmentationVertexNode": GeminiSegmentationVertexNode}
NODE_DISPLAY_NAME_MAPPINGS = {"GeminiSegmentationVertexNode": "Gemini Segmentation (Vertex AI)"}
nano_banana_vertex.py
+5 -5
View File
@@ -13,6 +13,7 @@ class NanoBananaVertexNode:
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"multiline": True, "default": ""}),
"project_id": ("STRING", {"multiline": False, "default": ""}),
"location": ([
"global", "us-central1", "us-east1", "us-east4", "us-east5", "us-south1",
@@ -28,16 +29,15 @@ class NanoBananaVertexNode:
"me-central1", "me-central2", "me-west1"
], {"default": "us-central1"}),
"service_account": ("STRING", {"multiline": True, "default": ""}),
"model": (["gemini-3-pro-image-preview", "gemini-2.5-flash-image", "gemini-3.1-flash-image-preview"],),
"model": (["gemini-3-pro-image", "gemini-2.5-flash-image", "gemini-3.1-flash-image", "gemini-3.1-flash-lite-image"],),
"aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9", "21:9"],),
"resolution": (["1K", "2K", "4K"], {"default": "1K"}),
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
"temperature": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_p": ("FLOAT", {"default": 0.85, "min": 0.0, "max": 1.0, "step": 0.01}),
"google_search": ("BOOLEAN", {"default": False}),
"seed": ("INT", {"default": 69, "min": -1, "max": 2147483646, "step": 1}),
},
"optional": {
"prompt": ("STRING", {"multiline": True, "default": ""}),
"system_instruction": ("STRING", {"multiline": True, "default": ""}),
"image_1": ("IMAGE",),
"image_2": ("IMAGE",),
@@ -106,7 +106,7 @@ class NanoBananaVertexNode:
tools = None
if google_search:
if "gemini-2.5" in model:
if "gemini-2.5" in model or "gemini-3.1-flash-lite-image" in model:
print(f"Ignoring google_search: {model} does not support it.")
else:
tools = [types.Tool(googleSearch=types.GoogleSearch())]
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "vertexapi"
description = "A collection of powerful custom nodes for ComfyUI that connect your local workflows to closed-source AI models via Vertex AI."
version = "1.1.2"
version = "1.1.3"
license = {file = "LICENSE"}
# classifiers = [
# # For OS-independent nodes (works on all operating systems)
veo_vertex.py
+3 -6
View File
@@ -31,15 +31,12 @@ class GoogleVeoVertexVideoGenerator:
], {"default": "us-central1"}),
"service_account": ("STRING", {"multiline": True, "default": ""}),
"model": ([
"veo-2.0-generate-001",
"veo-2.0-generate-exp",
"veo-2.0-generate-preview",
"veo-3.0-generate-001",
"veo-3.0-fast-generate-001",
"veo-3.1-generate-001",
"veo-3.1-fast-generate-001",
"veo-3.1-generate-preview",
"veo-3.1-fast-generate-preview",
"veo-3.1-lite-generate-preview"
], {"default": "veo-3.0-generate-001"}),
], {"default": "veo-3.1-generate-001"}),
"resolution": (["720p", "1080p"], {"default": "720p"}),
"aspect_ratio": (["16:9", "9:16"], {"default": "16:9"}),
"duration_seconds": ("INT", {"default": 4, "min": 4, "max": 8, "step": 1}),