feat:Init commit

.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
+

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 <repository-url> 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.
+

__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']
+

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"
+

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()

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"
+}
+