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

First commit

Browse Source
This commit is contained in:
Aryan185
2025-11-27 18:04:22 +05:30
commit 08185f659c
10 changed files with 1213 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -0,0 +1 @@
/__pycache__
README.md
+221
View File
@@ -0,0 +1,221 @@
# ComfyUI-ExternalAPI-Helpers
A collection of powerful custom nodes for ComfyUI that connect your local workflows to closed-source AI models via their APIs. Use Google's Gemini, Imagen, Veo, OpenAI's GPT-Image-1, and Black Forest Labs' FLUX models directly within ComfyUI.
## Key Features
* **FLUX Kontext Pro & Max:** Image-to-image transformations using the FLUX models via the Replicate API.
* **Gemini Chat:** Google's powerful multimodal AI. Ask questions about an image, generate detailed descriptions or create prompts for other models. Supports thinking budget controls for applicable models.
* **Gemini Segmentation:** Generate segmentation masks for objects in an image using Gemini.
* **GPT Image Edit:** OpenAI's `gpt-image-1` for prompt-based image editing and inpainting. Simply mask an area and describe the change you want to see.
* **Google Imagen Generator & Edit:** Create and edit images with Google's Imagen models, with support for Vertex AI.
* **Nano Banana:** A creative image generation node using a specialized Gemini model.
* **Veo Text-to-Video:** Generate high-quality video clips from text prompts using Google's Veo model via Vertex AI.
* **ElevenLabs TTS:** Generate high-quality speech from text using ElevenLabs' diverse range of voices and models.
* **Gemini TTS:** Create speech from text using Google's Gemini models.
* **Seamless Integration:** All nodes are designed to work seamlessly with standard ComfyUI inputs (IMAGE, MASK, STRING) and outputs, allowing you to chain them into complex and creative workflows.
* **Secure & Simple:** Simply provide your API key in the node's input field to get started.
---
## 🚀 Installation
1. Navigate to your ComfyUI installation directory.
2. Go into the `custom_nodes` folder:
```bash
cd ComfyUI/custom_nodes/
```
3. Clone this repository:
```bash
git clone https://github.com/Aryan185/ComfyUI-ExternalAPI-Helpers.git
```
4. Install the required Python packages. Navigate into the newly cloned directory and use pip to install the dependencies:
```bash
cd ComfyUI-ExternalAPI-Helpers
pip install -r requirements.txt
```
5. **Restart ComfyUI.** After restarting, you should find the new nodes in the "Add Node" menu.
---
## 🔑 Prerequisites: API Keys
All nodes in this collection require API keys to function.
* **FLUX Nodes (Replicate):** You will need a [Replicate API Token](https://replicate.com/account/api-tokens).
* **Gemini, Imagen, Nano Banana, and Gemini TTS Nodes:** You will need a [Google AI Studio API Key](https://aistudio.google.com/app/api-keys).
* **GPT Image Edit Node:** You will need an [OpenAI API Key](https://platform.openai.com/api-keys).
* **ElevenLabs TTS Node:** You will need an [ElevenLabs API Key](https://elevenlabs.io/).
* **Vertex AI Nodes (Imagen Edit, Veo):** You will need a Google Cloud Project ID, a service account with appropriate permissions, and the location for the resources.
You can paste your key directly into the `api_key` field on the corresponding node. For Vertex AI nodes, you will need to provide the project ID, location, and path to your service account JSON file.
---
## 📚 Node Guide
### Flux Kontext Pro / Max
These nodes allow you to transform an input image based on a text prompt. They are ideal for applying artistic styles or making significant conceptual changes to an existing image.
* **Category:** `image/edit`
* **Inputs:**
* `image`: The source image to transform.
* `prompt`: A text description of the desired output (e.g., "A vibrant Van Gogh painting", "Make this a 90s cartoon").
* `replicate_api_token`: Your API token from Replicate.
* `aspect_ratio`: The desired output aspect ratio. `match_input_image` is highly recommended to preserve the original composition.
* `output_format`: `jpg` or `png`.
* `safety_tolerance`: Adjust the content safety filter level.
* **Output:**
* `image`: The generated image.
### Gemini Chat
A versatile node for text generation and image analysis. Use it to understand an image's content or to generate creative text for other nodes.
* **Category:** `text/generation`
* **Inputs:**
* `prompt`: The text prompt or question you want to ask the model.
* `image` (Optional): An input image for the model to analyze.
* `api_key`: Your API key from Google AI Studio.
* `model`: The Gemini model to use (e.g., `gemini-2.5-pro`).
* `system_instruction` (Optional): Provide context or rules for how the model should behave.
* `temperature`: Controls the creativity of the output. Higher is more creative.
* `thinking`: Enables the model's thinking/reasoning process (Gemini 2.5 Pro).
* **Output:**
* `response`: The text generated by the Gemini model.
### Gemini Segmentation
This node uses a Gemini model to generate segmentation masks for specified objects within an image.
* **Category:** `image/generation`
* **Inputs:**
* `image`: The source image for segmentation.
* `segment_prompt`: A text description of the objects to segment (e.g., "the car", "all people").
* `api_key`: Your API key from Google AI Studio.
* `model`: The Gemini model to use.
* `...other_params`: Controls for temperature, thinking, and seed.
* **Output:**
* `mask`: A black and white mask of the segmented objects.
### GPT Image Edit
This node uses OpenAI's API to perform powerful, prompt-based inpainting and editing.
* **Category:** `image/edit`
* **Inputs:**
* `image`: The source image to edit.
* `mask` (Optional): A black and white mask. The model will edit the **white area** of the mask.
* `prompt`: A description of the edit to perform (e.g., "Add a small red boat on the water", "Remove the person on the left").
* `api_key`: Your API key from OpenAI.
* `...other_params`: Various quality and formatting options for the OpenAI API.
* **Output:**
* `image`: The edited image.
**Note:** If a mask is provided, the edits will be constrained to the masked region. If no mask is provided, the model will attempt to edit the entire image based on the prompt.
### Google Imagen Generator
Generate images from a text prompt using Google's Imagen models.
* **Category:** `image/generation`
* **Inputs:**
* `prompt`: A text description of the image to generate.
* `api_key`: Your API key from Google AI Studio.
* `model`: The Imagen model to use.
* `...other_params`: Options for number of images, aspect ratio, and image size.
* **Output:**
* `images`: The generated image(s).
### Google Imagen Edit (Vertex AI only)
Perform advanced image editing, inpainting, outpainting, and background swapping using Imagen on Google's Vertex AI platform.
* **Category:** `image/edit`
* **Inputs:**
* `image`: The source image to edit.
* `mask`: A mask defining the area to edit.
* `prompt`: A description of the desired edit.
* `project_id`: Your Google Cloud Project ID.
* `location`: The Google Cloud location for the model.
* `service_account`: Path to your Google Cloud service account JSON file.
* `edit_mode`: The type of edit to perform (e.g., inpainting, outpainting).
* `...other_params`: Controls for negative prompt, seed, and steps.
* **Output:**
* `edited_images`: The edited image(s).
### Nano Banana
A creative image generation node that can take a combination of text and up to five images as input.
* **Category:** `image/generation`
* **Inputs:**
* `api_key`: Your API key from Google AI Studio.
* `prompt` (Optional): A text prompt.
* `image_1` to `image_5` (Optional): Up to five source images.
* `...other_params`: Controls for aspect ratio, temperature, top_p, and seed.
* **Output:**
* `image`: The generated image.
### Veo Text-to-Video (Vertex AI)
Generate short, high-quality video clips from a text description using Google's Veo model on Vertex AI.
* **Category:** `video/generation`
* **Inputs:**
* `prompt`: A text description of the video to generate.
* `project_id`: Your Google Cloud Project ID.
* `location`: The Google Cloud location for the model.
* `service_account`: Path to your Google Cloud service account JSON file.
* `...other_params`: Controls for negative prompt, aspect ratio, audio generation, and seed.
* **Output:**
* `frames`: The generated video frames, output as an image batch.
---
### ElevenLabs TTS
Generate speech from text using the ElevenLabs API.
* **Category:** `audio/generation`
* **Inputs:**
* `text`: The text to convert to speech.
* `api_key`: Your API key from ElevenLabs.
* `voice_id`: The ID of the voice to use for generation.
* `model_id`: The ElevenLabs model to use.
* `output_format`: The desired output audio format.
* `stability`: Controls the stability and variability of the generated speech.
* `similarity_boost`: Enhances the similarity of the generated speech to the chosen voice.
* `speed`: Adjusts the speaking rate.
* `style`: Controls the expressiveness of the speech.
* `use_speaker_boost`: A boolean to enable or disable speaker boost.
* `seed`: A seed for ensuring reproducible results.
* **Output:**
* `audio`: The generated audio waveform and sample rate.
### Gemini TTS
Generate speech from text using Google's Gemini TTS models.
* **Category:** `audio/generation`
* **Inputs:**
* `text`: The text to be converted into speech.
* `api_key`: Your API key from Google AI Studio.
* `model`: The specific Gemini model to use for generation.
* `voice_id`: The prebuilt voice to use for the output.
* `temperature`: Controls the randomness and creativity of the output.
* `seed`: A seed for ensuring reproducible results.
* `system_prompt` (Optional): A system-level instruction to guide the model's behavior.
* **Output:**
* `audio`: The generated audio waveform and sample rate.
## Acknowledgements
* The [ComfyUI](https://github.com/comfyanonymous/ComfyUI) team for creating such a flexible and powerful platform.
* [Google](https://deepmind.google/technologies/gemini/), [OpenAI](https://openai.com/), and [Black Forest Labs](https://www.blackforestlabs.ai/) for developing these incredible models.
* [Replicate](https://replicate.com/) for providing easy API access to a wide range of models.
__init__.py
+13
View File
@@ -0,0 +1,13 @@
from .gemini_node_vertex import NODE_CLASS_MAPPINGS as GEMINI_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GEMINI_DISPLAY
from .gemini_segment_vertex import NODE_CLASS_MAPPINGS as GEMINI_SEGMENT_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GEMINI_SEGMENT_DISPLAY
from .nano_banana_vertex import NODE_CLASS_MAPPINGS as NANO_BANANA_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as NANO_BANANA_DISPLAY
from .imagen_edit_vertex import NODE_CLASS_MAPPINGS as IMAGEN_EDIT_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as IMAGEN_EDIT_DISPLAY
from .gemini_tts_vertex import NODE_CLASS_MAPPINGS as GEMINI_TTS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GEMINI_TTS_DISPLAY
from .imagen import NODE_CLASS_MAPPINGS as IMAGEN_IMAGE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as IMAGEN_IMAGE_DISPLAY
NODE_CLASS_MAPPINGS = {**GEMINI_MAPPINGS, **IMAGEN_IMAGE_MAPPINGS, **IMAGEN_EDIT_MAPPINGS, **GEMINI_SEGMENT_MAPPINGS, **NANO_BANANA_MAPPINGS, **GEMINI_TTS_MAPPINGS}
NODE_DISPLAY_NAME_MAPPINGS = {**GEMINI_DISPLAY, **IMAGEN_IMAGE_DISPLAY, **IMAGEN_EDIT_DISPLAY, **GEMINI_SEGMENT_DISPLAY, **NANO_BANANA_DISPLAY, **GEMINI_TTS_DISPLAY}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
gemini_node_vertex.py
+148
View File
@@ -0,0 +1,148 @@
import os
import io
import json
import tempfile
import numpy as np
import torch
from PIL import Image
from typing import Optional
from google import genai
from google.genai import types
class GeminiChatVertexNode:
"""ComfyUI Node for Gemini Chat via Vertex AI with optional image input"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"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-3-pro-preview",
"gemini-2.0-flash-lite",
"gemini-2.0-flash"
], {"default": "gemini-2.5-pro"}),
"temperature": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 2.0, "step": 0.1}),
"thinking": ("BOOLEAN", {"default": True}),
"seed": ("INT", {"default": 69, "min": -1, "max": 2147483646, "step": 1}),
},
"optional": {
"system_instruction": ("STRING", {"multiline": True, "default": ""}),
"thinking_budget": ("INT", {"default": -1, "min": -1, "max": 24576, "step": 1}),
"image": ("IMAGE",),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("response",)
FUNCTION = "generate"
CATEGORY = "text/generation"
def setup_client(self, service_account_json, project_id, location):
"""Setup Vertex AI client with service account JSON content"""
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.")
# Validate and write JSON content to temporary file
try:
json.loads(service_account_json) # Validate JSON format
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON content: {str(e)}")
# Create temporary file with JSON content
temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
temp_file.write(service_account_json.strip())
temp_file.close()
# Set credentials path
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = temp_file.name
return genai.Client(vertexai=True, project=project_id.strip(), location=location.strip())
def generate(self, prompt: str, project_id: str, location: str, service_account: str,
model: str, temperature: float, thinking: bool, seed: int,
system_instruction: Optional[str] = None, thinking_budget: int = -1,
image: Optional[torch.Tensor] = None) -> tuple:
# Initialize Vertex AI client
client = self.setup_client(service_account, project_id, location)
parts = [types.Part.from_text(text=prompt)]
# Handle image input
if image is not None:
img_array = image.cpu().numpy() if isinstance(image, torch.Tensor) else image
if len(img_array.shape) == 4:
img_array = img_array[0]
if img_array.dtype in [np.float32, np.float64]:
img_array = (img_array * 255).astype(np.uint8)
buffered = io.BytesIO()
Image.fromarray(img_array).save(buffered, format="PNG")
parts.append(types.Part.from_bytes(mime_type="image/png", data=buffered.getvalue()))
model_lower = model.lower()
# Gemini 2.0 models don't support thinking at all
if "gemini-2.0" in model_lower:
print("Gemini-2.0 models do not support thinking - disabling thinking config")
final_thinking_budget = None
# Gemini Pro models (2.5-pro, 3-pro) cannot turn thinking off
elif "pro" in model_lower and ("2.5" in model_lower or "gemini-3" in model_lower):
print(f"{model} cannot have thinking turned off - thinking is always enabled")
final_thinking_budget = thinking_budget if thinking_budget != 0 else -1
# Flash models can toggle thinking on/off
elif not thinking:
final_thinking_budget = 0
else:
final_thinking_budget = thinking_budget
config = types.GenerateContentConfig(
temperature=temperature,
seed=seed,
response_mime_type="text/plain"
)
if "gemini-2.0" not in model_lower:
config.thinking_config = types.ThinkingConfig(thinking_budget=final_thinking_budget)
if system_instruction and system_instruction.strip():
config.system_instruction = [types.Part.from_text(text=system_instruction.strip())]
response = client.models.generate_content(
model=model,
contents=[types.Content(role="user", parts=parts)],
config=config
)
return (response.text,)
# Node mappings
NODE_CLASS_MAPPINGS = {
"GeminiChatVertexNode": GeminiChatVertexNode
}
NODE_DISPLAY_NAME_MAPPINGS = {
"GeminiChatVertexNode": "Gemini Chat (Vertex AI)"
}
gemini_segment_vertex.py
+229
View File
@@ -0,0 +1,229 @@
import base64
import os
import io
import json
import tempfile
import numpy as np
import torch
from PIL import Image
from google import genai
from google.genai import types
class GeminiSegmentationVertexNode:
"""ComfyUI Node for Gemini Image Segmentation via Vertex AI"""
@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": True}),
"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):
"""Setup Vertex AI client with service account JSON content"""
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.")
# Validate and write JSON content to temporary file
try:
json.loads(service_account_json) # Validate JSON format
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON content: {str(e)}")
# Create temporary file with JSON content
temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
temp_file.write(service_account_json.strip())
temp_file.close()
# Set credentials path
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = temp_file.name
return genai.Client(vertexai=True, project=project_id.strip(), location=location.strip())
def generate_segmentation(self, image: torch.Tensor, segment_prompt: str, project_id: str,
location: str, service_account: str, model: str,
temperature: float, thinking: bool, seed: int,
thinking_budget: int = 0) -> tuple:
# Initialize Vertex AI client
client = self.setup_client(service_account, project_id, location)
img_array = image.cpu().numpy() if isinstance(image, torch.Tensor) else image
if len(img_array.shape) == 4:
img_array = img_array[0] # Remove batch dimension
if img_array.dtype in [np.float32, np.float64]:
img_array = (img_array * 255).astype(np.uint8)
original_image = Image.fromarray(img_array).convert('RGB')
original_width, original_height = original_image.size
max_size = 1024
scale = min(max_size / original_width, max_size / original_height)
if scale < 1:
new_width = int(original_width * scale)
new_height = int(original_height * scale)
processed_image = original_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
else:
processed_image = original_image
buffer = io.BytesIO()
processed_image.save(buffer, format='PNG')
image_data = buffer.getvalue()
base_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\". Use descriptive labels. You are to only return only text output."
parts = [
types.Part.from_bytes(mime_type="image/png", data=image_data),
types.Part.from_text(text=base_prompt)
]
model_lower = model.lower()
# Gemini 2.0 models don't support thinking at all
if "gemini-2.0" in model_lower:
print("Gemini-2.0 models do not support thinking - disabling thinking config")
final_thinking_budget = None
# Gemini Pro models (2.5-pro) cannot turn thinking off
elif "pro" in model_lower and "2.5" in model_lower:
print(f"{model} cannot have thinking turned off - thinking is always enabled")
final_thinking_budget = thinking_budget if thinking_budget != 0 else -1
# Flash models can toggle thinking on/off
elif not thinking:
final_thinking_budget = 0
else:
final_thinking_budget = thinking_budget
config = types.GenerateContentConfig(
temperature=temperature,
seed=seed,
response_mime_type="text/plain",
response_modalities=["TEXT"]
)
if "gemini-2.0" not in model_lower:
config.thinking_config = types.ThinkingConfig(thinking_budget=final_thinking_budget)
# Generate content
try:
response = client.models.generate_content(
model=model,
contents=[types.Content(role="user", parts=parts)],
config=config
)
response_text = response.text
if '```json' in response_text:
response_text = response_text.split('```json')[1].split('```')[0]
segments = json.loads(response_text)
except Exception as e:
raise RuntimeError(f"Error calling Gemini API: {str(e)}")
# Create mask from segments
proc_width, proc_height = processed_image.size
mask_image = Image.new('L', (proc_width, proc_height), 0)
# Sort segments by size (largest first)
segments_with_size = []
for segment in segments:
box_2d = segment['box_2d']
ymin, xmin, ymax, xmax = box_2d
w = (xmax - xmin) / 1000
h = (ymax - ymin) / 1000
segments_with_size.append((segment, w * h))
segments_with_size.sort(key=lambda x: x[1], reverse=True)
# Process each segment
for i, (segment, _) in enumerate(segments_with_size):
try:
box_2d = segment['box_2d']
ymin, xmin, ymax, xmax = box_2d
x = int(xmin / 1000 * proc_width)
y = int(ymin / 1000 * proc_height)
w = int((xmax - xmin) / 1000 * proc_width)
h = int((ymax - ymin) / 1000 * proc_height)
mask_data = segment['mask']
if isinstance(mask_data, str):
if mask_data.startswith('data:image'):
mask_data = mask_data.split(',')[1]
mask_bytes = base64.b64decode(mask_data)
mask_img = Image.open(io.BytesIO(mask_bytes)).convert('L')
else:
continue
if mask_img.size != (w, h):
mask_img = mask_img.resize((w, h), Image.Resampling.LANCZOS)
mask_array = list(mask_img.getdata())
final_pixels = [255 if alpha > 128 else 0 for alpha in mask_array]
segment_mask = Image.new('L', (w, h))
segment_mask.putdata(final_pixels)
if x + w <= proc_width and y + h <= proc_height and x >= 0 and y >= 0:
region = mask_image.crop((x, y, x + w, y + h))
region_pixels = list(region.getdata())
segment_pixels = list(segment_mask.getdata())
combined_pixels = [max(r, s) for r, s in zip(region_pixels, segment_pixels)]
combined_region = Image.new('L', (w, h))
combined_region.putdata(combined_pixels)
mask_image.paste(combined_region, (x, y))
except Exception:
continue
if processed_image.size != original_image.size:
mask_image = mask_image.resize(original_image.size, Image.Resampling.LANCZOS)
# Convert PIL mask to ComfyUI mask format
mask_array = np.array(mask_image, dtype=np.float32) / 255.0
mask_tensor = torch.from_numpy(mask_array).unsqueeze(0) # Add batch dimension
return (mask_tensor,)
# Node mappings
NODE_CLASS_MAPPINGS = {"GeminiSegmentationVertexNode": GeminiSegmentationVertexNode}
NODE_DISPLAY_NAME_MAPPINGS = {"GeminiSegmentationVertexNode": "Gemini Segmentation (Vertex AI)"}
gemini_tts_vertex.py
+146
View File
@@ -0,0 +1,146 @@
import os
import struct
import io
import json
import tempfile
import torchaudio
from google.genai import Client, types
class GeminiTTSVertexNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING", {"multiline": True, "default": ""}),
"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-preview-tts", "gemini-2.5-pro-preview-tts"],),
"voice_id": (["Zephyr", "Puck", "Charon", "Kore", "Fenrir", "Leda", "Orus", "Aoede", "Callirrhoe", "Autonoe", "Enceladus", "Iapetus", "Umbriel", "Algieba", "Despina", "Erinome", "Achernar", "Laomedeia", "Rasalgethi", "Algenib", "Achird", "Pulcherrima", "Gacrux", "Schedar", "Alnilam", "Sulafat", "Sadaltager", "Sadachbia", "Vindemiatrix", "Zubenelgenubi"],),
"seed": ("INT", {"default": 69, "min": -1, "max": 2147483646, "step": 1}),
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
},
"optional": {
"system_prompt": ("STRING", {"multiline": True, "default": ""}),
}
}
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("audio",)
FUNCTION = "generate_speech"
CATEGORY = "audio/generation"
def setup_client(self, service_account_json, project_id, location):
"""Setup Vertex AI client with service account JSON content"""
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.")
# Validate and write JSON content to temporary file
try:
json.loads(service_account_json) # Validate JSON format
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON content: {str(e)}")
# Create temporary file with JSON content
temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
temp_file.write(service_account_json.strip())
temp_file.close()
# Set credentials path
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = temp_file.name
return Client(
vertexai=True,
project=project_id.strip(),
location=location.strip(),
http_options=types.HttpOptions(
retry_options=types.HttpRetryOptions(attempts=10, jitter=10)
)
)
def generate_speech(self, text, project_id, location, service_account, voice_id,
temperature, model, seed, system_prompt=""):
if not text.strip():
raise ValueError("Text input cannot be empty.")
# Initialize Vertex AI client
client = self.setup_client(service_account, project_id, location)
# Build prompt
prompt_text = text
if system_prompt.strip():
prompt_text = system_prompt.strip() + ":\n\n\"" + text + "\""
contents = [types.Content(role="user", parts=[types.Part.from_text(text=prompt_text)])]
config = types.GenerateContentConfig(
temperature=temperature,
seed=seed,
response_modalities=["audio"],
speech_config=types.SpeechConfig(
voice_config=types.VoiceConfig(
prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name=voice_id)
)
),
)
# Generate audio
audio_data = b""
mime_type = None
for chunk in client.models.generate_content_stream(
model=model,
contents=contents,
config=config
):
if (chunk.candidates and chunk.candidates[0].content and
chunk.candidates[0].content.parts and
chunk.candidates[0].content.parts[0].inline_data):
inline_data = chunk.candidates[0].content.parts[0].inline_data
audio_data += inline_data.data
if mime_type is None:
mime_type = inline_data.mime_type
if not audio_data:
raise ValueError("No audio data received from API.")
if mime_type and "L16" in mime_type:
header = struct.pack(
"<4sI4s4sIHHIIHH4sI",
b"RIFF", 36 + len(audio_data),
b"WAVE", b"fmt ", 16, 1, 1, 24000, 48000, 2, 16,
b"data", len(audio_data)
)
audio_data = header + audio_data
# Decode audio from memory
audio_buffer = io.BytesIO(audio_data)
waveform, sample_rate = torchaudio.load(audio_buffer)
# Return in ComfyUI audio format
return ({"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate},)
@classmethod
def IS_CHANGED(cls, **kwargs):
return f"{kwargs.get('text', '')}-{kwargs.get('voice_id', '')}-{kwargs.get('temperature', 1.0)}-{kwargs.get('model', '')}-{kwargs.get('seed', 69)}-{kwargs.get('system_prompt', '')}"
NODE_CLASS_MAPPINGS = {"GeminiTTSVertexNode": GeminiTTSVertexNode}
NODE_DISPLAY_NAME_MAPPINGS = {"GeminiTTSVertexNode": "Gemini TTS (Vertex AI)"}
imagen.py
+111
View File
@@ -0,0 +1,111 @@
import os
import torch
import numpy as np
from PIL import Image
from google import genai
from google.genai import types
class GoogleImagenNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"multiline": True, "default": "A majestic lion in the savanna"}),
"api_key": ("STRING", {"multiline": False, "default": ""}),
"model": (["models/imagen-4.0-ultra-generate-001", "models/imagen-4.0-generate-001", "models/imagen-4.0-fast-generate-001", "models/imagen-3.0-generate-002"],),
"number_of_images": ("INT", {"default": 1, "min": 1, "max": 4, "step": 1}),
"aspect_ratio": (["1:1", "9:16", "16:9", "4:3", "3:4"],),
"image_size": (["1K", "2K"],),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("images",)
FUNCTION = "generate_images"
CATEGORY = "image/generation"
def pil_to_tensor(self, images):
if not isinstance(images, list):
images = [images]
tensors = []
for image in images:
if image.mode != 'RGB':
image = image.convert('RGB')
array = np.array(image).astype(np.float32) / 255.0
tensor = torch.from_numpy(array)
tensors.append(tensor)
return torch.stack(tensors)
def generate_images(self, prompt, api_key, model, number_of_images, aspect_ratio, image_size):
try:
key = api_key.strip() or os.environ.get("GEMINI_API_KEY")
if not key:
raise ValueError("No API key provided.")
client = genai.Client(api_key=key)
config = types.GenerateImagesConfig(
number_of_images=number_of_images,
include_rai_reason=True,
output_mime_type="image/jpeg",
aspect_ratio=aspect_ratio,
)
if model in ["models/imagen-4.0-ultra-generate-001", "models/imagen-4.0-generate-001"]:
config.image_size = image_size
else:
print("2K resolution not supported for this model, using default image size")
result = client.models.generate_images(
model=model,
prompt=prompt,
config=config
)
if not result.generated_images:
raise ValueError("No images generated by the API")
pil_images = []
for generated_image in result.generated_images:
image_data = generated_image.image
if hasattr(image_data, 'mode') and hasattr(image_data, 'size'):
pil_images.append(image_data)
elif hasattr(image_data, '_pil_image'):
pil_images.append(image_data._pil_image)
elif hasattr(image_data, 'show'):
try:
from io import BytesIO
buffer = BytesIO()
image_data.save(buffer, format='PNG')
buffer.seek(0)
pil_image = Image.open(buffer)
pil_images.append(pil_image)
except:
pil_images.append(Image.new('RGB', (512, 512), color='gray'))
elif hasattr(image_data, 'read') or isinstance(image_data, bytes):
from io import BytesIO
image_bytes = image_data.read() if hasattr(image_data, 'read') else image_data
pil_images.append(Image.open(BytesIO(image_bytes)))
else:
try:
pil_images.append(Image.open(image_data))
except:
pil_images.append(Image.new('RGB', (512, 512), color='gray'))
return (self.pil_to_tensor(pil_images),)
except Exception as e:
print(f"Google Imagen Error: {str(e)}")
error_image = Image.new('RGB', (512, 512), color='black')
return (self.pil_to_tensor([error_image]),)
@classmethod
def IS_CHANGED(cls, **kwargs):
return f"{kwargs.get('prompt', '')}-{kwargs.get('model', '')}-{kwargs.get('number_of_images', 1)}-{kwargs.get('aspect_ratio', '1:1')}-{kwargs.get('image_size', '1K')}"
NODE_CLASS_MAPPINGS = {"GoogleImagenNode": GoogleImagenNode}
NODE_DISPLAY_NAME_MAPPINGS = {"GoogleImagenNode": "Google Imagen Generator"}
imagen_edit_vertex.py
+172
View File
@@ -0,0 +1,172 @@
import os
import json
import tempfile
import torch
import numpy as np
from PIL import Image
import base64
from io import BytesIO
from google import genai
from google.genai import types
from google.genai.types import RawReferenceImage, MaskReferenceImage
class GoogleImagenEditNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"mask": ("MASK",),
"prompt": ("STRING", {"multiline": True, "default": "Edit this image"}),
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
"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": ""}),
"edit_mode": (["EDIT_MODE_INPAINT_INSERTION", "EDIT_MODE_INPAINT_REMOVAL", "EDIT_MODE_OUTPAINT", "EDIT_MODE_BGSWAP"], {"default": "EDIT_MODE_INPAINT_INSERTION"}),
"number_of_images": ("INT", {"default": 1, "min": 1, "max": 4, "step": 1}),
"seed": ("INT", {"default": 12345, "min": 1, "max": 4294967295, "step": 1}),
"base_steps": ("INT", {"default": 50, "min": 10, "max": 100, "step": 1})
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("edited_images",)
FUNCTION = "edit_image"
CATEGORY = "image/edit"
def setup_client(self, service_account_json, project_id, location):
"""Setup Vertex AI client with service account JSON content"""
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.")
# Validate and write JSON content to temporary file
try:
json.loads(service_account_json) # Validate JSON format
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON content: {str(e)}")
# Create temporary file with JSON content
temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
temp_file.write(service_account_json.strip())
temp_file.close()
# Set credentials path
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = temp_file.name
return genai.Client(vertexai=True, project=project_id.strip(), location=location.strip())
def tensor_to_pil(self, tensor):
array = (tensor.cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(array)
def mask_to_pil(self, mask):
if len(mask.shape) == 3 and mask.shape[0] == 1:
mask = mask.squeeze(0)
array = (mask.cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(array, mode='L')
def pil_to_tensor(self, images):
if not isinstance(images, list):
images = [images]
tensors = []
for image in images:
if image.mode != 'RGB':
image = image.convert('RGB')
array = np.array(image).astype(np.float32) / 255.0
tensors.append(torch.from_numpy(array))
return torch.stack(tensors)
def edit_image(self, image, mask, prompt, project_id, location, service_account, edit_mode, number_of_images, negative_prompt, seed, base_steps):
try:
# Initialize Vertex AI client with improved authentication
client = self.setup_client(service_account, project_id, location)
input_image = self.tensor_to_pil(image[0])
input_mask = self.mask_to_pil(mask)
img_buffer = BytesIO()
input_image.save(img_buffer, format='PNG')
img_b64 = base64.b64encode(img_buffer.getvalue()).decode('utf-8')
mask_buffer = BytesIO()
input_mask.save(mask_buffer, format='PNG')
mask_b64 = base64.b64encode(mask_buffer.getvalue()).decode('utf-8')
raw_ref_image = RawReferenceImage(
reference_image={'image_bytes': img_b64},
reference_id=0
)
mask_ref_image = MaskReferenceImage(
reference_id=1,
reference_image={'image_bytes': mask_b64},
config=types.MaskReferenceConfig(
mask_mode="MASK_MODE_USER_PROVIDED",
mask_dilation=0.03,
),
)
config_params = {
"edit_mode": edit_mode,
"number_of_images": number_of_images,
"include_rai_reason": True,
"output_mime_type": "image/jpeg",
"base_steps": base_steps,
"seed": seed
}
if negative_prompt.strip():
config_params["negative_prompt"] = negative_prompt.strip()
response = client.models.edit_image(
model="imagen-3.0-capability-001",
prompt=prompt,
reference_images=[raw_ref_image, mask_ref_image],
config=types.EditImageConfig(**config_params),
)
if not response.generated_images:
raise ValueError("No images generated by the API")
pil_images = []
for generated_image in response.generated_images:
image_data = generated_image.image
if hasattr(image_data, 'mode') and hasattr(image_data, 'size'):
pil_images.append(image_data)
elif hasattr(image_data, '_pil_image'):
pil_images.append(image_data._pil_image)
elif hasattr(image_data, 'show'):
try:
buffer = BytesIO()
image_data.save(buffer, format='PNG')
buffer.seek(0)
pil_images.append(Image.open(buffer))
except:
pil_images.append(Image.new('RGB', (512, 512), color='gray'))
elif hasattr(image_data, 'read') or isinstance(image_data, bytes):
image_bytes = image_data.read() if hasattr(image_data, 'read') else image_data
pil_images.append(Image.open(BytesIO(image_bytes)))
else:
try:
pil_images.append(Image.open(image_data))
except:
pil_images.append(Image.new('RGB', (512, 512), color='gray'))
return (self.pil_to_tensor(pil_images),)
except Exception as e:
print(f"Google Imagen Edit Error: {str(e)}")
error_image = Image.new('RGB', (512, 512), color='black')
return (self.pil_to_tensor([error_image]),)
@classmethod
def IS_CHANGED(cls, **kwargs):
return f"{kwargs.get('prompt', '')}-{kwargs.get('negative_prompt', '')}-{kwargs.get('edit_mode', '')}-{kwargs.get('number_of_images', 1)}-{kwargs.get('seed', 12345)}-{kwargs.get('base_steps', 50)}"
NODE_CLASS_MAPPINGS = {"GoogleImagenEditNode": GoogleImagenEditNode}
NODE_DISPLAY_NAME_MAPPINGS = {"GoogleImagenEditNode": "Imagen Edit (Vertex AI)"}
nano_banana_vertex.py
+165
View File
@@ -0,0 +1,165 @@
import os
import io
import json
import tempfile
import torch
import numpy as np
from PIL import Image
from google import genai
from google.genai import types
class NanoBananaVertexNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"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-3-pro-image-preview", "gemini-2.5-flash-image"],),
"aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9", "21:9"],),
"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}),
"seed": ("INT", {"default": 69, "min": -1, "max": 2147483646, "step": 1}),
"resolution": (["1K", "2K", "4K"], {"default": "1K"}),
},
"optional": {
"prompt": ("STRING", {"multiline": True, "default": ""}),
"system_instruction": ("STRING", {"multiline": True, "default": ""}),
"image_1": ("IMAGE",),
"image_2": ("IMAGE",),
"image_3": ("IMAGE",),
"image_4": ("IMAGE",),
"image_5": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("image",)
FUNCTION = "generate"
CATEGORY = "image/generation"
def setup_client(self, service_account_json, project_id, location):
"""Setup Vertex AI client with service account JSON content"""
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.")
# Validate and write JSON content to temporary file
try:
json.loads(service_account_json) # Validate JSON format
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON content: {str(e)}")
# Create temporary file with JSON content
temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
temp_file.write(service_account_json.strip())
temp_file.close()
# Set credentials path
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = temp_file.name
return genai.Client(vertexai=True, project=project_id.strip(), location=location.strip())
def tensor_to_pil(self, tensor):
if tensor.dim() == 4:
tensor = tensor[0]
array = (tensor.cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(array)
def pil_to_tensor(self, image):
if image.mode != 'RGB':
image = image.convert('RGB')
array = np.array(image).astype(np.float32) / 255.0
tensor = torch.from_numpy(array)
return tensor.unsqueeze(0)
def generate(self, project_id, location, service_account, model, aspect_ratio,
resolution, temperature, top_p, seed,
prompt="", system_instruction="",
image_1=None, image_2=None, image_3=None, image_4=None, image_5=None):
# Initialize Vertex AI client
client = self.setup_client(service_account, project_id, location)
# Build parts list
parts = []
# Add images
for img_tensor in [image_1, image_2, image_3, image_4, image_5]:
if img_tensor is not None:
pil_img = self.tensor_to_pil(img_tensor)
buffer = io.BytesIO()
pil_img.save(buffer, format='PNG')
parts.append(types.Part.from_bytes(
mime_type="image/png",
data=buffer.getvalue()
))
# Add prompt if provided
if prompt.strip():
parts.append(types.Part.from_text(text=prompt))
if not parts:
raise ValueError("At least one image or prompt must be provided.")
contents = [types.Content(role="user", parts=parts)]
config_dict = {
"temperature": temperature,
"seed": seed,
"top_p": top_p,
"response_modalities": ["IMAGE"],
"image_config": types.ImageConfig(aspect_ratio=aspect_ratio),
}
if "gemini-3-pro" in model:
config_dict["image_config"] = types.ImageConfig(aspect_ratio=aspect_ratio, image_size=resolution)
else:
print("gemini-2.5-flash-image does not support resolution parameter, using default resolution")
config = types.GenerateContentConfig(**config_dict)
if system_instruction.strip():
config.system_instruction = [types.Part.from_text(text=system_instruction)]
# Generate
response = client.models.generate_content(
model=model,
contents=contents,
config=config,
)
result_image = None
if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
for part in response.candidates[0].content.parts:
if part.inline_data and part.inline_data.data:
result_image = Image.open(io.BytesIO(part.inline_data.data))
break
if result_image is None:
raise ValueError("No image generated by the API.")
return (self.pil_to_tensor(result_image),)
@classmethod
def IS_CHANGED(cls, **kwargs):
return f"{kwargs.get('prompt', '')}-{kwargs.get('temperature', 0.5)}-{kwargs.get('top_p', 0.85)}-{kwargs.get('seed', 69)}-{kwargs.get('aspect_ratio', '1:1')}-{kwargs.get('model', 'gemini-3-pro-image-preview')}-{kwargs.get('image_1')}-{kwargs.get('image_2')}-{kwargs.get('image_3')}-{kwargs.get('image_4')}-{kwargs.get('image_5')}"
NODE_CLASS_MAPPINGS = {"NanoBananaVertexNode": NanoBananaVertexNode}
NODE_DISPLAY_NAME_MAPPINGS = {"NanoBananaVertexNode": "Nano Banana (Vertex AI)"}
requirements.txt
+7
View File
@@ -0,0 +1,7 @@
replicate
pillow
numpy
torch
google-genai
opencv-python
torchaudio