273 lines
11 KiB
Python
273 lines
11 KiB
Python
import asyncio
|
|
import random
|
|
import re
|
|
import time
|
|
import threading
|
|
import hashlib
|
|
|
|
import torch
|
|
from comfy.model_management import InterruptProcessingException
|
|
|
|
from ..services.llm import LLMService
|
|
from ..services.baidu import BaiduTranslateService
|
|
from ..utils.common import format_api_error, format_model_with_thinking, generate_request_id, log_prepare, log_error, TASK_TRANSLATE, SOURCE_NODE
|
|
from ..services.thinking_control import build_thinking_suppression
|
|
from .base import LLMNodeBase
|
|
|
|
|
|
class PromptTranslate(LLMNodeBase):
|
|
"""
|
|
提示词翻译节点
|
|
自动识别输入语言并翻译成目标语言，支持多种翻译服务
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# ---动态获取翻译服务/模型列表(包含硬编码的百度翻译)---
|
|
service_options = cls.get_translate_service_options()
|
|
default_service = service_options[0] if service_options else "百度翻译"
|
|
|
|
return {
|
|
"required": {
|
|
"source_text": ("STRING", {"forceInput": True, "default": "", "multiline": True, "placeholder": "Input text to translate..."}),
|
|
"target_language": (["English", "Chinese"], {"default": "English"}),
|
|
"translate_service": (service_options, {"default": default_service, "tooltip": "Select translation service and model"}),
|
|
# Ollama Automatic VRAM Unload
|
|
"ollama_auto_unload": ("BOOLEAN", {"default": True, "label_on": "Enable", "label_off": "Disable", "tooltip": "Auto unload Ollama model after generation"}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("translated_text",)
|
|
FUNCTION = "translate"
|
|
CATEGORY = "✨Prompt Assistant"
|
|
OUTPUT_NODE = False
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, source_text=None, target_language=None, translate_service=None, ollama_auto_unload=None, unique_id=None):
|
|
"""
|
|
只在输入内容真正变化时才触发重新执行
|
|
使用输入参数的哈希值作为判断依据
|
|
"""
|
|
# 计算文本的哈希值
|
|
text_hash = ""
|
|
if source_text:
|
|
# 使用hashlib计算文本的哈希值，更安全和一致
|
|
text_hash = hashlib.md5(source_text.encode('utf-8')).hexdigest()
|
|
|
|
# 组合所有输入的哈希值
|
|
input_hash = hash((
|
|
text_hash,
|
|
target_language,
|
|
translate_service,
|
|
bool(ollama_auto_unload)
|
|
))
|
|
|
|
return input_hash
|
|
|
|
def _contains_chinese(self, text: str) -> bool:
|
|
"""检查文本是否包含中文字符"""
|
|
if not text:
|
|
return False
|
|
return bool(re.search('[\u4e00-\u9fa5]', text))
|
|
|
|
def _detect_language(self, text: str) -> str:
|
|
"""自动检测文本语言"""
|
|
if not text:
|
|
return "auto"
|
|
|
|
# 检查是否为纯英文 (只包含ASCII可打印字符)
|
|
is_pure_english = bool(re.fullmatch(r'[ -~]+', text))
|
|
# 检查是否包含中文字符
|
|
contains_chinese = self._contains_chinese(text)
|
|
|
|
if contains_chinese:
|
|
return "zh"
|
|
elif is_pure_english:
|
|
return "en"
|
|
else:
|
|
return "auto"
|
|
|
|
def translate(self, source_text, target_language, translate_service, ollama_auto_unload, unique_id=None):
|
|
"""
|
|
翻译文本函数
|
|
"""
|
|
request_id = None # 提升到方法级别作用域
|
|
try:
|
|
# 检查输入
|
|
if not source_text or not source_text.strip():
|
|
return ("",)
|
|
|
|
# 自动检测源语言
|
|
detected_lang = self._detect_language(source_text)
|
|
to_lang = "en" if target_language == "English" else "zh"
|
|
|
|
# 智能跳过翻译逻辑
|
|
skip_translation = False
|
|
if to_lang == 'en' and detected_lang == 'en':
|
|
from ..utils.common import _ANSI_CLEAR_EOL
|
|
print(f"\r{_ANSI_CLEAR_EOL}{self.REQUEST_PREFIX} 检测到英文输入，目标为英文，无需翻译", flush=True)
|
|
skip_translation = True
|
|
elif to_lang == 'zh' and detected_lang == 'zh':
|
|
from ..utils.common import _ANSI_CLEAR_EOL
|
|
print(f"\r{_ANSI_CLEAR_EOL}{self.REQUEST_PREFIX} 检测到中文输入，目标为中文，无需翻译", flush=True)
|
|
skip_translation = True
|
|
|
|
if skip_translation:
|
|
return (source_text,)
|
|
|
|
# 映射语言名称
|
|
lang_map = {'zh': '中文', 'en': '英文', 'auto': '原文'}
|
|
from_lang_name = lang_map.get(detected_lang, detected_lang)
|
|
to_lang_name = lang_map.get(to_lang, to_lang)
|
|
|
|
# ---解析服务/模型字符串---
|
|
service_id, model_name = self.parse_service_model(translate_service)
|
|
if not service_id:
|
|
raise ValueError(f"Invalid service selection: {translate_service}")
|
|
|
|
# ---百度翻译特殊处理---
|
|
if service_id == 'baidu':
|
|
request_id, result = self._translate_with_baidu(source_text, detected_lang, to_lang, translate_service, from_lang_name, to_lang_name, unique_id)
|
|
else:
|
|
# ---LLM翻译:获取服务配置---
|
|
from ..config_manager import config_manager
|
|
service = config_manager.get_service(service_id)
|
|
if not service:
|
|
raise ValueError(f"Service config not found: {translate_service}")
|
|
|
|
request_id, result = self._translate_with_llm(source_text, detected_lang, to_lang, service_id, model_name, service, translate_service, from_lang_name, to_lang_name, ollama_auto_unload, unique_id)
|
|
|
|
if result and result.get('success'):
|
|
translated_text = result.get('data', {}).get('translated', '').strip()
|
|
if not translated_text:
|
|
error_msg = 'API returned empty result'
|
|
raise RuntimeError(f"❌Translation failed: {error_msg}")
|
|
|
|
# 结果阶段日志由服务层统一输出，节点层不再重复打印
|
|
return (translated_text,)
|
|
else:
|
|
error_msg = result.get('error', 'Unknown error') if result else 'No result returned'
|
|
# 如果是中断错误,直接抛出InterruptProcessingException,不打印日志(由基类打印)
|
|
if error_msg == "任务被中断":
|
|
raise InterruptProcessingException()
|
|
log_error(TASK_TRANSLATE, request_id, error_msg)
|
|
raise RuntimeError(f"Translation failed: {error_msg}")
|
|
|
|
except InterruptProcessingException:
|
|
# 不打印日志,由基类统一打印
|
|
raise
|
|
except Exception as e:
|
|
error_msg = format_api_error(e, translate_service)
|
|
log_error(TASK_TRANSLATE, request_id, error_msg)
|
|
raise RuntimeError(f"Translation error: {error_msg}")
|
|
|
|
def _translate_with_baidu(self, text, from_lang, to_lang, service_name, from_lang_name, to_lang_name, unique_id):
|
|
"""使用百度翻译服务"""
|
|
# 创建请求ID
|
|
request_id = generate_request_id("trans", "baidu", unique_id)
|
|
|
|
# 准备阶段日志
|
|
log_prepare(TASK_TRANSLATE, request_id, SOURCE_NODE, "百度翻译", None, None, {"方向": f"{from_lang_name}→{to_lang_name}", "长度": len(text)})
|
|
|
|
# 执行翻译（异步线程 + 可中断）
|
|
result = self._run_llm_task(
|
|
BaiduTranslateService.translate,
|
|
service_name,
|
|
text=text,
|
|
from_lang=from_lang,
|
|
to_lang=to_lang,
|
|
request_id=request_id,
|
|
task_type=TASK_TRANSLATE,
|
|
source=SOURCE_NODE
|
|
)
|
|
|
|
return request_id, result
|
|
|
|
def _translate_with_llm(self, text, from_lang, to_lang, service_id, model_name, service, service_display_name, from_lang_name, to_lang_name, auto_unload, unique_id):
|
|
"""使用LLM翻译服务"""
|
|
# ---构建provider_config---
|
|
# 查找指定的模型或默认模型
|
|
llm_models = service.get('llm_models', [])
|
|
target_model = None
|
|
|
|
if model_name:
|
|
# 查找指定的模型
|
|
target_model = next((m for m in llm_models if m.get('name') == model_name), None)
|
|
|
|
if not target_model:
|
|
# 使用默认模型或第一个模型
|
|
target_model = next((m for m in llm_models if m.get('is_default')),
|
|
llm_models[0] if llm_models else None)
|
|
|
|
if not target_model:
|
|
return {"success": False, "error": f"Service {service_display_name} has no available models"}
|
|
|
|
# 构建配置对象
|
|
provider_config = {
|
|
'provider': service_id,
|
|
'model': target_model.get('name', ''),
|
|
'base_url': service.get('base_url', ''),
|
|
'api_key': service.get('api_key', ''),
|
|
'temperature': target_model.get('temperature', 0.7),
|
|
'max_tokens': target_model.get('max_tokens', 1000),
|
|
'top_p': target_model.get('top_p', 0.9),
|
|
}
|
|
|
|
# Ollama特殊处理:添加auto_unload配置
|
|
if service.get('type') == 'ollama':
|
|
provider_config['auto_unload'] = auto_unload
|
|
|
|
# 创建请求ID
|
|
request_id = generate_request_id("trans", "llm", unique_id)
|
|
|
|
# 检查是否关闭思维链
|
|
model_full_name = provider_config.get('model')
|
|
disable_thinking_enabled = service.get('disable_thinking', True)
|
|
thinking_extra = build_thinking_suppression(service_id, model_full_name) if disable_thinking_enabled else None
|
|
model_display = format_model_with_thinking(model_full_name, bool(thinking_extra))
|
|
|
|
# 获取服务显示名称
|
|
service_display_name = service.get('name', service_id)
|
|
|
|
# 准备阶段日志
|
|
log_prepare(TASK_TRANSLATE, request_id, SOURCE_NODE, service_display_name, model_display, None, {"方向": f"{from_lang_name}→{to_lang_name}", "长度": len(text)})
|
|
|
|
# 检查API密钥和模型
|
|
api_key = provider_config.get('api_key', '')
|
|
model = provider_config.get('model', '')
|
|
|
|
if not api_key or not model:
|
|
return {"success": False, "error": f"Please configure API key and model for {service_display_name}"}
|
|
|
|
# 执行翻译（异步线程 + 可中断）
|
|
result = self._run_llm_task(
|
|
LLMService.translate,
|
|
service_id,
|
|
text=text,
|
|
from_lang=from_lang,
|
|
to_lang=to_lang,
|
|
request_id=request_id,
|
|
stream_callback=None,
|
|
custom_provider=service_id,
|
|
custom_provider_config=provider_config,
|
|
task_type=TASK_TRANSLATE,
|
|
source=SOURCE_NODE
|
|
)
|
|
|
|
return request_id, result
|
|
|
|
|
|
# 节点映射，用于向ComfyUI注册节点
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PromptTranslate": PromptTranslate,
|
|
}
|
|
|
|
# 节点显示名称映射
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PromptTranslate": "✨Prompt Translate",
|
|
}
|