584 lines
23 KiB
Python
584 lines
23 KiB
Python
import http.client
|
|
import json
|
|
import base64
|
|
import time
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
import io
|
|
import urllib.request
|
|
|
|
|
|
class VectorEngineBase:
|
|
"""
|
|
Base class with shared utilities for Vector Engine nodes
|
|
"""
|
|
|
|
def __init__(self):
|
|
# Read API key from environment variable
|
|
self.api_key = os.getenv("VECTOR_ENGINE_API_KEY")
|
|
if not self.api_key:
|
|
raise ValueError(
|
|
"VECTOR_ENGINE_API_KEY environment variable is not set. "
|
|
"Please set it before using this node: "
|
|
"export VECTOR_ENGINE_API_KEY='your-api-key-here'"
|
|
)
|
|
|
|
def tensor_to_base64(self, tensor, mime_type="image/jpeg", max_size=2048, quality=85):
|
|
"""
|
|
Convert ComfyUI image tensor to base64 string with optimization
|
|
"""
|
|
if len(tensor.shape) == 4:
|
|
tensor = tensor[0]
|
|
|
|
numpy_image = (tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(numpy_image)
|
|
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
|
|
width, height = pil_image.size
|
|
if max(width, height) > max_size:
|
|
if width > height:
|
|
new_width = max_size
|
|
new_height = int(height * (max_size / width))
|
|
else:
|
|
new_height = max_size
|
|
new_width = int(width * (max_size / height))
|
|
pil_image = pil_image.resize((new_width, new_height), Image.BILINEAR)
|
|
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format="JPEG", quality=quality, progressive=True)
|
|
|
|
img_bytes = buffer.getvalue()
|
|
base64_str = base64.b64encode(img_bytes).decode('utf-8')
|
|
|
|
return base64_str, "image/jpeg"
|
|
|
|
def base64_to_tensor(self, base64_str):
|
|
"""
|
|
Convert base64 string to ComfyUI image tensor
|
|
"""
|
|
img_bytes = base64.b64decode(base64_str)
|
|
pil_image = Image.open(io.BytesIO(img_bytes))
|
|
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
|
|
numpy_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
tensor = torch.from_numpy(numpy_image)[None,]
|
|
|
|
return tensor
|
|
|
|
def _tensor_to_png_bytes(self, tensor, max_size=2048):
|
|
"""
|
|
Convert ComfyUI image tensor to PNG bytes for multipart upload
|
|
"""
|
|
if len(tensor.shape) == 4:
|
|
tensor = tensor[0]
|
|
|
|
numpy_image = (tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(numpy_image)
|
|
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
|
|
width, height = pil_image.size
|
|
if max(width, height) > max_size:
|
|
if width > height:
|
|
new_width = max_size
|
|
new_height = int(height * (max_size / width))
|
|
else:
|
|
new_height = max_size
|
|
new_width = int(width * (max_size / height))
|
|
pil_image = pil_image.resize((new_width, new_height), Image.BILINEAR)
|
|
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format="PNG")
|
|
|
|
return buffer.getvalue()
|
|
|
|
|
|
class VectorEngineGemini(VectorEngineBase):
|
|
"""
|
|
ComfyUI node for Gemini Image Generation via Vector Engine API
|
|
Supports multi-image input and system prompts
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": (["gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview"], {
|
|
"default": "gemini-3-pro-image-preview"
|
|
}),
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "Generate a creative image based on the provided pictures."
|
|
}),
|
|
"system_prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "You are an AI assistant skilled in generating images and editing pictures."
|
|
}),
|
|
"aspect_ratio": (["Auto", "1:1", "2:3", "3:2", "4:3", "3:4", "16:9", "9:16"], {
|
|
"default": "Auto"
|
|
}),
|
|
"image_size": (["1K", "2K", "4K"], {
|
|
"default": "1K"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
"image_3": ("IMAGE",),
|
|
"image_4": ("IMAGE",),
|
|
"image_5": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "info")
|
|
FUNCTION = "generate_image"
|
|
CATEGORY = "VectorEngine"
|
|
|
|
def generate_image(self, model, prompt, system_prompt, aspect_ratio, image_size, seed,
|
|
image_1=None, image_2=None, image_3=None, image_4=None, image_5=None):
|
|
"""
|
|
Generate image using Gemini API via Vector Engine
|
|
"""
|
|
|
|
try:
|
|
conn = http.client.HTTPSConnection("api.vectorengine.ai")
|
|
|
|
# Build parts array
|
|
parts = []
|
|
|
|
# Add system prompt and user prompt as text
|
|
full_prompt = f"{system_prompt}\n\nUser request: {prompt}"
|
|
parts.append({"text": full_prompt})
|
|
|
|
# Add images if provided
|
|
images = [image_1, image_2, image_3, image_4, image_5]
|
|
image_count = 0
|
|
total_encode_time = 0.0
|
|
|
|
print(f"[VectorEngine Gemini] Processing {sum(1 for img in images if img is not None)} input images...")
|
|
|
|
for idx, img in enumerate(images, 1):
|
|
if img is not None:
|
|
image_count += 1
|
|
encode_start = time.time()
|
|
base64_data, mime_type = self.tensor_to_base64(img, max_size=2048, quality=85)
|
|
encode_time = time.time() - encode_start
|
|
total_encode_time += encode_time
|
|
|
|
data_size_kb = len(base64_data) / 1024
|
|
print(f"[VectorEngine Gemini] Image {idx}: encoded in {encode_time:.2f}s, size: {data_size_kb:.1f}KB")
|
|
|
|
parts.append({
|
|
"inline_data": {
|
|
"mime_type": mime_type,
|
|
"data": base64_data
|
|
}
|
|
})
|
|
|
|
# Build request payload
|
|
payload = json.dumps({
|
|
"contents": [{"role": "user", "parts": parts}],
|
|
"generationConfig": {
|
|
"responseModalities": ["TEXT", "IMAGE"],
|
|
"imageConfig": {
|
|
"aspectRatio": "" if aspect_ratio == "Auto" else aspect_ratio,
|
|
"imageSize": image_size
|
|
}
|
|
}
|
|
})
|
|
|
|
headers = {
|
|
'Authorization': f'Bearer {self.api_key}',
|
|
'Content-Type': 'application/json'
|
|
}
|
|
|
|
payload_size_mb = len(payload) / (1024 * 1024)
|
|
print(f"[VectorEngine Gemini] Sending request (payload: {payload_size_mb:.2f}MB)...")
|
|
|
|
start_time = time.time()
|
|
|
|
conn.request("POST",
|
|
f"/v1beta/models/{model}:generateContent?key={self.api_key}",
|
|
payload, headers)
|
|
|
|
res = conn.getresponse()
|
|
data = res.read()
|
|
api_generation_time = time.time() - start_time
|
|
|
|
print(f"[VectorEngine Gemini] API request completed in {api_generation_time:.2f}s")
|
|
|
|
response_json = json.loads(data.decode("utf-8"))
|
|
candidates = response_json.get("candidates", [])
|
|
|
|
if not candidates:
|
|
error_msg = response_json.get("error", {}).get("message", "Unknown error")
|
|
info_text = self._format_info(model, aspect_ratio, image_size, api_generation_time,
|
|
False, None, error_msg, image_count, seed,
|
|
total_encode_time, 0.0)
|
|
black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
return (black_image, info_text)
|
|
|
|
# Extract image from response
|
|
for candidate in candidates:
|
|
content = candidate.get("content", {})
|
|
parts = content.get("parts", [])
|
|
|
|
for part in parts:
|
|
img_base64 = None
|
|
|
|
if "inline_data" in part:
|
|
img_base64 = part["inline_data"]["data"]
|
|
elif "inlineData" in part:
|
|
img_base64 = part["inlineData"]["data"]
|
|
elif "data" in part and "text" not in part:
|
|
img_base64 = part["data"]
|
|
|
|
if img_base64:
|
|
decode_start = time.time()
|
|
output_tensor = self.base64_to_tensor(img_base64)
|
|
decode_time = time.time() - decode_start
|
|
|
|
print(f"[VectorEngine Gemini] Image decoded in {decode_time:.2f}s")
|
|
|
|
_, height, width, _ = output_tensor.shape
|
|
|
|
info_text = self._format_info(model, aspect_ratio, image_size, api_generation_time,
|
|
True, f"{width}x{height}", None, image_count, seed,
|
|
total_encode_time, decode_time)
|
|
|
|
return (output_tensor, info_text)
|
|
|
|
# No image found
|
|
info_text = self._format_info(model, aspect_ratio, image_size, api_generation_time,
|
|
False, None, "No image found in API response", image_count, seed,
|
|
total_encode_time, 0.0)
|
|
black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
return (black_image, info_text)
|
|
|
|
except Exception as e:
|
|
try:
|
|
error_time = time.time() - start_time
|
|
except:
|
|
error_time = 0.0
|
|
try:
|
|
error_encode_time = total_encode_time
|
|
except:
|
|
error_encode_time = 0.0
|
|
|
|
info_text = self._format_info(model, aspect_ratio, image_size, error_time,
|
|
False, None, str(e),
|
|
sum(1 for img in [image_1, image_2, image_3, image_4, image_5] if img is not None),
|
|
seed, error_encode_time, 0.0)
|
|
black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
return (black_image, info_text)
|
|
|
|
def _format_info(self, model_name, aspect_ratio, image_size, generation_time,
|
|
success, resolution, error_message, input_images, seed,
|
|
encode_time, decode_time):
|
|
lines = [
|
|
"=" * 60,
|
|
"Vector Engine Gemini - Generation Result",
|
|
"=" * 60,
|
|
f"Model: {model_name}",
|
|
f"Aspect Ratio: {aspect_ratio}",
|
|
f"Image Size: {image_size}",
|
|
]
|
|
|
|
if resolution:
|
|
lines.append(f"Resolution: {resolution}")
|
|
|
|
lines.extend([
|
|
f"Seed: {seed}",
|
|
f"Input Images: {input_images}",
|
|
"",
|
|
"--- Timing ---",
|
|
f"Encode: {encode_time:.3f}s | API: {generation_time:.2f}s | Decode: {decode_time:.3f}s",
|
|
f"Total: {(encode_time + generation_time + decode_time):.2f}s",
|
|
"",
|
|
f"Status: {'SUCCESS' if success else 'FAILED'}",
|
|
])
|
|
|
|
if not success and error_message:
|
|
lines.extend(["", "Error:", error_message])
|
|
|
|
lines.append("=" * 60)
|
|
return "\n".join(lines)
|
|
|
|
|
|
class VectorEngineGPT(VectorEngineBase):
|
|
"""
|
|
ComfyUI node for GPT Image Generation via Vector Engine API
|
|
Supports image generation and editing
|
|
"""
|
|
|
|
def __init__(self):
|
|
# GPT node uses a different API key
|
|
self.api_key = os.getenv("VECTOR_ENGINE_API_KEY_GPT")
|
|
if not self.api_key:
|
|
raise ValueError(
|
|
"VECTOR_ENGINE_API_KEY_GPT environment variable is not set. "
|
|
"Please set it before using this node: "
|
|
"export VECTOR_ENGINE_API_KEY_GPT='your-api-key-here'"
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": (["gpt-image-2", "gpt-image-1.5", "gpt-image-1"], {
|
|
"default": "gpt-image-2"
|
|
}),
|
|
"prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "A beautiful landscape with mountains and a lake."
|
|
}),
|
|
"size": (["auto", "1024x1024", "1536x1024", "1024x1536"], {
|
|
"default": "auto"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
"image_3": ("IMAGE",),
|
|
"image_4": ("IMAGE",),
|
|
"image_5": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "info")
|
|
FUNCTION = "generate_image"
|
|
CATEGORY = "VectorEngine"
|
|
|
|
def generate_image(self, model, prompt, size, seed,
|
|
image_1=None, image_2=None, image_3=None, image_4=None, image_5=None):
|
|
"""
|
|
Generate or edit image using GPT API via Vector Engine
|
|
"""
|
|
|
|
# Collect input images
|
|
images = [image_1, image_2, image_3, image_4, image_5]
|
|
input_images = [img for img in images if img is not None]
|
|
|
|
try:
|
|
conn = http.client.HTTPSConnection("api.vectorengine.ai")
|
|
|
|
image_count = len(input_images)
|
|
total_encode_time = 0.0
|
|
|
|
# Determine endpoint based on whether images are provided
|
|
if image_count > 0:
|
|
# Use /v1/images/edits with multipart/form-data
|
|
endpoint = "/v1/images/edits"
|
|
print(f"[VectorEngine GPT] Edit mode - Size: {size}, Images: {image_count}")
|
|
|
|
boundary = 'wL36Yn8afVp8Ag7AmP8qZ0SA4n1v9T'
|
|
dataList = []
|
|
|
|
# Add images
|
|
for idx, img in enumerate(input_images):
|
|
encode_start = time.time()
|
|
img_bytes = self._tensor_to_png_bytes(img)
|
|
encode_time = time.time() - encode_start
|
|
total_encode_time += encode_time
|
|
|
|
data_size_kb = len(img_bytes) / 1024
|
|
print(f"[VectorEngine GPT] Image {idx + 1}: encoded in {encode_time:.2f}s, size: {data_size_kb:.1f}KB")
|
|
|
|
dataList.append(f'--{boundary}'.encode())
|
|
dataList.append(f'Content-Disposition: form-data; name=image; filename=image_{idx + 1}.png'.encode())
|
|
dataList.append(b'Content-Type: image/png')
|
|
dataList.append(b'')
|
|
dataList.append(img_bytes)
|
|
|
|
# Add form fields
|
|
for name, value in [("prompt", prompt), ("model", model), ("n", "1"), ("size", size)]:
|
|
dataList.append(f'--{boundary}'.encode())
|
|
dataList.append(f'Content-Disposition: form-data; name={name};'.encode())
|
|
dataList.append(b'Content-Type: text/plain')
|
|
dataList.append(b'')
|
|
dataList.append(value.encode() if isinstance(value, str) else value)
|
|
|
|
dataList.append(f'--{boundary}--'.encode())
|
|
dataList.append(b'')
|
|
|
|
payload = b'\r\n'.join(dataList)
|
|
|
|
headers = {
|
|
'Accept': 'application/json',
|
|
'Authorization': f'Bearer {self.api_key}',
|
|
'Content-Type': f'multipart/form-data; boundary={boundary}'
|
|
}
|
|
else:
|
|
# Use /v1/images/generations with JSON
|
|
endpoint = "/v1/images/generations"
|
|
print(f"[VectorEngine GPT] Generation mode - Size: {size}")
|
|
|
|
payload = json.dumps({
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"size": size,
|
|
"n": 1
|
|
})
|
|
|
|
headers = {
|
|
'Accept': 'application/json',
|
|
'Authorization': f'Bearer {self.api_key}',
|
|
'Content-Type': 'application/json'
|
|
}
|
|
|
|
payload_size_kb = len(payload) / 1024
|
|
print(f"[VectorEngine GPT] Sending request (payload: {payload_size_kb:.2f}KB)...")
|
|
|
|
start_time = time.time()
|
|
conn.request("POST", endpoint, payload, headers)
|
|
|
|
res = conn.getresponse()
|
|
data = res.read()
|
|
api_generation_time = time.time() - start_time
|
|
|
|
print(f"[VectorEngine GPT] API request completed in {api_generation_time:.2f}s")
|
|
|
|
response_json = json.loads(data.decode("utf-8"))
|
|
|
|
# Check for error
|
|
if "error" in response_json:
|
|
error_msg = response_json["error"].get("message", str(response_json["error"]))
|
|
info_text = self._format_info(model, size, api_generation_time,
|
|
False, None, error_msg, image_count, seed,
|
|
total_encode_time, 0.0)
|
|
black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
return (black_image, info_text)
|
|
|
|
# Extract image from response
|
|
data_list = response_json.get("data", [])
|
|
|
|
if not data_list:
|
|
info_text = self._format_info(model, size, api_generation_time,
|
|
False, None, "No image data in response", image_count, seed,
|
|
total_encode_time, 0.0)
|
|
black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
return (black_image, info_text)
|
|
|
|
image_data = data_list[0]
|
|
decode_start = time.time()
|
|
|
|
# Handle URL response
|
|
if "url" in image_data:
|
|
image_url = image_data["url"]
|
|
print(f"[VectorEngine GPT] Downloading image from URL...")
|
|
|
|
req = urllib.request.Request(image_url)
|
|
with urllib.request.urlopen(req) as response:
|
|
img_bytes = response.read()
|
|
|
|
pil_image = Image.open(io.BytesIO(img_bytes))
|
|
|
|
# Handle base64 response
|
|
elif "b64_json" in image_data:
|
|
img_base64 = image_data["b64_json"]
|
|
img_bytes = base64.b64decode(img_base64)
|
|
pil_image = Image.open(io.BytesIO(img_bytes))
|
|
else:
|
|
info_text = self._format_info(model, size, api_generation_time,
|
|
False, None, "Unknown response format", image_count, seed,
|
|
total_encode_time, 0.0)
|
|
black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
return (black_image, info_text)
|
|
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
|
|
numpy_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
output_tensor = torch.from_numpy(numpy_image)[None,]
|
|
|
|
decode_time = time.time() - decode_start
|
|
print(f"[VectorEngine GPT] Image downloaded and decoded in {decode_time:.2f}s")
|
|
|
|
_, height, width, _ = output_tensor.shape
|
|
|
|
info_text = self._format_info(model, size, api_generation_time,
|
|
True, f"{width}x{height}", None, image_count, seed,
|
|
total_encode_time, decode_time)
|
|
|
|
return (output_tensor, info_text)
|
|
|
|
except Exception as e:
|
|
try:
|
|
error_time = time.time() - start_time
|
|
except:
|
|
error_time = 0.0
|
|
try:
|
|
error_encode_time = total_encode_time
|
|
error_image_count = image_count
|
|
except:
|
|
error_encode_time = 0.0
|
|
error_image_count = 0
|
|
|
|
info_text = self._format_info(model, size, error_time,
|
|
False, None, str(e), error_image_count, seed,
|
|
error_encode_time, 0.0)
|
|
black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
return (black_image, info_text)
|
|
|
|
def _format_info(self, model_name, size, generation_time,
|
|
success, resolution, error_message, input_images, seed,
|
|
encode_time, decode_time):
|
|
lines = [
|
|
"=" * 60,
|
|
"Vector Engine GPT - Generation Result",
|
|
"=" * 60,
|
|
f"Model: {model_name}",
|
|
f"Size: {size}",
|
|
]
|
|
|
|
if resolution:
|
|
lines.append(f"Resolution: {resolution}")
|
|
|
|
lines.extend([
|
|
f"Seed: {seed}",
|
|
f"Input Images: {input_images}",
|
|
"",
|
|
"--- Timing ---",
|
|
f"Encode: {encode_time:.3f}s | API: {generation_time:.2f}s | Decode: {decode_time:.3f}s",
|
|
f"Total: {(encode_time + generation_time + decode_time):.2f}s",
|
|
"",
|
|
f"Status: {'SUCCESS' if success else 'FAILED'}",
|
|
])
|
|
|
|
if not success and error_message:
|
|
lines.extend(["", "Error:", error_message])
|
|
|
|
lines.append("=" * 60)
|
|
return "\n".join(lines)
|
|
|
|
|
|
# ComfyUI node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"VectorEngineGemini": VectorEngineGemini,
|
|
"VectorEngineGPT": VectorEngineGPT,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"VectorEngineGemini": "Vector Engine Gemini",
|
|
"VectorEngineGPT": "Vector Engine GPT",
|
|
}
|