comnote-max-builmenlabo/prompt_translator.py

"""
ComfyUI Prompt Translator Node
Multi-provider translation node with token counting functionality.
"""

import json
import os
import requests
import hashlib
import time
import threading
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, Any, Tuple, Optional
from pathlib import Path

try:
    from transformers import CLIPTokenizer
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False

try:
    from transformers import MarianMTModel, MarianTokenizer
    HAS_MARIANMT = True
except ImportError:
    HAS_MARIANMT = False


class PromptTranslator:
    """
    ComfyUI Prompt Translator Node
    
    Provides multi-provider translation with token counting functionality.
    Supports DeepL, Google Translate, Yandex, Baidu, and MarianMT providers.
    """
    
    def __init__(self):
        self.providers = ["DeepL", "Google", "Yandex", "Baidu", "MarianMT"]
        self.cache_file = Path.home() / ".comfyui" / "prompt_translator_cache.json"
        self.cache_file.parent.mkdir(parents=True, exist_ok=True)
        self.api_cache = self.load_cache()
        
        self.tokenizer = None
        if HAS_TRANSFORMERS:
            self.init_clip_tokenizer()
        else:
            print("Warning: transformers not available, using fallback token counting")
        
        self.marianmt_models = {}
        self.executor = ThreadPoolExecutor(max_workers=4)
        self.translation_cache = {}
        self.last_translation_time = 0
        
        # Language code mappings for different providers
        self.language_codes = {
            "Auto": {"google": "auto", "deepl": None, "yandex": None, "baidu": "auto"},
            "Japanese": {"google": "ja", "deepl": "JA", "yandex": "ja", "baidu": "jp"},
            "English": {"google": "en", "deepl": "EN-US", "yandex": "en", "baidu": "en"},
            "Chinese": {"google": "zh", "deepl": "ZH", "yandex": "zh", "baidu": "zh"},
            "Korean": {"google": "ko", "deepl": "KO", "yandex": "ko", "baidu": "kor"},
            "Spanish": {"google": "es", "deepl": "ES", "yandex": "es", "baidu": "spa"},
            "French": {"google": "fr", "deepl": "FR", "yandex": "fr", "baidu": "fra"},
            "German": {"google": "de", "deepl": "DE", "yandex": "de", "baidu": "de"},
            "Italian": {"google": "it", "deepl": "IT", "yandex": "it", "baidu": "it"},
            "Portuguese": {"google": "pt", "deepl": "PT-PT", "yandex": "pt", "baidu": "pt"},
            "Russian": {"google": "ru", "deepl": "RU", "yandex": "ru", "baidu": "ru"},
            "Arabic": {"google": "ar", "deepl": "AR", "yandex": "ar", "baidu": "ara"},
            "Hindi": {"google": "hi", "deepl": "HI", "yandex": "hi", "baidu": "hi"},
            "Thai": {"google": "th", "deepl": "TH", "yandex": "th", "baidu": "th"},
            "Vietnamese": {"google": "vi", "deepl": "VI", "yandex": "vi", "baidu": "vie"},
            "Indonesian": {"google": "id", "deepl": "ID", "yandex": "id", "baidu": "id"},
            "Dutch": {"google": "nl", "deepl": "NL", "yandex": "nl", "baidu": "nl"},
            "Swedish": {"google": "sv", "deepl": "SV", "yandex": "sv", "baidu": "swe"},
            "Polish": {"google": "pl", "deepl": "PL", "yandex": "pl", "baidu": "pl"},
            "Turkish": {"google": "tr", "deepl": "TR", "yandex": "tr", "baidu": "tr"},
            "Czech": {"google": "cs", "deepl": "CS", "yandex": "cs", "baidu": "cs"},
            "Hungarian": {"google": "hu", "deepl": "HU", "yandex": "hu", "baidu": "hu"}
        }
    
    def get_language_code(self, language_name: str, provider: str) -> str:
        """Get language code for specific provider."""
        provider_key = provider.lower()
        if provider_key == "deepl":
            provider_key = "deepl"
        elif provider_key in ["yandex"]:
            provider_key = "yandex"  
        elif provider_key in ["baidu"]:
            provider_key = "baidu"
        else:
            provider_key = "google"  # Default to Google format
            
        if language_name in self.language_codes:
            code = self.language_codes[language_name].get(provider_key)
            if code is None:
                # Fallback to Google format if provider doesn't support
                code = self.language_codes[language_name].get("google", "en")
            return code
        return "en"  # Default fallback
        
    @classmethod
    def INPUT_TYPES(cls):
        # ComfyUI focused: English ↔ Other Languages
        translation_modes = [
            "English → Other Language", 
            "Other Language → English"
        ]
        
        # Major languages for ComfyUI prompts
        languages = [
            "Japanese", "Chinese", "Korean", "Spanish", "French", 
            "German", "Italian", "Portuguese", "Russian", "Arabic", "Hindi", 
            "Thai", "Vietnamese", "Indonesian", "Dutch", "Swedish", "Polish", 
            "Turkish", "Czech", "Hungarian"
        ]
        
        return {
            "required": {
                "provider": (["Google", "DeepL", "Yandex", "Baidu", "MarianMT"], {"default": "Google"}),
                "input_text": ("STRING", {"multiline": True, "default": "", "placeholder": "Enter text to translate..."}),
                "translation_mode": (translation_modes, {"default": "Other Language → English"}),
                "other_language": (languages, {"default": "Japanese"}),
                "api_key": ("STRING", {"default": "", "placeholder": "API Key (Google/MarianMT不要)"}),
                "execute_translation": ("BOOLEAN", {"default": False, "label_on": "Translating...", "label_off": "Translation"}),
                "translated_text": ("STRING", {"multiline": True, "default": "", "placeholder": "Translation will appear here..."}),
            },
            "optional": {
                "input_text_connection": ("STRING", {"forceInput": True}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "INT")
    RETURN_NAMES = ("original_text", "translated_text", "token_count")
    FUNCTION = "translate"
    CATEGORY = "text/translation"
    
    def load_cache(self) -> Dict[str, str]:
        """Load cached API keys from file."""
        if self.cache_file.exists():
            try:
                with open(self.cache_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                print(f"Warning: Could not load cache file: {e}")
        return {}
    
    def save_cache(self):
        """Save API keys to cache file."""
        try:
            with open(self.cache_file, 'w', encoding='utf-8') as f:
                json.dump(self.api_cache, f, indent=2)
        except Exception as e:
            print(f"Warning: Could not save cache file: {e}")
    
    def count_tokens(self, text: str) -> int:
        """Count tokens using CLIP tokenizer with fallback methods."""
        if not text or not text.strip():
            return 0
        
        # Try CLIP tokenizer first
        if self.tokenizer:
            try:
                tokens = self.tokenizer.encode(text)
                return len(tokens)
            except Exception as e:
                print(f"Warning: CLIP tokenizer failed: {e}")
        
        # Fallback to simple tokenization
        return self.simple_token_count(text)
    
    def simple_token_count(self, text: str) -> int:
        """Simple token counting fallback method."""
        if not text:
            return 0
        
        # Remove extra whitespace and split
        words = text.strip().split()
        
        # Rough estimation: 1 token per word, plus punctuation
        token_count = 0
        for word in words:
            # Count word tokens
            token_count += 1
            
            # Add extra tokens for punctuation
            punctuation_count = sum(1 for char in word if char in '.,!?;:"()[]{}/')
            token_count += punctuation_count
        
        return max(1, token_count)  # Ensure at least 1 token for non-empty text
    
    def init_clip_tokenizer(self):
        """Initialize CLIP tokenizer with multiple fallback options."""
        tokenizer_models = [
            "openai/clip-vit-base-patch32",
            "openai/clip-vit-base-patch16",
            "openai/clip-vit-large-patch14"
        ]
        
        for model_name in tokenizer_models:
            try:
                self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
                print(f"Successfully loaded CLIP tokenizer: {model_name}")
                return
            except Exception as e:
                print(f"Failed to load {model_name}: {e}")
                continue
        
        print("Warning: Could not load any CLIP tokenizer, using fallback method")
        self.tokenizer = None
    
    def translate_deepl(self, text: str, api_key: str, target_lang: str) -> str:
        """Translate using DeepL API with enhanced parameters."""
        # Determine URL based on API key type
        if api_key.endswith(":fx"):
            url = "https://api-free.deepl.com/v2/translate"
        else:
            url = "https://api.deepl.com/v2/translate"
        
        headers = {
            "Authorization": f"DeepL-Auth-Key {api_key}",
            "Content-Type": "application/json"
        }
        
        # Detect source language
        source_lang = "JA" if target_lang == "EN-US" else "EN"
        
        data = {
            "text": [text],
            "source_lang": source_lang,
            "target_lang": target_lang,
            "formality": "default",  # Use default formality
            "preserve_formatting": True,  # Preserve text formatting
            "split_sentences": "1"  # Split sentences for better translation
        }
        
        print(f"DeepL API Request: {url}")
        print(f"DeepL Data: {data}")
        
        try:
            response = requests.post(url, headers=headers, json=data, timeout=15)
            print(f"DeepL Response Status: {response.status_code}")
            
            if response.status_code == 403:
                raise Exception("DeepL API key invalid or quota exceeded")
            elif response.status_code == 456:
                raise Exception("DeepL API quota exceeded")
            
            response.raise_for_status()
            result = response.json()
            
            print(f"DeepL Response: {result}")
            
            if "translations" not in result or not result["translations"]:
                raise Exception("No translation result from DeepL")
            
            translated_text = result["translations"][0]["text"]
            detected_source = result["translations"][0].get("detected_source_language", "Unknown")
            
            print(f"DeepL Translation: '{text}' -> '{translated_text}' (detected: {detected_source})")
            
            return translated_text
            
        except requests.exceptions.RequestException as e:
            raise Exception(f"DeepL API request failed: {str(e)}")
        except KeyError as e:
            raise Exception(f"DeepL API response format error: {str(e)}")
        except Exception as e:
            raise Exception(f"DeepL translation failed: {str(e)}")
    
    def translate_google(self, text: str, api_key: str, target_lang: str) -> str:
        """Translate using Google Translate API."""
        if not api_key.strip():
            # Use free Google Translate service (unofficial)
            return self.translate_google_free(text, target_lang)
            
        url = "https://translation.googleapis.com/language/translate/v2"
        
        params = {
            "key": api_key,
            "q": text,
            "target": target_lang,
            "format": "text"
        }
        
        try:
            response = requests.post(url, params=params, timeout=10)
            response.raise_for_status()
            result = response.json()
            return result["data"]["translations"][0]["translatedText"]
        except Exception as e:
            # Fallback to free service if API fails
            return self.translate_google_free(text, target_lang)
    
    def translate_google_free(self, text: str, target_lang: str) -> str:
        """Translate using free Google Translate service."""
        try:
            import urllib.parse
            import re
            
            # Detect source language
            source_lang = "ja" if target_lang == "en" else "en"
            
            # Construct URL
            encoded_text = urllib.parse.quote(text)
            url = f"https://translate.googleapis.com/translate_a/single?client=gtx&sl={source_lang}&tl={target_lang}&dt=t&q={encoded_text}"
            
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            
            result = response.json()
            if result and result[0] and result[0][0]:
                return result[0][0][0]
            else:
                raise Exception("No translation result")
                
        except Exception as e:
            raise Exception(f"Free Google translation failed: {str(e)}")
    
    def translate_yandex(self, text: str, api_key: str, lang_pair: str) -> str:
        """Translate using Yandex Translate API."""
        url = "https://translate.yandex.net/api/v1.5/tr.json/translate"
        
        params = {
            "key": api_key,
            "text": text,
            "lang": lang_pair,
            "format": "plain"
        }
        
        try:
            response = requests.post(url, params=params, timeout=10)
            response.raise_for_status()
            result = response.json()
            
            if result["code"] != 200:
                raise Exception(f"Yandex API error: {result.get('message', 'Unknown error')}")
            
            return result["text"][0]
        except Exception as e:
            raise Exception(f"Yandex translation failed: {str(e)}")
    
    def translate_baidu(self, text: str, api_key: str, from_lang: str, to_lang: str) -> str:
        """Translate using Baidu Translate API."""
        app_id, secret_key = api_key.split(":")
        salt = str(int(time.time()))
        sign_str = app_id + text + salt + secret_key
        sign = hashlib.md5(sign_str.encode()).hexdigest()
        
        url = "https://fanyi-api.baidu.com/api/trans/vip/translate"
        
        params = {
            "q": text,
            "from": from_lang,
            "to": to_lang,
            "appid": app_id,
            "salt": salt,
            "sign": sign
        }
        
        try:
            response = requests.post(url, params=params, timeout=10)
            response.raise_for_status()
            result = response.json()
            
            if "error_code" in result:
                raise Exception(f"Baidu API error: {result['error_msg']}")
            
            return result["trans_result"][0]["dst"]
        except Exception as e:
            raise Exception(f"Baidu translation failed: {str(e)}")
    
    def translate_marianmt(self, text: str, direction: str) -> str:
        """Translate using MarianMT local models."""
        if not HAS_MARIANMT:
            raise Exception("MarianMT not available. Install with: pip install transformers torch")
        
        model_name = "Helsinki-NLP/opus-mt-jap-en" if direction == "日→英" else "Helsinki-NLP/opus-mt-en-jap"
        
        try:
            if model_name not in self.marianmt_models:
                self.marianmt_models[model_name] = {
                    "model": MarianMTModel.from_pretrained(model_name),
                    "tokenizer": MarianTokenizer.from_pretrained(model_name)
                }
            
            model = self.marianmt_models[model_name]["model"]
            tokenizer = self.marianmt_models[model_name]["tokenizer"]
            
            inputs = tokenizer.encode(text, return_tensors="pt", max_length=512, truncation=True)
            translated = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True)
            result = tokenizer.decode(translated[0], skip_special_tokens=True)
            
            return result
        except Exception as e:
            raise Exception(f"MarianMT translation failed: {str(e)}")
    
    def get_cache_key(self, text: str, provider: str, direction: str) -> str:
        """Generate cache key for translation."""
        return hashlib.md5(f"{text}_{provider}_{direction}".encode()).hexdigest()
    
    def translate_async(self, provider: str, text: str, api_key: str, direction: str) -> str:
        """Perform translation in background thread."""
        cache_key = self.get_cache_key(text, provider, direction)
        
        if cache_key in self.translation_cache:
            return self.translation_cache[cache_key]
        
        try:
            if provider == "DeepL":
                # DeepL uses specific language codes
                target_lang = "EN-US" if direction == "日→英" else "JA"
                translated = self.translate_deepl(text, api_key, target_lang)
                
            elif provider == "Google":
                target_lang = "en" if direction == "日→英" else "ja"
                translated = self.translate_google(text, api_key, target_lang)
                
            elif provider == "Yandex":
                lang_pair = "ja-en" if direction == "日→英" else "en-ja"
                translated = self.translate_yandex(text, api_key, lang_pair)
                
            elif provider == "Baidu":
                from_lang = "jp" if direction == "日→英" else "en"
                to_lang = "en" if direction == "日→英" else "jp"
                translated = self.translate_baidu(text, api_key, from_lang, to_lang)
                
            elif provider == "MarianMT":
                translated = self.translate_marianmt(text, direction)
                
            else:
                translated = f"Error: Unknown provider {provider}"
            
            self.translation_cache[cache_key] = translated
            return translated
            
        except Exception as e:
            error_msg = f"Translation failed: {str(e)}"
            self.translation_cache[cache_key] = error_msg
            return error_msg
    
    def translate(self, provider: str, input_text: str, translation_mode: str, other_language: str, api_key: str, execute_translation: bool, translated_text: str, input_text_connection: Optional[str] = None) -> Tuple[str, str, int]:
        """
        ComfyUI-focused translation: English ↔ Other Languages
        
        Args:
            provider: Translation provider
            input_text: Text to translate
            translation_mode: "English → Other Language" or "Other Language → English"
            other_language: The non-English language (Japanese, Chinese, etc.)
            api_key: API key for the provider
            execute_translation: Execute translation flag
            translated_text: Previous translated text (for UI display)
            input_text_connection: Optional connected input text
            
        Returns:
            Tuple of (original_text, translated_text, token_count)
        """
        text_to_translate = input_text_connection if input_text_connection else input_text
        
        if not text_to_translate.strip():
            return "", "", 0
        
        if not execute_translation:
            # No translation requested, return original text and any existing translation
            token_count = self.count_tokens(text_to_translate)
            if translated_text and translated_text.strip():
                translated_tokens = self.count_tokens(translated_text)
                token_count += translated_tokens
                return text_to_translate, translated_text, token_count
            return text_to_translate, "", token_count
        
        # Determine source and target languages from mode
        if translation_mode == "English → Other Language":
            source_language = "English"
            target_language = other_language
        else:  # "Other Language → English"
            source_language = other_language
            target_language = "English"
        
        # For Google, we can use free service, others need API key
        if provider not in ["MarianMT", "Google"] and not api_key.strip():
            cached_key = self.api_cache.get(provider, "")
            if not cached_key:
                help_text = {
                    "DeepL": "Get API key from https://www.deepl.com/pro-api",
                    "Yandex": "Get API key from https://translate.yandex.com/developers", 
                    "Baidu": "Get App ID and Secret from https://fanyi-api.baidu.com/ (format: app_id:secret_key)"
                }
                return text_to_translate, f"Error: {provider} API key required. {help_text.get(provider, '')}", 0
            api_key = cached_key
        elif provider == "Google" and not api_key.strip():
            # Use free Google Translate
            api_key = ""
        
        # Validate API key format for specific providers
        if provider == "DeepL" and api_key and not api_key.endswith(":fx") and not api_key.startswith("DeepL-Auth-Key"):
            if ":" in api_key:
                return text_to_translate, "Error: DeepL API key format incorrect. Use DeepL key only (not Baidu format)", 0
        elif provider == "Yandex" and api_key and ":" in api_key:
            return text_to_translate, "Error: Yandex requires standalone API key, not Baidu format", 0
        elif provider == "Baidu" and api_key and ":" not in api_key:
            return text_to_translate, "Error: Baidu requires format 'app_id:secret_key'", 0
        
        if api_key.strip() and provider != "MarianMT":
            self.api_cache[provider] = api_key
            self.save_cache()
        
        # Create translation direction for caching
        translation_direction = f"{source_language}→{target_language}"
        cache_key = self.get_cache_key(text_to_translate, provider, translation_direction)
        
        if cache_key in self.translation_cache:
            translated = self.translation_cache[cache_key]
        else:
            current_time = time.time()
            if current_time - self.last_translation_time < 0.5:
                time.sleep(0.1)
            
            self.last_translation_time = current_time
            
            try:
                # Get language codes for this provider
                source_code = self.get_language_code(source_language, provider)
                target_code = self.get_language_code(target_language, provider)
                
                # Call provider-specific translation
                if provider == "DeepL":
                    translated = self.translate_deepl(text_to_translate, api_key, target_code)
                elif provider == "Google":
                    translated = self.translate_google(text_to_translate, api_key, target_code)
                elif provider == "Yandex":
                    lang_pair = f"{source_code}-{target_code}"
                    translated = self.translate_yandex(text_to_translate, api_key, lang_pair)
                elif provider == "Baidu":
                    translated = self.translate_baidu(text_to_translate, api_key, source_code, target_code)
                elif provider == "MarianMT":
                    # MarianMT still uses old direction format for now
                    old_direction = "日→英" if source_language == "Japanese" else "英→日"
                    translated = self.translate_marianmt(text_to_translate, old_direction)
                else:
                    translated = f"Error: Unknown provider {provider}"
                
                self.translation_cache[cache_key] = translated
                
            except Exception as e:
                translated = f"Translation failed: {str(e)}"
        
        # Count tokens for original and translated text separately, then combine
        original_tokens = self.count_tokens(text_to_translate)
        translated_tokens = self.count_tokens(translated)
        total_tokens = original_tokens + translated_tokens
        
        print(f"Token count: Original={original_tokens}, Translated={translated_tokens}, Total={total_tokens}")
        
        # Update the translated_text field directly
        self.last_translated_text = translated
        
        return text_to_translate, translated, total_tokens


NODE_CLASS_MAPPINGS = {
    "PromptTranslator": PromptTranslator
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PromptTranslator": "Prompt Translator"
}