202 lines
8.0 KiB
Python
202 lines
8.0 KiB
Python
import os
|
|
import json
|
|
from typing import Dict, Any
|
|
import requests
|
|
|
|
# Supported Models List (Must be synced with prompt_optimizer.js if manual sync is required)
|
|
SUPPORTED_MODELS = {
|
|
"openai": [
|
|
"gpt-5.1", # Example placeholder
|
|
"gpt-4o",
|
|
"gpt-5-nano",
|
|
"gpt-5-mini",
|
|
"gpt-5.2",
|
|
],
|
|
"anthropic": [
|
|
"claude-sonnet-4-5",
|
|
"claude-haiku-4-5",
|
|
"claude-opus-4-5"
|
|
],
|
|
"google_gemini": [
|
|
"gemini-2.5-pro",
|
|
"gemini-3-flash-preview",
|
|
"gemini-3-pro-preview"
|
|
],
|
|
"xai": [
|
|
"grok-4-1-fast",
|
|
"grok-4-1-fast-reasoning-latest",
|
|
"grok-4-1-fast-non-reasoning-latest",
|
|
],
|
|
"deepseek": [
|
|
"deepseek-chat",
|
|
"deepseek-reasoner",
|
|
],
|
|
"qwen": [
|
|
"qwen-plus",
|
|
"qwen3-max",
|
|
"qwen-flash"
|
|
],
|
|
}
|
|
|
|
class PromptOptimizer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# API Providers list
|
|
api_providers = [
|
|
"openai", "anthropic", "google_gemini", "xai",
|
|
"deepseek", "qwen"
|
|
]
|
|
|
|
# Flatten models list for validation
|
|
all_models = []
|
|
for manufacturer_models in SUPPORTED_MODELS.values():
|
|
all_models.extend(manufacturer_models)
|
|
|
|
# Ensure unique and sorted
|
|
all_models = sorted(list(set(all_models)))
|
|
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": "a cute cat"}),
|
|
"provider": (api_providers, {"default": "deepseek"}),
|
|
"model": (all_models, {"default": "deepseek-chat"}),
|
|
"api_key": ("STRING", {"default": "", "multiline": False}),
|
|
"temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.1}),
|
|
"max_tokens": ("INT", {"default": 1024, "min": 50, "max": 8192}),
|
|
"system_prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "You are a professional AI prompt engineer for image generation. "
|
|
"Turn the user's short description into a detailed, vivid, artistic prompt "
|
|
"suitable for Stable Diffusion / Flux / SDXL / Pony etc. "
|
|
"Use descriptive language, styles, lighting, composition, quality enhancers. "
|
|
"Output ONLY the optimized prompt, no explanations."
|
|
}),
|
|
},
|
|
"hidden": {"extra_pnginfo": "extra_pnginfo", "prompt": "prompt"}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("optimized_prompt",)
|
|
FUNCTION = "optimize"
|
|
CATEGORY = "utils/text"
|
|
|
|
def optimize(self, prompt, provider, model, api_key, temperature, max_tokens, system_prompt, **kwargs):
|
|
# API Key 检查与获取
|
|
if not api_key.strip():
|
|
env_key = f"{provider.upper()}_API_KEY"
|
|
api_key = os.getenv(env_key) or os.getenv("API_KEY")
|
|
if not api_key:
|
|
# 尝试其他可能的环境变量命名
|
|
if provider == "google_gemini":
|
|
api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
|
|
elif provider == "qwen":
|
|
api_key = os.getenv("DASHSCOPE_API_KEY") or os.getenv("QWEN_API_KEY")
|
|
|
|
if not api_key:
|
|
raise ValueError(f"No API key for {provider} (tried env: {env_key})")
|
|
|
|
# OpenAI 兼容的 providers
|
|
openai_compatible = ["openai", "groq", "xai", "deepseek", "qwen"]
|
|
|
|
if provider in openai_compatible:
|
|
base_urls = {
|
|
"openai": "https://api.openai.com/v1",
|
|
"groq": "https://api.groq.com/openai/v1",
|
|
"xai": "https://api.x.ai/v1",
|
|
"deepseek": "https://api.deepseek.com",
|
|
"qwen": "https://dashscope.aliyuncs.com/compatible-mode/v1", # 已修改为国内版 Endpoint
|
|
}
|
|
url = f"{base_urls[provider]}/chat/completions"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
}
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
"temperature": temperature,
|
|
}
|
|
|
|
# OpenAI 新模型强制要求使用 max_completion_tokens
|
|
if provider == "openai":
|
|
payload["max_completion_tokens"] = max_tokens
|
|
else:
|
|
payload["max_tokens"] = max_tokens
|
|
|
|
# Anthropic 特殊处理
|
|
elif provider == "anthropic":
|
|
url = "https://api.anthropic.com/v1/messages"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"x-api-key": api_key,
|
|
"anthropic-version": "2023-06-01",
|
|
}
|
|
payload = {
|
|
"model": model,
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"system": system_prompt,
|
|
}
|
|
|
|
# Gemini 特殊处理
|
|
elif provider == "google_gemini":
|
|
url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
|
|
headers = {"Content-Type": "application/json"}
|
|
payload = {
|
|
"contents": [
|
|
{"role": "user", "parts": [{"text": system_prompt + "\n\n" + prompt}]}
|
|
],
|
|
"generationConfig": {
|
|
"temperature": temperature,
|
|
"maxOutputTokens": max_tokens,
|
|
}
|
|
}
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported provider: {provider}")
|
|
|
|
try:
|
|
response = requests.post(url, headers=headers, json=payload, timeout=90)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
if provider in openai_compatible:
|
|
optimized = data["choices"][0]["message"]["content"].strip()
|
|
elif provider == "anthropic":
|
|
optimized = data["content"][0]["text"].strip()
|
|
elif provider == "google_gemini":
|
|
# 兼容 Gemini 可能的返回结构
|
|
try:
|
|
candidates = data.get("candidates", [])
|
|
# 检查是否因为 Token 耗尽而导致截断（常见于 reasoning 模型）
|
|
if candidates and candidates[0].get("finishReason") == "MAX_TOKENS":
|
|
# 如果没有 content 或者是空的
|
|
content = candidates[0].get("content", {})
|
|
if "parts" not in content or not content["parts"]:
|
|
err_msg = "[WARNING] Output truncated! Gemini used all tokens for thinking. Please increase 'max_tokens' widget value."
|
|
print(err_msg)
|
|
raise ValueError(err_msg)
|
|
|
|
optimized = candidates[0]["content"]["parts"][0]["text"].strip()
|
|
except (KeyError, IndexError, TypeError):
|
|
# 有时候可能是 block 等情况
|
|
print(f"Gemini raw response: {data}")
|
|
raise ValueError(f"Unexpected response format from Gemini: {data}")
|
|
else:
|
|
optimized = data.get("text", data.get("content", "No content found")).strip()
|
|
|
|
return (optimized,)
|
|
|
|
except Exception as e:
|
|
error_msg = f"API call failed for {provider}: {str(e)}"
|
|
if 'response' in locals():
|
|
error_msg += f"\nResponse: {response.text}"
|
|
print(error_msg)
|
|
# 在 ComfyUI 里报错最好直接 raise，但也有些用户喜欢返回错误字符串防止流程中断
|
|
# 这里按照原设计返回错误字符串
|
|
return (f"[ERROR] {error_msg}",)
|