commit 23e8af58fb0e44e0cb68a0a830a119f668e27482 Author: dzy1128 <1241877130@qq.com> Date: Mon Dec 1 15:50:54 2025 +0800 feat:Init commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f395abc --- /dev/null +++ b/.gitignore @@ -0,0 +1,48 @@ +# Configuration files with API keys +config.py + +# Python cache +__pycache__/ +*.py[cod] +*$py.class +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# Testing +.pytest_cache/ +.coverage +htmlcov/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db + +# Output files +assets/output_*.png +assets/output_*.jpg +assets/output_*.jpeg + diff --git a/README.md b/README.md new file mode 100644 index 0000000..c5b1dd8 --- /dev/null +++ b/README.md @@ -0,0 +1,156 @@ +# ComfyUI Vector Engine Image Generator Node + +This is a custom node for ComfyUI that integrates Vector Engine's Gemini API for advanced image generation and editing. + +## Features + +- 🖼️ Support up to 5 optional input images +- 🎨 Configurable aspect ratios (1:1, 2:3, 3:2, 4:3, 3:4, 16:9, 9:16) +- 📐 Multiple resolution options (1K, 2K, 4K) +- 🎲 Seed control for reproducible results +- 📝 Custom system and user prompts +- ℹ️ Detailed generation information output + +## Installation + +1. Navigate to your ComfyUI custom_nodes directory: +```bash +cd ComfyUI/custom_nodes/ +``` + +2. Clone or copy this repository: +```bash +git clone ComfyUI-Vector-Engine +# or copy the folder manually +``` + +3. Restart ComfyUI + +## Node Inputs + +### Required Parameters + +- **prompt** (STRING, multiline): User prompt describing what to generate + - Default: "Generate a creative image based on the provided pictures." + +- **system_prompt** (STRING, multiline): System instructions for the AI + - Default: "You are an AI assistant skilled in generating images and editing pictures." + +- **aspect_ratio** (DROPDOWN): Image aspect ratio + - Options: 1:1, 2:3, 3:2, 4:3, 3:4, 16:9, 9:16 + - Default: 1:1 + +- **image_size** (DROPDOWN): Output resolution + - Options: 1K, 2K, 4K + - Default: 1K + +- **seed** (INT): Random seed for reproducibility + - Range: 0 to max int + - Default: 0 + +### Optional Parameters + +- **image_1** to **image_5** (IMAGE): Up to 5 optional input images + - All images are optional, you can use 0 to 5 images + +## Node Outputs + +1. **image** (IMAGE): The generated image tensor + - Can be connected to Preview Image, Save Image, or other image processing nodes + +2. **info** (STRING): Generation information including: + - Model name + - Aspect ratio + - Image size + - Actual resolution + - Seed value + - Number of input images + - Generation time + - Success/failure status + - Error message (if failed) + +## Usage Examples + +### Example 1: Generate from Text Only +``` +1. Add "Vector Engine Image Generator" node +2. Set prompt: "A beautiful sunset over mountains" +3. Leave all image inputs empty +4. Connect output to "Preview Image" +5. Run the workflow +``` + +### Example 2: Edit/Combine Multiple Images +``` +1. Load 2-5 images using "Load Image" nodes +2. Add "Vector Engine Image Generator" node +3. Connect images to image_1, image_2, etc. +4. Set prompt: "Combine these people in a group photo at the beach" +5. Connect output to "Save Image" +6. Run the workflow +``` + +### Example 3: Generate with Specific Aspect Ratio +``` +1. Add "Vector Engine Image Generator" node +2. Set aspect_ratio: "16:9" +3. Set image_size: "2K" +4. Set prompt: "A cinematic landscape photo" +5. Run the workflow +``` + +## Configuration + +The API key is currently hardcoded in the node. To change it: + +1. Open `vector_engine_node.py` +2. Find the line: `self.api_key = "sk-..."` +3. Replace with your API key + +For production use, consider using environment variables: +```python +import os +self.api_key = os.getenv("VECTOR_ENGINE_API_KEY", "default-key") +``` + +## Model Information + +- **Current Model**: gemini-3-pro-image-preview +- **API Endpoint**: api.vectorengine.ai +- **Features**: Text-to-image and image-to-image generation + +## Troubleshooting + +### No image generated (black output) +- Check the info output for error messages +- Verify your API key is valid +- Check network connectivity +- Ensure input images are in correct format + +### Image quality issues +- Try increasing image_size (1K → 2K → 4K) +- Adjust aspect_ratio to match desired output +- Refine your prompt for better results + +### Slow generation +- Higher resolutions (4K) take longer +- Multiple input images increase processing time +- Check your network connection + +## Dependencies + +- torch +- numpy +- PIL (Pillow) +- Standard Python libraries (http.client, json, base64, time, io) + +These should already be available in a standard ComfyUI installation. + +## License + +[Add your license information here] + +## Support + +For issues and questions, please open an issue on the repository. + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ffa6e23 --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +""" +Vector Engine Image Generator Node for ComfyUI + +This custom node integrates Vector Engine's Gemini API for image generation +into ComfyUI workflow. +""" + +from .vector_engine_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + diff --git a/assets/boy.jpg b/assets/boy.jpg new file mode 100644 index 0000000..3abe9fe Binary files /dev/null and b/assets/boy.jpg differ diff --git a/assets/girl.png b/assets/girl.png new file mode 100644 index 0000000..56f2377 Binary files /dev/null and b/assets/girl.png differ diff --git a/config.example.py b/config.example.py new file mode 100644 index 0000000..1d775c6 --- /dev/null +++ b/config.example.py @@ -0,0 +1,23 @@ +""" +Configuration file for Vector Engine API + +To use: +1. Copy this file to config.py +2. Update with your actual API key +3. Import in vector_engine_node.py +""" + +# Vector Engine API Configuration +VECTOR_ENGINE_API_KEY = "sk-your-api-key-here" + +# API Endpoint +VECTOR_ENGINE_API_HOST = "api.vectorengine.ai" + +# Default Model +DEFAULT_MODEL = "gemini-3-pro-image-preview" + +# Default Generation Config +DEFAULT_SYSTEM_PROMPT = "You are an AI assistant skilled in generating images and editing pictures." +DEFAULT_ASPECT_RATIO = "1:1" +DEFAULT_IMAGE_SIZE = "1K" + diff --git a/test.py b/test.py new file mode 100644 index 0000000..4538981 --- /dev/null +++ b/test.py @@ -0,0 +1,138 @@ +import http.client +import json +import base64 +import time + +conn = http.client.HTTPSConnection("api.vectorengine.ai") + + +with open("./assets/girl.png", "rb") as f: + girl_data = base64.b64encode(f.read()).decode('utf-8') + +with open("./assets/boy.jpg", "rb") as f: + boy_data = base64.b64encode(f.read()).decode('utf-8') + +payload = json.dumps({ + "contents": [ + { + "role": "user", + "parts": [ + { + "text": "'Picture 1 shows a girl, and Picture 2 shows a boy. Generate a photo of them hugging." + }, + { + "inline_data": { + "mime_type": "image/png", + "data": girl_data + } + }, + { + "inline_data": { + "mime_type": "image/jpeg", + "data": boy_data + } + } + ] + } + ], + "generationConfig": { + "responseModalities": [ + "TEXT", + "IMAGE" + ], + "imageConfig": { + "aspectRatio": "2:3", + "imageSize": "1K" + } + } +}) +headers = { + 'Authorization': 'Bearer sk-RZBLe4v8MolmD3fwOd6vdRTCaj7PohDhf4f44UNgWHEAA4zF', + 'Content-Type': 'application/json' +} + +# 记录开始时间 +print("开始生成图片...") +start_time = time.time() + +conn.request("POST", "/v1beta/models/gemini-3-pro-image-preview:generateContent?key=sk-RZBLe4v8MolmD3fwOd6vdRTCaj7PohDhf4f44UNgWHEAA4zF", payload, headers) +res = conn.getresponse() +data = res.read() + +# 记录结束时间并计算耗时 +end_time = time.time() +generation_time = end_time - start_time + +print(f"✓ 图片生成完成！") +print(f"⏱️ 总耗时: {generation_time:.2f} 秒 ({generation_time:.3f}s)") + +# 解析响应 +response_json = json.loads(data.decode("utf-8")) +#print("响应内容：", json.dumps(response_json, indent=2, ensure_ascii=False)) + +# 提取并保存图片 +saved_images = 0 +try: + # 根据 API 响应结构提取图片数据（可能需要根据实际响应调整路径） + candidates = response_json.get("candidates", []) + print(f"\n候选数量: {len(candidates)}") + + for i, candidate in enumerate(candidates): + print(f"\n处理候选 {i}:") + print(f"候选结构: {list(candidate.keys())}") + + content = candidate.get("content", {}) + print(f"内容结构: {list(content.keys())}") + + parts = content.get("parts", []) + print(f"Parts 数量: {len(parts)}") + + for j, part in enumerate(parts): + print(f"\nPart {j} 的键: {list(part.keys())}") + + # 检查不同的可能字段名 + if "inline_data" in part: + print("找到 inline_data!") + img_base64 = part["inline_data"]["data"] + mime_type = part["inline_data"].get("mime_type", "image/jpeg") + elif "inlineData" in part: + print("找到 inlineData!") + img_base64 = part["inlineData"]["data"] + mime_type = part["inlineData"].get("mimeType", "image/jpeg") + elif "data" in part: + print("找到 data 字段!") + img_base64 = part["data"] + mime_type = part.get("mimeType", part.get("mime_type", "image/jpeg")) + else: + print(f"未找到图片数据，part 内容: {part}") + continue + + # 解码 base64 并保存图片 + img_bytes = base64.b64decode(img_base64) + + # 根据 mime_type 确定文件扩展名 + ext = "jpg" if "jpeg" in mime_type.lower() else mime_type.split("/")[-1] + output_path = f"./assets/output_image_{i}_{j}.{ext}" + + with open(output_path, "wb") as f: + f.write(img_bytes) + + # 计算图片大小 + img_size_kb = len(img_bytes) / 1024 + print(f"✓ 图片已保存到: {output_path}") + print(f" 文件大小: {img_size_kb:.2f} KB") + saved_images += 1 + + # 显示总结信息 + print(f"\n{'='*50}") + print(f"📊 生成总结:") + print(f" - 成功生成并保存: {saved_images} 张图片") + print(f" - 总耗时: {generation_time:.2f} 秒") + if saved_images > 0: + print(f" - 平均每张耗时: {generation_time/saved_images:.2f} 秒") + print(f"{'='*50}") + +except Exception as e: + print(f"\n❌ 提取图片时出错: {e}") + import traceback + traceback.print_exc() \ No newline at end of file diff --git a/vector_engine_node.py b/vector_engine_node.py new file mode 100644 index 0000000..20e137b --- /dev/null +++ b/vector_engine_node.py @@ -0,0 +1,317 @@ +import http.client +import json +import base64 +import time +import numpy as np +import torch +from PIL import Image +import io + + +class VectorEngineImageGenerator: + """ + ComfyUI node for Vector Engine Image Generation API + """ + + def __init__(self): + self.api_key = "sk-RZBLe4v8MolmD3fwOd6vdRTCaj7PohDhf4f44UNgWHEAA4zF" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "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": (["1:1", "2:3", "3:2", "4:3", "3:4", "16:9", "9:16"], { + "default": "1:1" + }), + "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 tensor_to_base64(self, tensor, mime_type="image/png"): + """ + Convert ComfyUI image tensor to base64 string + tensor shape: [batch, height, width, channels] with values in [0, 1] + """ + # 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 bytes + buffer = io.BytesIO() + if "png" in mime_type: + pil_image.save(buffer, format="PNG") + else: + pil_image.save(buffer, format="JPEG") + + # Encode to base64 + img_bytes = buffer.getvalue() + base64_str = base64.b64encode(img_bytes).decode('utf-8') + + return base64_str, mime_type + + 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 generate_image(self, prompt, system_prompt, aspect_ratio, image_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 + """ + start_time = time.time() + + 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 + for idx, img in enumerate(images, 1): + if img is not None: + image_count += 1 + base64_data, mime_type = self.tensor_to_base64(img, "image/png") + 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 + }, + "seed": seed + } + }) + + headers = { + 'Authorization': f'Bearer {self.api_key}', + 'Content-Type': 'application/json' + } + + # Make API request + conn.request("POST", + f"/v1beta/models/gemini-3-pro-image-preview:generateContent?key={self.api_key}", + payload, headers) + + res = conn.getresponse() + data = res.read() + + # Parse response + response_json = json.loads(data.decode("utf-8")) + + # Calculate generation time + generation_time = time.time() - start_time + + # 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="gemini-3-pro-image-preview", + aspect_ratio=aspect_ratio, + image_size=image_size, + generation_time=generation_time, + success=False, + error_message=error_msg, + input_images=image_count, + seed=seed + ) + + # 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 + output_tensor = self.base64_to_tensor(img_base64) + + # Get actual dimensions + _, height, width, _ = output_tensor.shape + + # Format info text + info_text = self._format_info( + model_name="gemini-3-pro-image-preview", + aspect_ratio=aspect_ratio, + image_size=image_size, + resolution=f"{width}x{height}", + generation_time=generation_time, + success=True, + input_images=image_count, + seed=seed + ) + + return (output_tensor, info_text) + + # If we reach here, no image was found + info_text = self._format_info( + model_name="gemini-3-pro-image-preview", + aspect_ratio=aspect_ratio, + image_size=image_size, + generation_time=generation_time, + success=False, + error_message="No image found in API response", + input_images=image_count, + seed=seed + ) + + black_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) + return (black_image, info_text) + + except Exception as e: + generation_time = time.time() - start_time + + info_text = self._format_info( + model_name="gemini-3-pro-image-preview", + aspect_ratio=aspect_ratio, + image_size=image_size, + generation_time=generation_time, + success=False, + error_message=str(e), + input_images=sum(1 for img in images if img is not None), + seed=seed + ) + + # 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): + """ + 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}", + f"Generation Time: {generation_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" +} +