update for custom_size

2025-09-03 22:31:06 +08:00
parent c6f07f7742
commit b82efc1950
2 changed files with 197 additions and 43 deletions
@@ -5,8 +5,10 @@ import os
 import io
 import re
 import torch
-from PIL import Image
+import random
+from PIL import Image, ImageOps
 import numpy as np
+import math

 class OpenRouterLLM:
    """
@@ -76,10 +78,15 @@ class OpenRouterLLM:
                "api_key": ("STRING", {"multiline": False, "default": "", "placeholder": "輸入您的OpenRouter API金鑰"}),
                "user_prompt": ("STRING", {"multiline": True, "default": "Please analyze the provided content."}),
                "text_model": (text_models_with_add, {"default": text_models_with_add[0]}),
+                "seed": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff, "step": 1, "tooltip": "隨機種子控制(-1為隨機)。注意：圖像生成模型(如Gemini)不支援seed參數"}),
            },
            "optional": {
                "custom_model": ("STRING", {"multiline": False, "default": "", "placeholder": "輸入自定義模型名稱 (選擇Add Custom Model...時使用)"}),
-                "system_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "system_prompt": ("STRING",),
+                "enable_resize": ("BOOLEAN", {"default": False, "tooltip": "啟用圖像尺寸調整功能"}),
+                "target_width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8, "tooltip": "需要啟用resize才會生效"}),
+                "target_height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8, "tooltip": "需要啟用resize才會生效"}),
+                "resize_method": (["bicubic", "lanczos", "bilinear", "nearest"], {"default": "lanczos", "tooltip": "需要啟用resize才會生效"}),
                "image_input_1": ("IMAGE",),
                "image_input_2": ("IMAGE",),
                "image_input_3": ("IMAGE",),
@@ -93,6 +100,27 @@ class OpenRouterLLM:
    
    def __init__(self):
        self.config_file = "config.json"
+    
+    def _is_image_generation_model(self, model_name):
+        """檢查是否為圖像生成模型"""
+        image_generation_models = [
+            "gemini-2.5-flash-image-preview",
+            "gemini-image",  # 可能的其他變體
+        ]
+        return any(img_model in model_name.lower() for img_model in image_generation_models)
+    
+    def _get_default_tensor_size(self, enable_resize, target_width, target_height):
+        """獲取默認tensor尺寸"""
+        if enable_resize:
+            return (1, target_height, target_width, 3)
+        else:
+            return (1, 512, 512, 3)
+    
+    def _create_blank_image_tensor(self, width, height):
+        """創建指定尺寸的空白圖像tensor（白色背景）"""
+        # 創建白色圖像 (1.0 = 白色 in 0-1 range)
+        blank_tensor = torch.ones(1, height, width, 3, dtype=torch.float32)
+        return blank_tensor
        
    def _load_config(self):
        """載入配置文件"""
@@ -125,8 +153,8 @@ class OpenRouterLLM:
        # 否則使用配置文件中的金鑰
        return config.get('openrouter_api_key', '')
    
-    def _tensor_to_base64(self, tensor):
-        """將ComfyUI圖像tensor轉換為base64編碼"""
+    def _tensor_to_base64(self, tensor, enable_resize=False, target_width=None, target_height=None):
+        """將ComfyUI圖像tensor轉換為base64編碼，並可選擇性調整尺寸"""
        # tensor shape: [H, W, C] (0-1 range)
        if tensor.dim() == 4:
            tensor = tensor.squeeze(0)  # 移除批次維度
@@ -137,6 +165,10 @@ class OpenRouterLLM:
        # 轉換為PIL圖像
        pil_image = Image.fromarray(numpy_image)
        
+        # 只有在啟用resize且指定了目標尺寸時，才進行padding處理
+        if enable_resize and target_width is not None and target_height is not None:
+            pil_image = self._pad_and_resize_image(pil_image, target_width, target_height)
+        
        # 轉換為base64
        buffer = io.BytesIO()
        pil_image.save(buffer, format='PNG')
@@ -144,8 +176,8 @@ class OpenRouterLLM:
        
        return f"data:image/png;base64,{image_base64}"
    
-    def _base64_to_tensor(self, base64_str):
-        """將base64編碼轉換為ComfyUI圖像tensor"""
+    def _base64_to_tensor(self, base64_str, enable_resize=False, target_width=None, target_height=None, resize_method="lanczos"):
+        """將base64編碼轉換為ComfyUI圖像tensor，並可選擇性調整尺寸"""
        try:
            # 移除各種可能的前綴
            prefixes = [
@@ -173,6 +205,10 @@ class OpenRouterLLM:
            if pil_image.mode != 'RGB':
                pil_image = pil_image.convert('RGB')
            
+            # 只有在啟用resize且指定了目標尺寸時，才進行縮放處理
+            if enable_resize and target_width is not None and target_height is not None:
+                pil_image = self._resize_image(pil_image, target_width, target_height, resize_method)
+            
            # 轉換為numpy數組
            numpy_image = np.array(pil_image).astype(np.float32) / 255.0
            
@@ -183,10 +219,11 @@ class OpenRouterLLM:
            
        except Exception as e:
            print(f"❌ base64轉tensor失敗: {e}")
-            return torch.zeros(1, 1, 1, 3)
+            # 返回一個較大的默認圖像而不是1x1
+            return torch.zeros(1, 512, 512, 3)
    
-    def _url_to_tensor(self, image_url):
-        """從URL下載圖像並轉換為ComfyUI圖像tensor"""
+    def _url_to_tensor(self, image_url, enable_resize=False, target_width=None, target_height=None, resize_method="lanczos"):
+        """從URL下載圖像並轉換為ComfyUI圖像tensor，並可選擇性調整尺寸"""
        try:
            # 下載圖像
            response = requests.get(image_url, timeout=30)
@@ -199,6 +236,10 @@ class OpenRouterLLM:
            if pil_image.mode != 'RGB':
                pil_image = pil_image.convert('RGB')
            
+            # 只有在啟用resize且指定了目標尺寸時，才進行縮放處理
+            if enable_resize and target_width is not None and target_height is not None:
+                pil_image = self._resize_image(pil_image, target_width, target_height, resize_method)
+            
            # 轉換為numpy數組
            numpy_image = np.array(pil_image).astype(np.float32) / 255.0
            
@@ -209,9 +250,10 @@ class OpenRouterLLM:
            
        except Exception as e:
            print(f"❌ 從URL載入圖像失敗: {e}")
-            return torch.zeros(1, 1, 1, 3)
+            # 返回一個較大的默認圖像而不是1x1
+            return torch.zeros(1, 512, 512, 3)
    
-    def _call_openrouter_api(self, api_key, model, messages):
+    def _call_openrouter_api(self, api_key, model, messages, seed=None):
        """調用OpenRouter API"""
        headers = {
            "Authorization": f"Bearer {api_key}",
@@ -225,6 +267,20 @@ class OpenRouterLLM:
            "messages": messages
        }
        
+        # 檢查是否為圖像生成模型
+        is_image_generation_model = self._is_image_generation_model(model)
+        
+        # 對於圖像生成模型，添加特殊參數
+        if is_image_generation_model:
+            # 確保包含image輸出模式
+            data["modalities"] = ["image", "text"]
+            # 不設置max_tokens，讓模型有足夠空間生成圖像
+            
+        # 只對非圖像生成模型設置seed參數
+        if seed is not None and not is_image_generation_model:
+            data["seed"] = seed
+        
+        
        try:
            response = requests.post(
                "https://openrouter.ai/api/v1/chat/completions",
@@ -234,7 +290,8 @@ class OpenRouterLLM:
            )
            
            if response.status_code == 200:
-                return response.json()
+                result = response.json()
+                return result
            elif response.status_code == 429:
                # Rate limit錯誤
                try:
@@ -253,19 +310,100 @@ class OpenRouterLLM:
            print(f"❌ API調用異常: {e}")
            return None
    
-    def process_llm(self, api_key, user_prompt, text_model, custom_model="", system_prompt="",
+    def _pad_and_resize_image(self, pil_image, target_width, target_height):
+        """使用比例縮放+padding的方式處理圖像到指定尺寸"""
+        try:
+            original_width, original_height = pil_image.size
+            target_ratio = target_width / target_height
+            original_ratio = original_width / original_height
+            
+            # 先根據比例關係決定如何縮放
+            if original_ratio > target_ratio:
+                # 原圖較寬，以寬度為準縮放
+                scale_ratio = target_width / original_width
+                new_width = target_width
+                new_height = int(original_height * scale_ratio)
+            else:
+                # 原圖較高，以高度為準縮放
+                scale_ratio = target_height / original_height
+                new_width = int(original_width * scale_ratio)
+                new_height = target_height
+            
+            # 縮放圖像
+            pil_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
+            
+            # 計算需要的padding
+            pad_width = target_width - new_width
+            pad_height = target_height - new_height
+            
+            # 應用padding（居中）
+            if pad_width > 0 or pad_height > 0:
+                left_pad = pad_width // 2
+                right_pad = pad_width - left_pad
+                top_pad = pad_height // 2
+                bottom_pad = pad_height - top_pad
+                
+                # 使用黑色填充padding區域
+                pil_image = ImageOps.expand(pil_image, (left_pad, top_pad, right_pad, bottom_pad), fill=(0, 0, 0))
+            
+            # 檢查是否需要進一步縮放（如果任一邊超過1536）
+            current_width, current_height = pil_image.size
+            if current_width > 1536 or current_height > 1536:
+                # 保持比例縮放到1536以內
+                final_scale_ratio = min(1536 / current_width, 1536 / current_height)
+                final_width = int(current_width * final_scale_ratio)
+                final_height = int(current_height * final_scale_ratio)
+                pil_image = pil_image.resize((final_width, final_height), Image.LANCZOS)
+            
+            return pil_image
+            
+        except Exception as e:
+            print(f"❌ 圖像padding處理失敗: {e}")
+            return pil_image
+    
+    def _resize_image(self, pil_image, target_width, target_height, resize_method="lanczos"):
+        """使用指定方法縮放圖像到目標尺寸"""
+        try:
+            # 選擇縮放方法
+            method_map = {
+                "bicubic": Image.BICUBIC,
+                "lanczos": Image.LANCZOS,
+                "bilinear": Image.BILINEAR,
+                "nearest": Image.NEAREST
+            }
+            
+            resize_filter = method_map.get(resize_method, Image.LANCZOS)
+            return pil_image.resize((target_width, target_height), resize_filter)
+            
+        except Exception as e:
+            print(f"❌ 圖像縮放失敗: {e}")
+            return pil_image
+
+    def process_llm(self, api_key, user_prompt, text_model, seed=-1, custom_model="", system_prompt=None, 
+                   enable_resize=False, target_width=512, target_height=512, resize_method="lanczos",
                   image_input_1=None, image_input_2=None, image_input_3=None):
        """處理LLM請求"""
        
        # 獲取API金鑰
        actual_api_key = self._get_api_key(api_key)
        if not actual_api_key:
-            return (torch.zeros(1, 1, 1, 3), "❌ 錯誤: 請提供OpenRouter API金鑰")
+            # 使用默認尺寸512x512當enable_resize關閉時
+            default_height = target_height if enable_resize else 512
+            default_width = target_width if enable_resize else 512
+            return (torch.zeros(*self._get_default_tensor_size(enable_resize, target_width, target_height)), "❌ 錯誤: 請提供OpenRouter API金鑰")
+        
+        # 處理種子設定 - Control After Generate支援
+        if seed == -1:
+            # 隨機種子
+            actual_seed = random.randint(0, 2147483647)
+        else:
+            # 固定種子
+            actual_seed = seed
        
        # 處理模型選擇
        if text_model == "Add Custom Model...":
            if not custom_model or not custom_model.strip():
-                return (torch.zeros(1, 1, 1, 3), "❌ 請在custom_model欄位輸入自定義模型名稱")
+                return (torch.zeros(1, target_height, target_width, 3), "❌ 請在custom_model欄位輸入自定義模型名稱")
            
            # 使用自定義模型
            selected_model = custom_model.strip()
@@ -294,6 +432,15 @@ class OpenRouterLLM:
            })
        
        # 檢查是否有圖像輸入
+        original_has_images = any([image_input_1 is not None, image_input_2 is not None, image_input_3 is not None])
+        
+        # 如果enable_resize開啟且沒有圖像輸入，自動生成空白圖像給image_input_1
+        auto_generated_blank = False
+        if enable_resize and not original_has_images:
+            image_input_1 = self._create_blank_image_tensor(target_width, target_height)
+            auto_generated_blank = True
+        
+        # 重新檢查圖像輸入狀況（包含自動生成的空白圖像）
        has_images = any([image_input_1 is not None, image_input_2 is not None, image_input_3 is not None])
        
        # 確定使用的模型
@@ -305,17 +452,25 @@ class OpenRouterLLM:
            actual_model = selected_model
        
        # 檢查是否使用圖像生成模型（無論是否有圖像輸入）
-        is_image_model = actual_model == "google/gemini-2.5-flash-image-preview:free"
+        is_image_model = self._is_image_generation_model(actual_model)
        
        if has_images:
            # 有圖像輸入的情況
-            # 為圖像生成添加特殊提示
-            if not user_prompt or not user_prompt.strip():
-                combined_text = "請分析這些圖像並生成一個相關的新圖像。請返回生成的圖像。"
+            if auto_generated_blank:
+                # 使用自動生成的空白圖像時的特殊處理
+                if not user_prompt or not user_prompt.strip():
+                    combined_text = "Please take a look at the size of image 1 and give a picture of a beautiful landscape with the same dimensions."
+                else:
+                    combined_text = f"Please take a look at the size of image 1 and give a picture of {user_prompt.strip()}"
+                combined_text += "\n\n請返回生成的圖像。"
            else:
-                combined_text = user_prompt.strip()
-                if "生成" in combined_text or "創建" in combined_text or "製作" in combined_text:
-                    combined_text += "\n\n請返回生成的圖像。"
+                # 有實際圖像輸入時的處理
+                if not user_prompt or not user_prompt.strip():
+                    combined_text = "請分析這些圖像並生成一個相關的新圖像。請返回生成的圖像。"
+                else:
+                    combined_text = user_prompt.strip()
+                    if "生成" in combined_text or "創建" in combined_text or "製作" in combined_text:
+                        combined_text += "\n\n請返回生成的圖像。"
            
            # 重新添加文字內容
            user_content = [{
@@ -323,11 +478,11 @@ class OpenRouterLLM:
                "text": combined_text
            }]
            
-            # 添加圖像到用戶消息
+            # 添加圖像到用戶消息（帶padding處理）
            for i, image_input in enumerate([image_input_1, image_input_2, image_input_3], 1):
                if image_input is not None:
                    try:
-                        base64_image = self._tensor_to_base64(image_input)
+                        base64_image = self._tensor_to_base64(image_input, enable_resize, target_width, target_height)
                        user_content.append({
                            "type": "image_url",
                            "image_url": {
@@ -372,10 +527,10 @@ class OpenRouterLLM:
        
        
        # 調用API
-        response = self._call_openrouter_api(actual_api_key, actual_model, messages)
+        response = self._call_openrouter_api(actual_api_key, actual_model, messages, actual_seed)
        
        if not response:
-            return (torch.zeros(1, 1, 1, 3), "❌ API調用失敗")
+            return (torch.zeros(1, target_height, target_width, 3), "❌ API調用失敗")
        
        # 檢查rate limit錯誤
        if 'error' in response:
@@ -383,24 +538,23 @@ class OpenRouterLLM:
            if isinstance(error_info, dict) and 'code' in error_info:
                if error_info['code'] == 429 or 'rate_limit' in str(error_info).lower():
                    error_msg = f"⚠️ Rate Limit達到: {error_info.get('message', '請稍後再試')}"
-                    print(error_msg)
-                    return (torch.zeros(1, 1, 1, 3), error_msg)
+                    return (torch.zeros(1, target_height, target_width, 3), error_msg)
                else:
                    error_msg = f"❌ API錯誤: {error_info.get('message', str(error_info))}"
-                    return (torch.zeros(1, 1, 1, 3), error_msg)
+                    return (torch.zeros(1, target_height, target_width, 3), error_msg)
            else:
                error_msg = f"❌ API錯誤: {str(error_info)}"
-                return (torch.zeros(1, 1, 1, 3), error_msg)
+                return (torch.zeros(1, target_height, target_width, 3), error_msg)
        
        if 'choices' not in response or not response['choices']:
-            return (torch.zeros(1, 1, 1, 3), "❌ API回應格式異常")
+            return (torch.zeros(1, target_height, target_width, 3), "❌ API回應格式異常")
        
        # 獲取回應訊息
        message = response['choices'][0]['message']
        response_content = message.get('content', '')
        
        # 檢查回應是否包含圖像數據
-        output_image = torch.zeros(1, 1, 1, 3)  # 預設空圖像
+        output_image = torch.zeros(1, target_height, target_width, 3)  # 預設指定尺寸的空圖像
        output_text = response_content
        found_image = False
        
@@ -411,15 +565,15 @@ class OpenRouterLLM:
            try:
                if isinstance(image_data, str):
                    if image_data.startswith('data:image/'):
-                        output_image = self._base64_to_tensor(image_data)
+                        output_image = self._base64_to_tensor(image_data, enable_resize, target_width, target_height, resize_method)
                        found_image = True
                    elif image_data.startswith('http'):
-                        output_image = self._url_to_tensor(image_data)
+                        output_image = self._url_to_tensor(image_data, enable_resize, target_width, target_height, resize_method)
                        found_image = True
                    else:
                        # 可能是純 base64，添加 data URL 前綴
                        full_data_url = f"data:image/png;base64,{image_data}"
-                        output_image = self._base64_to_tensor(full_data_url)
+                        output_image = self._base64_to_tensor(full_data_url, enable_resize, target_width, target_height, resize_method)
                        found_image = True
                        
                elif isinstance(image_data, dict):
@@ -429,21 +583,21 @@ class OpenRouterLLM:
                        if isinstance(nested_image_url, dict) and 'url' in nested_image_url:
                            image_url_value = nested_image_url['url']
                            if str(image_url_value).startswith('data:image/'):
-                                output_image = self._base64_to_tensor(str(image_url_value))
+                                output_image = self._base64_to_tensor(str(image_url_value), enable_resize, target_width, target_height, resize_method)
                            else:
-                                output_image = self._url_to_tensor(str(image_url_value))
+                                output_image = self._url_to_tensor(str(image_url_value), enable_resize, target_width, target_height, resize_method)
                            found_image = True
                    # 處理直接包含 url 的格式
                    elif 'url' in image_data:
                        url_value = str(image_data['url'])
                        if url_value.startswith('data:image/'):
-                            output_image = self._base64_to_tensor(url_value)
+                            output_image = self._base64_to_tensor(url_value, enable_resize, target_width, target_height, resize_method)
                        else:
-                            output_image = self._url_to_tensor(url_value)
+                            output_image = self._url_to_tensor(url_value, enable_resize, target_width, target_height, resize_method)
                        found_image = True
                    # 處理直接包含 data 的格式
                    elif 'data' in image_data:
-                        output_image = self._base64_to_tensor(str(image_data['data']))
+                        output_image = self._base64_to_tensor(str(image_data['data']), enable_resize, target_width, target_height, resize_method)
                        found_image = True
                    
            except Exception as e:
@@ -470,7 +624,7 @@ class OpenRouterLLM:
                matches = re.findall(pattern, response_content, re.DOTALL)
                if matches:
                    try:
-                        output_image = self._url_to_tensor(matches[0])
+                        output_image = self._url_to_tensor(matches[0], enable_resize, target_width, target_height, resize_method)
                        found_image = True
                        output_text = re.sub(pattern, "[Generated Image]", response_content, flags=re.DOTALL)
                        break
@@ -485,7 +639,7 @@ class OpenRouterLLM:
                        try:
                            clean_base64 = re.sub(r'[\n\r\s]', '', matches[0])
                            base64_image = f"data:image/png;base64,{clean_base64}"
-                            output_image = self._base64_to_tensor(base64_image)
+                            output_image = self._base64_to_tensor(base64_image, enable_resize, target_width, target_height, resize_method)
                            found_image = True
                            output_text = re.sub(pattern, "[Generated Image]", response_content, flags=re.DOTALL)
                            break
@@ -1,7 +1,7 @@
 [project]
 name = "Listhelper"
 description = "The ListHelper collection is a comprehensive set of custom nodes for ComfyUI that provides powerful list manipulation capabilities. This collection includes audio processing, text splitting, and number generation tools for enhanced workflow automation.And OpenRouter API Function"
-version = "1.0.4"
+version = "1.0.5"
 license = {file = "LICENSE"}
 dependencies = ["regex"]