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 VectorEngineImageGenerator: """ ComfyUI node for Vector Engine Image Generation API """ # Model configuration: defines API endpoints and parameter formats for each model MODEL_CONFIG = { "gemini-3-pro-image-preview": { "api_type": "gemini", "endpoint": "/v1beta/models/{model}:generateContent", "supports_multi_image": True, "supports_system_prompt": True, }, "gpt-image-1.5": { "api_type": "openai", "endpoint": "/v1/images/generations", "supports_multi_image": False, "supports_system_prompt": False, }, } 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'" ) @classmethod def INPUT_TYPES(cls): return { "required": { "model": (["gemini-3-pro-image-preview", "gpt-image-1.5"], { "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. (Only for Gemini models)" }), # Gemini model size options "aspect_ratio": (["1:1", "2:3", "3:2", "4:3", "3:4", "16:9", "9:16"], { "default": "1:1" }), "image_size": (["1K", "2K", "4K"], { "default": "1K" }), # GPT model size options (only used when model is gpt-image-1.5) "gpt_size": (["1024x1024", "1536x1024", "1024x1536", "auto"], { "default": "1024x1024" }), "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 tensor_to_base64(self, tensor, mime_type="image/jpeg", max_size=2048, quality=85): """ Convert ComfyUI image tensor to base64 string with optimization tensor shape: [batch, height, width, channels] with values in [0, 1] Args: tensor: Image tensor mime_type: Output mime type (default JPEG for better compression) max_size: Maximum dimension (width or height) in pixels quality: JPEG quality (1-100, default 85 for good balance) """ # Take first image if batch if len(tensor.shape) == 4: tensor = tensor[0] # Convert to numpy and scale to 0-255 numpy_image = (tensor.cpu().numpy() * 255).astype(np.uint8) # Convert to PIL Image pil_image = Image.fromarray(numpy_image) # Convert to RGB first (more efficient to do before resize) if pil_image.mode != 'RGB': pil_image = pil_image.convert('RGB') # Resize if image is too large 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)) # Use BILINEAR for faster resizing (3-4x faster than LANCZOS) # Quality difference is minimal for photo content pil_image = pil_image.resize((new_width, new_height), Image.BILINEAR) # Save with compression buffer = io.BytesIO() # Remove optimize=True for faster encoding (file size increase is minimal ~2-5%) # Use progressive for better streaming (no speed impact) pil_image.save(buffer, format="JPEG", quality=quality, progressive=True) # Encode to base64 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 """ # Decode base64 img_bytes = base64.b64decode(base64_str) # Convert to PIL Image pil_image = Image.open(io.BytesIO(img_bytes)) # Convert to RGB if necessary if pil_image.mode != 'RGB': pil_image = pil_image.convert('RGB') # Convert to numpy array numpy_image = np.array(pil_image).astype(np.float32) / 255.0 # Convert to tensor [1, height, width, channels] 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 """ # Take first image if batch if len(tensor.shape) == 4: tensor = tensor[0] # Convert to numpy and scale to 0-255 numpy_image = (tensor.cpu().numpy() * 255).astype(np.uint8) # Convert to PIL Image pil_image = Image.fromarray(numpy_image) # Convert to RGB if pil_image.mode != 'RGB': pil_image = pil_image.convert('RGB') # Resize if image is too large 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) # Save to bytes buffer as PNG buffer = io.BytesIO() pil_image.save(buffer, format="PNG") return buffer.getvalue() def generate_image(self, model, prompt, system_prompt, aspect_ratio, image_size, gpt_size, seed, image_1=None, image_2=None, image_3=None, image_4=None, image_5=None): """ Main function to generate image using Vector Engine API Routes to appropriate API based on model type """ # Get model configuration model_config = self.MODEL_CONFIG.get(model, self.MODEL_CONFIG["gemini-3-pro-image-preview"]) api_type = model_config["api_type"] print(f"[VectorEngine] Using {api_type} API for model: {model}") if api_type == "openai": # Collect images for OpenAI API images = [image_1, image_2, image_3, image_4, image_5] input_images = [img for img in images if img is not None] return self._generate_image_openai(model, prompt, gpt_size, seed, input_images) else: return self._generate_image_gemini(model, prompt, system_prompt, aspect_ratio, image_size, seed, image_1, image_2, image_3, image_4, image_5) def _generate_image_openai(self, model, prompt, gpt_size, seed, input_images=None): """ Generate image using OpenAI-style API (for gpt-image-1.5 and similar models) Supports both generation (no images) and editing (with images) """ if input_images is None: input_images = [] try: # Prepare connection conn = http.client.HTTPSConnection("api.vectorengine.ai") # Parse gpt_size to get actual size value # Options: "1024x1024 (square)", "1536x1024 (landscape)", "1024x1536 (portrait)", "auto" size_mapping = { "1024x1024 (square)": "1024x1024", "1536x1024 (landscape)": "1536x1024", "1024x1536 (portrait)": "1024x1536", "auto": "auto" } size = size_mapping.get(gpt_size, "1024x1024") image_count = len(input_images) total_encode_time = 0.0 # Determine endpoint and format 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] OpenAI Edit API - Size: {size}, Images: {image_count}") # Build multipart form data boundary = 'wL36Yn8afVp8Ag7AmP8qZ0SA4n1v9T' dataList = [] # Add images for idx, img in enumerate(input_images): encode_start = time.time() # Convert tensor to PNG bytes 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] 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 prompt dataList.append(f'--{boundary}'.encode()) dataList.append(b'Content-Disposition: form-data; name=prompt;') dataList.append(b'Content-Type: text/plain') dataList.append(b'') dataList.append(prompt.encode()) # Add model dataList.append(f'--{boundary}'.encode()) dataList.append(b'Content-Disposition: form-data; name=model;') dataList.append(b'Content-Type: text/plain') dataList.append(b'') dataList.append(model.encode()) # Add n dataList.append(f'--{boundary}'.encode()) dataList.append(b'Content-Disposition: form-data; name=n;') dataList.append(b'Content-Type: text/plain') dataList.append(b'') dataList.append(b'1') # Add size dataList.append(f'--{boundary}'.encode()) dataList.append(b'Content-Disposition: form-data; name=size;') dataList.append(b'Content-Type: text/plain') dataList.append(b'') dataList.append(size.encode()) # End boundary 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] OpenAI Generation API - 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' } # Calculate payload size payload_size_kb = len(payload) / 1024 print(f"[VectorEngine] Sending request to API (payload size: {payload_size_kb:.2f}KB)...") # Record start time start_time = time.time() # Make API request conn.request("POST", endpoint, payload, headers) res = conn.getresponse() data = res.read() # Record end time api_generation_time = time.time() - start_time print(f"[VectorEngine] API request completed in {api_generation_time:.2f}s") # Parse response 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_name=model, aspect_ratio="-", image_size=size, generation_time=api_generation_time, success=False, error_message=error_msg, input_images=image_count, seed=seed, encode_time=total_encode_time, decode_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_name=model, aspect_ratio="-", image_size=size, generation_time=api_generation_time, success=False, error_message="No image data in response", input_images=image_count, seed=seed, encode_time=total_encode_time, decode_time=0.0 ) black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) return (black_image, info_text) # Get first image 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] Downloading image from URL...") # Download image from URL req = urllib.request.Request(image_url) with urllib.request.urlopen(req) as response: img_bytes = response.read() # Convert to PIL Image 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_name=model, aspect_ratio="-", image_size=size, generation_time=api_generation_time, success=False, error_message="Unknown response format", input_images=image_count, seed=seed, encode_time=total_encode_time, decode_time=0.0 ) black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) return (black_image, info_text) # Convert to RGB if necessary if pil_image.mode != 'RGB': pil_image = pil_image.convert('RGB') # Convert to tensor 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] Image downloaded and decoded in {decode_time:.2f}s") # Get actual dimensions _, height, width, _ = output_tensor.shape # Format info text info_text = self._format_info( model_name=model, aspect_ratio="-", image_size=size, resolution=f"{width}x{height}", generation_time=api_generation_time, success=True, input_images=image_count, seed=seed, encode_time=total_encode_time, decode_time=decode_time ) return (output_tensor, info_text) except Exception as e: try: error_generation_time = time.time() - start_time except: error_generation_time = 0.0 try: error_encode_time = total_encode_time error_image_count = image_count error_size = size except: error_encode_time = 0.0 error_image_count = 0 error_size = "1024x1024" info_text = self._format_info( model_name=model, aspect_ratio="-", image_size=error_size, generation_time=error_generation_time, success=False, error_message=str(e), input_images=error_image_count, seed=seed, encode_time=error_encode_time, decode_time=0.0 ) black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) return (black_image, info_text) def _generate_image_gemini(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-style API (for gemini models) """ try: # Prepare connection 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] 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() # Use optimized compression: max 2048px, 85% quality JPEG 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 # Calculate compressed size data_size_kb = len(base64_data) / 1024 print(f"[VectorEngine] 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": aspect_ratio, "imageSize": image_size } } }) headers = { 'Authorization': f'Bearer {self.api_key}', 'Content-Type': 'application/json' } # Calculate payload size payload_size_mb = len(payload) / (1024 * 1024) print(f"[VectorEngine] Sending request to API (payload size: {payload_size_mb:.2f}MB)...") # Record start time right before API request start_time = time.time() # Make API request conn.request("POST", f"/v1beta/models/{model}:generateContent?key={self.api_key}", payload, headers) res = conn.getresponse() data = res.read() # Record end time right after receiving response api_generation_time = time.time() - start_time print(f"[VectorEngine] API request completed in {api_generation_time:.2f}s") # Parse response response_json = json.loads(data.decode("utf-8")) # Extract image from response candidates = response_json.get("candidates", []) if not candidates: # No candidates, check for error error_msg = response_json.get("error", {}).get("message", "Unknown error") info_text = self._format_info( model_name=model, aspect_ratio=aspect_ratio, image_size=image_size, generation_time=api_generation_time, success=False, error_message=error_msg, input_images=image_count, seed=seed, encode_time=total_encode_time, decode_time=0.0 ) # Return a black image as placeholder black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) return (black_image, info_text) # Try to extract image from first candidate for candidate in candidates: content = candidate.get("content", {}) parts = content.get("parts", []) for part in parts: img_base64 = None mime_type = None # Check different possible field names if "inline_data" in part: img_base64 = part["inline_data"]["data"] mime_type = part["inline_data"].get("mime_type", "image/jpeg") elif "inlineData" in part: img_base64 = part["inlineData"]["data"] mime_type = part["inlineData"].get("mimeType", "image/jpeg") elif "data" in part and not "text" in part: img_base64 = part["data"] mime_type = part.get("mimeType", part.get("mime_type", "image/jpeg")) if img_base64: # Convert to tensor and record decode time decode_start = time.time() output_tensor = self.base64_to_tensor(img_base64) decode_time = time.time() - decode_start print(f"[VectorEngine] Image decoded in {decode_time:.2f}s") # Get actual dimensions _, height, width, _ = output_tensor.shape # Format info text info_text = self._format_info( model_name=model, aspect_ratio=aspect_ratio, image_size=image_size, resolution=f"{width}x{height}", generation_time=api_generation_time, success=True, input_images=image_count, seed=seed, encode_time=total_encode_time, decode_time=decode_time ) return (output_tensor, info_text) # If we reach here, no image was found info_text = self._format_info( model_name=model, aspect_ratio=aspect_ratio, image_size=image_size, generation_time=api_generation_time, success=False, error_message="No image found in API response", input_images=image_count, seed=seed, encode_time=total_encode_time, decode_time=0.0 ) black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) return (black_image, info_text) except Exception as e: # If error occurred before API request, set generation time to 0 try: error_generation_time = time.time() - start_time except: error_generation_time = 0.0 # Get encode time if available try: error_encode_time = total_encode_time except: error_encode_time = 0.0 info_text = self._format_info( model_name=model, aspect_ratio=aspect_ratio, image_size=image_size, generation_time=error_generation_time, success=False, error_message=str(e), input_images=sum(1 for img in [image_1, image_2, image_3, image_4, image_5] if img is not None), seed=seed, encode_time=error_encode_time, decode_time=0.0 ) # Return black image as placeholder 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=None, error_message=None, input_images=0, seed=0, encode_time=0.0, decode_time=0.0): """ Format generation information as text """ lines = [ "=" * 60, "Vector Engine Image Generation Result", "=" * 60, f"Model Name: {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 Information ---", f"Image Encode Time: {encode_time:.3f}s", f"API Generation Time: {generation_time:.2f}s", f"Image Decode Time: {decode_time:.3f}s", f"Total Time: {(encode_time + generation_time + decode_time):.2f}s", "", f"Status: {'SUCCESS' if success else 'FAILED'}", ]) if not success and error_message: lines.extend([ "", "Error Message:", error_message ]) lines.append("=" * 60) return "\n".join(lines) # ComfyUI node registration NODE_CLASS_MAPPINGS = { "VectorEngineImageGenerator": VectorEngineImageGenerator } NODE_DISPLAY_NAME_MAPPINGS = { "VectorEngineImageGenerator": "Vector Engine Image Generator" }