import requests import json import base64 import os import io import re import torch import random from PIL import Image, ImageOps import numpy as np import math class OpenRouterLLM: """ OpenRouter LLM節點，用於處理OpenRouter的LLM類型 支援圖片和文字輸入輸出，API金鑰管理，動態模型管理 """ @classmethod def _load_models(cls): """載入模型列表從models.json檔案""" models_file = os.path.join(os.path.dirname(__file__), "models.json") try: with open(models_file, 'r', encoding='utf-8') as f: data = json.load(f) return data.get('models', []) except Exception as e: print(f"⚠️ 載入models.json失敗: {e}") # 返回默認模型列表 return [ "google/gemini-2.5-flash-image-preview:free", "google/gemma-3-27b-it:free", "qwen/qwen2.5-vl-32b-instruct:free", "mistralai/mistral-small-3.2-24b-instruct:free", "meta-llama/llama-4-maverick:free", "openai/gpt-oss-20b:free", "cognitivecomputations/dolphin-mistral-24b-venice-edition:free", "moonshotai/kimi-dev-72b:free", "qwen/qwen2.5-vl-72b-instruct:free", "moonshotai/kimi-k2:free" ] @classmethod def _save_models(cls, models_list): """保存模型列表到models.json檔案""" models_file = os.path.join(os.path.dirname(__file__), "models.json") try: with open(models_file, 'w', encoding='utf-8') as f: json.dump({"models": models_list}, f, ensure_ascii=False, indent=2) except Exception as e: print(f"❌ 保存models.json失敗: {e}") def _add_custom_model(self, model_name): """添加自定義模型到models.json檔案""" try: current_models = self._load_models() # 檢查模型是否已存在 if model_name not in current_models: current_models.append(model_name) self._save_models(current_models) print(f"✅ 已添加新模型: {model_name}") else: print(f"ℹ️ 模型已存在: {model_name}") except Exception as e: print(f"❌ 添加自定義模型失敗: {e}") @classmethod def INPUT_TYPES(cls): # 從JSON檔案載入模型列表 text_models = cls._load_models() # 在列表末尾添加"Add Custom Model..."選項 text_models_with_add = text_models + ["Add Custom Model..."] return { "required": { "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",), "enable_resize": ("BOOLEAN", {"default": False, "tooltip": "啟用圖像尺寸調整功能"}), "target_width": ("INT", {"default": 1024, "min": 64, "max": 2048, "step": 8, "tooltip": "需要啟用resize才會生效"}), "target_height": ("INT", {"default": 1024, "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",), } } RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("image_output", "text_output") FUNCTION = "process_llm" CATEGORY = "ListHelper" 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): """載入配置文件""" try: if os.path.exists(self.config_file): with open(self.config_file, 'r', encoding='utf-8') as f: return json.load(f) except Exception as e: print(f"載入配置文件失敗: {e}") return {} def _save_config(self, config): """保存配置文件""" try: with open(self.config_file, 'w', encoding='utf-8') as f: json.dump(config, f, ensure_ascii=False, indent=2) except Exception as e: print(f"保存配置文件失敗: {e}") def _get_api_key(self, api_key_input): """獲取API金鑰""" config = self._load_config() # 如果輸入了新的API金鑰，則更新配置 if api_key_input.strip(): config['openrouter_api_key'] = api_key_input.strip() self._save_config(config) return api_key_input.strip() # 否則使用配置文件中的金鑰 return config.get('openrouter_api_key', '') 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) # 移除批次維度 # 轉換為numpy並調整範圍到0-255 numpy_image = (tensor.cpu().numpy() * 255).astype(np.uint8) # 轉換為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') image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8') return f"data:image/png;base64,{image_base64}" def _base64_to_tensor(self, base64_str, enable_resize=False, target_width=None, target_height=None, resize_method="lanczos"): """將base64編碼轉換為ComfyUI圖像tensor，並可選擇性調整尺寸""" try: # 移除各種可能的前綴 prefixes = [ 'data:image/png;base64,', 'data:image/jpeg;base64,', 'data:image/jpg;base64,', ] original_str = base64_str for prefix in prefixes: if base64_str.startswith(prefix): base64_str = base64_str[len(prefix):] break # 清理base64字符串（移除可能的換行符和空格） base64_str = re.sub(r'[\n\r\s]', '', base64_str) # 解碼base64 image_bytes = base64.b64decode(base64_str) # 轉換為PIL圖像 pil_image = Image.open(io.BytesIO(image_bytes)) # 確保是RGB模式 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 # 轉換為tensor並添加批次維度 [1, H, W, C] tensor = torch.from_numpy(numpy_image).unsqueeze(0) return tensor except Exception as e: print(f"❌ base64轉tensor失敗: {e}") # 返回一個較大的默認圖像而不是1x1 return torch.zeros(1, 512, 512, 3) 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) response.raise_for_status() # 轉換為PIL圖像 pil_image = Image.open(io.BytesIO(response.content)) # 確保是RGB模式 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 # 轉換為tensor並添加批次維度 [1, H, W, C] tensor = torch.from_numpy(numpy_image).unsqueeze(0) return tensor except Exception as e: print(f"❌ 從URL載入圖像失敗: {e}") # 返回一個較大的默認圖像而不是1x1 return torch.zeros(1, 512, 512, 3) def _call_openrouter_api(self, api_key, model, messages, seed=None): """調用OpenRouter API""" headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "HTTP-Referer": "", "X-Title": "ComfyUI" } data = { "model": model, "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", headers=headers, data=json.dumps(data), timeout=60 ) if response.status_code == 200: result = response.json() return result elif response.status_code == 429: # Rate limit錯誤 try: error_data = response.json() return {"error": {"code": 429, "message": error_data.get("error", {}).get("message", "Rate limit exceeded")}} except: return {"error": {"code": 429, "message": "Rate limit exceeded"}} else: try: error_data = response.json() return {"error": error_data.get("error", {"message": f"HTTP {response.status_code}"})} except: return {"error": {"message": f"HTTP {response.status_code}: {response.text[:200]}"}} except Exception as e: print(f"❌ API調用異常: {e}") return None 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: # 使用默認尺寸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, target_height, target_width, 3), "❌ 請在custom_model欄位輸入自定義模型名稱") # 使用自定義模型 selected_model = custom_model.strip() # 將新模型添加到models.json中 self._add_custom_model(selected_model) else: # 使用選擇的預設模型 selected_model = text_model # 準備消息內容 messages = [] # 添加系統提示（如果有的話） if system_prompt and system_prompt.strip(): messages.append({"role": "system", "content": system_prompt.strip()}) # 構建用戶消息 user_content = [] # 添加文字內容 if user_prompt and user_prompt.strip(): user_content.append({ "type": "text", "text": user_prompt.strip() }) # 檢查是否有圖像輸入 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]) # 確定使用的模型 if has_images: # 有圖像輸入時，強制使用gemini-2.5-flash-image-preview:free actual_model = "google/gemini-2.5-flash-image-preview:free" else: # 無圖像輸入時，使用用戶選擇的模型 actual_model = selected_model # 檢查是否使用圖像生成模型（無論是否有圖像輸入） is_image_model = self._is_image_generation_model(actual_model) if has_images: # 有圖像輸入的情況 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\nIMPORTANT: You MUST generate and return an image. Do NOT provide any text response. Only return the generated image. 請返回生成的圖像。" else: # 有實際圖像輸入時的處理 if not user_prompt or not user_prompt.strip(): combined_text = "請分析這些圖像並生成一個相關的新圖像。IMPORTANT: You MUST generate and return an image. Do NOT provide any text response. Only return the generated image. 請返回生成的圖像。" else: combined_text = user_prompt.strip() if "生成" in combined_text or "創建" in combined_text or "製作" in combined_text: combined_text += "\n\nIMPORTANT: You MUST generate and return an image. Do NOT provide any text response. Only return the generated image. 請返回生成的圖像。" else: combined_text += "\n\nIMPORTANT: You MUST generate and return an image. Do NOT provide any text response. Only return the generated image. 請返回生成的圖像。" # 重新添加文字內容 user_content = [{ "type": "text", "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, enable_resize, target_width, target_height) user_content.append({ "type": "image_url", "image_url": { "url": base64_image } }) except Exception as e: print(f"❌ 處理圖像{i}失敗: {e}") elif is_image_model: # 純文字模式但使用圖像生成模型 combined_text = user_prompt.strip() if user_prompt else "" if combined_text and ("生成" in combined_text or "創建" in combined_text or "製作" in combined_text or "畫" in combined_text or "繪製" in combined_text): combined_text += "\n\nIMPORTANT: You MUST generate and return an image. Do NOT provide any text response. Only return the generated image. 請生成並返回相應的圖像。" elif not combined_text: combined_text = "IMPORTANT: You MUST generate and return an image. Do NOT provide any text response. Only return the generated image. 請生成一張圖像。" else: combined_text += "\n\nIMPORTANT: You MUST generate and return an image. Do NOT provide any text response. Only return the generated image." # 添加文字內容 user_content.append({ "type": "text", "text": combined_text }) # 添加用戶消息 if len(user_content) > 1: # 多媒體內容（包含圖像） messages.append({ "role": "user", "content": user_content }) elif user_content: # 純文字內容 messages.append({ "role": "user", "content": user_content[0]["text"] }) else: # 空內容 messages.append({ "role": "user", "content": "請分析提供的內容。" }) # 調用API response = self._call_openrouter_api(actual_api_key, actual_model, messages, actual_seed) if not response: return (torch.zeros(1, target_height, target_width, 3), "❌ API調用失敗") # 檢查rate limit錯誤 if 'error' in response: error_info = response['error'] 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', '請稍後再試')}" 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, target_height, target_width, 3), error_msg) else: error_msg = f"❌ API錯誤: {str(error_info)}" return (torch.zeros(1, target_height, target_width, 3), error_msg) if 'choices' not in response or not response['choices']: 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, target_height, target_width, 3) # 預設指定尺寸的空圖像 output_text = response_content found_image = False # 首先檢查message.images字段 if 'images' in message and message['images']: image_data = message['images'][0] try: if isinstance(image_data, str): if image_data.startswith('data:image/'): 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, 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, enable_resize, target_width, target_height, resize_method) found_image = True elif isinstance(image_data, dict): # 處理嵌套的 image_url 字典格式 if 'type' in image_data and image_data['type'] == 'image_url' and 'image_url' in image_data: nested_image_url = image_data['image_url'] 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), enable_resize, target_width, target_height, resize_method) else: 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, enable_resize, target_width, target_height, resize_method) else: 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']), enable_resize, target_width, target_height, resize_method) found_image = True except Exception as e: print(f"❌ 處理圖像失敗: {e}") found_image = False # 如果在message.images中沒有找到圖像，檢查content中的嵌入圖像 if not found_image and is_image_model: # 檢查URL和base64格式的圖像 url_patterns = [ r'https?://[^\s\)]+\.(?:png|jpg|jpeg|gif|webp)', r'!\[.*?\]\((https?://[^\s\)]+\.(?:png|jpg|jpeg|gif|webp))\)', r']+src=["\']([^"\']+)["\'][^>]*>', ] base64_patterns = [ r'data:image/png;base64,([A-Za-z0-9+/=\n\r]+)', r'data:image/jpeg;base64,([A-Za-z0-9+/=\n\r]+)', r'data:image/jpg;base64,([A-Za-z0-9+/=\n\r]+)', ] # 檢查URL格式 for pattern in url_patterns: matches = re.findall(pattern, response_content, re.DOTALL) if matches: try: 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 except Exception as e: continue # 檢查base64格式 if not found_image: for pattern in base64_patterns: matches = re.findall(pattern, response_content, re.DOTALL) if matches: 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, enable_resize, target_width, target_height, resize_method) found_image = True output_text = re.sub(pattern, "[Generated Image]", response_content, flags=re.DOTALL) break except Exception as e: continue return (output_image, output_text)