diff --git a/README.md b/README.md index 42c951f..8a52a77 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/__init__.py b/__init__.py index b5aeb17..9ce04f3 100644 --- a/__init__.py +++ b/__init__.py @@ -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, diff --git a/gemini_diarisation_vertex.py b/gemini_diarisation_vertex.py index e938727..eb198f2 100644 --- a/gemini_diarisation_vertex.py +++ b/gemini_diarisation_vertex.py @@ -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( diff --git a/gemini_node_vertex.py b/gemini_node_vertex.py index 0888184..c928a9e 100644 --- a/gemini_node_vertex.py +++ b/gemini_node_vertex.py @@ -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) diff --git a/gemini_segment_vertex.py b/gemini_segment_vertex.py deleted file mode 100644 index d3944cc..0000000 --- a/gemini_segment_vertex.py +++ /dev/null @@ -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)"} \ No newline at end of file diff --git a/nano_banana_vertex.py b/nano_banana_vertex.py index 445b71e..47616d4 100644 --- a/nano_banana_vertex.py +++ b/nano_banana_vertex.py @@ -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())] diff --git a/pyproject.toml b/pyproject.toml index cee0cb2..1e3e73d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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) diff --git a/veo_vertex.py b/veo_vertex.py index e286948..a2fdb1e 100644 --- a/veo_vertex.py +++ b/veo_vertex.py @@ -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}),