diff --git a/.gitignore b/.gitignore index 978e7ab..0d35b66 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ __pycache__/ *.py[cod] *$py.class *.so +*.whl # ComfyUI models/ @@ -18,3 +19,7 @@ input/ # OS .DS_Store Thumbs.db + +# Misc +.tracking +.gemini/ diff --git a/README.md b/README.md index 16e08a3..b19952c 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,16 @@ # ComfyUI-PromptTranslator-Enhancer -**Autor**: EzeTojo | **Versión**: 0.0.1 (Nightly) +**Autor**: EzeTojo | **Versión**: 1.0.0 Custom node para ComfyUI que traduce y mejora prompts desde distintos idiomas al inglés usando un modelo de lenguaje local (GGUF). Los idiomas soportados dependen exclusivamente del modelo que elijas usar. ## Características -- 🌐 **Multi**: Traduce automáticamente desde los idiomas soportados por el modelo al inglés +- 🌐 **Multi-idioma**: Traduce automáticamente desde los idiomas soportados por el modelo al inglés - ✨ **Mejora de prompts** con 3 niveles: basic, detailed, creative - 🖥️ **100% local** — usa modelos GGUF vía `llama-cpp-python` - 🔄 **Modelo reutilizable** — carga una vez, usa en múltiples nodos +- 🛡️ **Carga robusta** — fallback automático GPU → CPU con gestión inteligente de VRAM +- 🚫 **Filtro de modelos VL** — excluye automáticamente modelos Vision-Language no compatibles ## Ejemplos en acción @@ -36,7 +38,7 @@ Aquí tienes cómo se ve el nodo en ComfyUI en sus diferentes estados: ## Instalación -1. Clonar o copiar este directorio en `ComfyUI/custom_nodes/ComfyUI-PromptEnhancer/` +1. Clonar o copiar este directorio en `ComfyUI/custom_nodes/comfyui-prompt-translator-enhancer/` 2. Instalar dependencias: ```bash pip install llama-cpp-python @@ -49,13 +51,22 @@ Aquí tienes cómo se ve el nodo en ComfyUI en sus diferentes estados: ## Modelos recomendados -> **IMPORTANTE**: Recomendamos **fuertemente** usar modelos **Instruct** o **Chat** (que suelen tener "Instruct" o "Chat" en su nombre, ej. `Qwen3-VL-4B-Instruct-Q4.gguf`) en lugar de modelos base. Los modelos Instruct siguen nuestras reglas mucho mejor y no requieren del *fallback* de completado raw que es más lento. -> *Nota: Los idiomas de entrada que el nodo puede entender dependen de si el modelo GGUF fue entrenado en esos idiomas.* +> **IMPORTANTE**: Usa modelos **Instruct** o **Chat** (que suelen tener "Instruct" o "Chat" en su nombre). Los modelos Vision-Language (VL) como `Qwen3VL`, `LLaVA`, etc. **no son compatibles** con este nodo y serán filtrados automáticamente de la lista. - **Qwen3-4B-Instruct-Q4_K_M.gguf** (~2.5 GB) — Buen balance velocidad/calidad, excelente soporte multi-idioma. - **Qwen2.5-3B-Instruct-Q4_K_M.gguf** (~2 GB) — Más ligero. - **Phi-3-mini-4k-instruct-Q4_K_M.gguf** (~2.3 GB) — Alternativa sólida. +## Gestión de VRAM + +Este nodo está diseñado para funcionar dentro de ComfyUI, donde la GPU suele estar ocupada con modelos de difusión. La estrategia de carga es: + +1. **Intenta GPU** con el contexto solicitado, luego con contextos más pequeños (2048, 512) +2. **Libera cache VRAM** de PyTorch antes de cada intento (`torch.cuda.empty_cache()`) +3. **Fallback CPU puro** — si la GPU no tiene VRAM suficiente, carga el modelo 100% en CPU sin usar CUDA para compute, KQV offload, ni Flash Attention. Esto garantiza que funcione sin importar cuánta VRAM esté ocupada. + +> El modo CPU es más lento pero siempre funcional. Para mejor rendimiento, descarga otros modelos de la GPU antes de ejecutar el nodo. + ## Uso ### Opción 1: Todo-en-uno @@ -73,3 +84,17 @@ Aquí tienes cómo se ve el nodo en ComfyUI en sus diferentes estados: - **basic**: Traduce y añade tags mínimos de calidad - **detailed**: Traduce y añade iluminación, composición, calidad - **creative**: Traduce con interpretación artística, estilos, efectos + +## Changelog + +### v1.0.0 +- ✅ Carga robusta con fallback automático GPU → CPU +- ✅ Gestión inteligente de VRAM (libera cache de PyTorch) +- ✅ Modo CPU puro: `op_offload=False`, `flash_attn=disabled`, `offload_kqv=False` +- ✅ Filtro automático de modelos Vision-Language (VL) +- ✅ Registro proactivo de DLLs CUDA para Windows +- ✅ Verbose mode siempre activo para diagnóstico +- ✅ Soporte para prompts en cualquier idioma (depende del modelo) + +### v0.0.1 +- Versión inicial (nightly) diff --git a/__init__.py b/__init__.py index 8fe7ca3..0d182d5 100644 --- a/__init__.py +++ b/__init__.py @@ -5,7 +5,7 @@ Translates and enhances prompts from multiple languages to English using local G from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -__version__ = "0.0.1-nightly" +__version__ = "1.0.0" __author__ = "EzeTojo" __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "__version__", "__author__"] diff --git a/node.zip b/node.zip deleted file mode 100644 index bc53016..0000000 Binary files a/node.zip and /dev/null differ diff --git a/nodes.py b/nodes.py index 19f3c4f..6f2e407 100644 --- a/nodes.py +++ b/nodes.py @@ -1,16 +1,88 @@ """ ComfyUI-PromptTranslator-Enhancer -Custom nodes that enhance and translate prompts from Spanish to English +Custom nodes that enhance and translate prompts from any language to English using a local GGUF model via llama-cpp-python. """ import os +import sys import glob +import ctypes import folder_paths +# --------------------------------------------------------------------------- +# CUDA DLL Pre-loading (critical for running inside ComfyUI on Windows) +# --------------------------------------------------------------------------- +_cuda_dlls_loaded = False + + +def _ensure_cuda_dlls(): + """Pre-register DLL directories so llama.cpp can find CUDA runtime libs. + + Inside ComfyUI, PyTorch loads its own copy of CUDA DLLs from + ``site-packages/torch/lib``. llama-cpp-python ships its own + ``ggml-cuda.dll`` that depends on ``cudart64_*.dll`` and + ``cublas64_*.dll`` from the *same* CUDA toolkit. If those DLLs + are not on PATH or registered via ``os.add_dll_directory``, + ``llama_model_load_from_file`` silently returns NULL. + + This function registers every directory that might contain the + needed CUDA runtime DLLs. + """ + global _cuda_dlls_loaded + if _cuda_dlls_loaded: + return + _cuda_dlls_loaded = True + + if sys.platform != "win32": + return + + dll_dirs_to_register = set() + + # 1. llama_cpp's own bin/ and lib/ directories + try: + import llama_cpp as _lc + pkg_dir = os.path.dirname(_lc.__file__) + for subdir in ("bin", "lib"): + d = os.path.join(pkg_dir, subdir) + if os.path.isdir(d): + dll_dirs_to_register.add(d) + except ImportError: + pass + + # 2. PyTorch's CUDA DLLs (cudart, cublas) + try: + import torch + torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib") + if os.path.isdir(torch_lib): + dll_dirs_to_register.add(torch_lib) + except ImportError: + pass + + # 3. NVIDIA CUDA toolkit on PATH (if installed) + cuda_path = os.environ.get("CUDA_PATH", "") + if cuda_path: + cuda_bin = os.path.join(cuda_path, "bin") + if os.path.isdir(cuda_bin): + dll_dirs_to_register.add(cuda_bin) + + # Register all directories + for d in dll_dirs_to_register: + try: + os.add_dll_directory(d) + except (OSError, AttributeError): + pass + # Also prepend to PATH as legacy fallback + if d not in os.environ.get("PATH", ""): + os.environ["PATH"] = d + os.pathsep + os.environ.get("PATH", "") + + if dll_dirs_to_register: + print(f"[PromptTranslator-Enhancer] Registered {len(dll_dirs_to_register)} CUDA DLL directories") + def _get_llama_class(): """Lazy import of Llama to avoid errors when llama-cpp-python is not installed.""" + _ensure_cuda_dlls() try: from llama_cpp import Llama return Llama @@ -19,9 +91,124 @@ def _get_llama_class(): "llama-cpp-python is not installed. Please install it with:\n" " pip install llama-cpp-python\n" "For GPU (CUDA) support:\n" - " CMAKE_ARGS=\"-DGGML_CUDA=on\" pip install llama-cpp-python --force-reinstall --no-cache-dir" + ' CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python --force-reinstall --no-cache-dir' ) + +# Vision-Language model filename patterns (these need special multimodal handlers) +_VL_PATTERNS = ("qwen3vl", "qwen2vl", "qwen2.5vl", "llava", "minicpm-v", "internvl") + + +def _is_vl_model(model_path): + """Check if a GGUF file is a Vision-Language model (not suitable for text-only use).""" + name_lower = os.path.basename(model_path).lower() + # Also check parent folder name + parent_lower = os.path.basename(os.path.dirname(model_path)).lower() + combined = f"{parent_lower}/{name_lower}" + return any(p in combined for p in _VL_PATTERNS) + + +def _free_torch_vram(): + """Ask PyTorch to release cached VRAM so llama.cpp can use it.""" + try: + import torch + if torch.cuda.is_available(): + before = torch.cuda.memory_reserved(0) / 1024**2 + torch.cuda.empty_cache() + after = torch.cuda.memory_reserved(0) / 1024**2 + freed = before - after + if freed > 1: + print(f"[PromptTranslator-Enhancer] Freed {freed:.0f} MiB of PyTorch VRAM cache") + except Exception: + pass + + +def _load_model_with_fallback(Llama, model_path, n_gpu_layers, n_ctx): + """Try to load a GGUF model with an aggressive fallback strategy. + + Fallback order (tries freeing PyTorch VRAM cache before each): + 1. GPU + requested n_ctx (flash_attn=auto) + 2. GPU + 2048 (flash_attn=auto) + 3. GPU + 512 (flash_attn=auto) + 4. CPU + requested n_ctx (flash_attn=DISABLED, offload_kqv=False) + 5. CPU + 2048 (flash_attn=DISABLED, offload_kqv=False) + 6. CPU + 512 (flash_attn=DISABLED, offload_kqv=False) + + The CPU fallbacks explicitly disable Flash Attention and KQV offloading + so that llama.cpp does NOT try to allocate CUDA compute buffers. + + Returns the loaded Llama model instance. + Raises RuntimeError if all attempts fail. + """ + # Warn about Vision-Language models + if _is_vl_model(model_path): + print( + f"[PromptTranslator-Enhancer] ⚠ WARNING: '{os.path.basename(model_path)}' appears to be a " + f"Vision-Language model (VL). These models require a multimodal handler and are NOT " + f"supported by this text-only node. Please select a text/instruct model instead " + f"(e.g. Qwen3-4B-Q4_K_M.gguf)." + ) + raise RuntimeError( + f"Model '{os.path.basename(model_path)}' is a Vision-Language model and cannot be used " + f"with this text-only prompt enhancement node. Please choose a regular text model." + ) + + # Flash Attention type values (llama.cpp enum) + FLASH_ATTN_DISABLED = 0 # force off – no CUDA compute buffers + + # Build list of attempts: (label, kwargs_dict) + ctx_sizes = sorted(set([n_ctx, 2048, 512]), reverse=True) # dedupe & descending + + attempts = [] + # GPU attempts (flash_attn left at default/auto) + if n_gpu_layers != 0: + for ctx in ctx_sizes: + attempts.append((f"GPU/ctx={ctx}", dict( + n_gpu_layers=n_gpu_layers, n_ctx=ctx, + ))) + # CPU attempts – disable flash_attn, kqv offload, AND op_offload + # op_offload=False is critical: it prevents llama.cpp from allocating + # ~606 MiB CUDA compute buffers even when the model is on CPU + for ctx in ctx_sizes: + attempts.append((f"CPU-pure/ctx={ctx}", dict( + n_gpu_layers=0, n_ctx=ctx, + flash_attn_type=FLASH_ATTN_DISABLED, + offload_kqv=False, + op_offload=False, + ))) + + last_error = None + for label, kwargs in attempts: + try: + _free_torch_vram() + gpu_layers = kwargs.get("n_gpu_layers", 0) + ctx = kwargs.get("n_ctx", n_ctx) + print(f"[PromptTranslator-Enhancer] Trying {label} (n_gpu_layers={gpu_layers}) ...") + model = Llama( + model_path=model_path, + verbose=True, # ALWAYS verbose so native errors show in console + **kwargs, + ) + if ctx < n_ctx: + print(f"[PromptTranslator-Enhancer] ⚠ Context reduced from {n_ctx} to {ctx} to fit in memory") + print(f"[PromptTranslator-Enhancer] ✓ Model loaded successfully ({label})") + return model + except Exception as e: + last_error = e + print(f"[PromptTranslator-Enhancer] ✗ {label} failed: {e}") + + raise RuntimeError( + f"Failed to load model after {len(attempts)} attempts.\n" + f" Model: {model_path}\n" + f" Last error: {last_error}\n" + f" Troubleshooting:\n" + f" 1. Make sure the .gguf file is a TEXT model (not Vision-Language)\n" + f" 2. Make sure the .gguf file is valid and not corrupted\n" + f" 3. Try freeing GPU VRAM by unloading other models first\n" + f" 4. Check the console output above for native llama.cpp error messages" + ) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -33,15 +220,17 @@ if not os.path.exists(LLM_DIR): def get_available_gguf_models(): - """Recursively find all .gguf files under models/LLM/, excluding mmproj files.""" + """Recursively find all .gguf files under models/LLM/, excluding mmproj and VL files.""" pattern = os.path.join(LLM_DIR, "**", "*.gguf") all_files = glob.glob(pattern, recursive=True) - # Filter out vision projector files (mmproj-*) and return relative paths + # Filter out vision projector files (mmproj-*), VL models, and return relative paths models = [] for f in all_files: basename = os.path.basename(f) if basename.lower().startswith("mmproj-"): continue + if _is_vl_model(f): + continue rel_path = os.path.relpath(f, LLM_DIR) models.append(rel_path) return sorted(models) if models else ["no_models_found"] @@ -187,14 +376,8 @@ class LoadLLMModel: print(f"[PromptTranslator-Enhancer] GPU layers: {n_gpu_layers}, Context: {n_ctx}") Llama = _get_llama_class() - model = Llama( - model_path=model_path, - n_gpu_layers=n_gpu_layers, - n_ctx=n_ctx, - verbose=False, - ) + model = _load_model_with_fallback(Llama, model_path, n_gpu_layers, n_ctx) - print(f"[PromptTranslator-Enhancer] Model loaded successfully!") return (model,) @@ -203,7 +386,7 @@ class LoadLLMModel: # --------------------------------------------------------------------------- class PromptEnhancerTranslator: - """Takes a Spanish prompt, enhances it for image generation, and translates to English. + """Takes a prompt in any language, enhances it for image generation, and translates to English. This node loads the model on each execution. For repeated use, prefer LoadLLMModel + PromptEnhancerFromModel.""" @classmethod @@ -211,7 +394,7 @@ class PromptEnhancerTranslator: return { "required": { "prompt": ("STRING", {"multiline": True, "default": "", - "placeholder": "Escribe tu prompt en español..."}), + "placeholder": "Escribe tu prompt en cualquier idioma..."}), "llm_model": (get_available_gguf_models(), {"default": get_available_gguf_models()[0]}), "enhancement_level": (["basic", "detailed", "creative"], {"default": "detailed"}), "n_gpu_layers": ("INT", {"default": -1, "min": -1, "max": 100, "step": 1}), @@ -225,7 +408,7 @@ class PromptEnhancerTranslator: RETURN_NAMES = ("enhanced_prompt", "original_prompt",) FUNCTION = "enhance" CATEGORY = "LLM/Prompt Enhancement" - DESCRIPTION = "Enhances and translates a Spanish prompt to English using a local GGUF model." + DESCRIPTION = "Enhances and translates a prompt to English using a local GGUF model." def enhance(self, prompt, llm_model, enhancement_level, n_gpu_layers, max_tokens, temperature, seed): if not prompt.strip(): @@ -237,12 +420,7 @@ class PromptEnhancerTranslator: print(f"[PromptTranslator-Enhancer] Loading model: {llm_model}") Llama = _get_llama_class() - model = Llama( - model_path=model_path, - n_gpu_layers=n_gpu_layers, - n_ctx=4096, - verbose=False, - ) + model = _load_model_with_fallback(Llama, model_path, n_gpu_layers, n_ctx=4096) enhanced = run_llm_enhancement(model, prompt, enhancement_level, max_tokens, temperature, seed) @@ -260,7 +438,7 @@ class PromptEnhancerTranslator: # --------------------------------------------------------------------------- class PromptEnhancerFromModel: - """Takes a Spanish prompt and a pre-loaded LLM model, enhances the prompt and translates to English.""" + """Takes a prompt in any language and a pre-loaded LLM model, enhances the prompt and translates to English.""" @classmethod def INPUT_TYPES(cls): @@ -268,7 +446,7 @@ class PromptEnhancerFromModel: "required": { "llm_model": ("LLM_MODEL",), "prompt": ("STRING", {"multiline": True, "default": "", - "placeholder": "Escribe tu prompt en español..."}), + "placeholder": "Escribe tu prompt en cualquier idioma..."}), "enhancement_level": (["basic", "detailed", "creative"], {"default": "detailed"}), "max_tokens": ("INT", {"default": 256, "min": 32, "max": 1024, "step": 32}), "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.05}), @@ -280,7 +458,7 @@ class PromptEnhancerFromModel: RETURN_NAMES = ("enhanced_prompt", "original_prompt",) FUNCTION = "enhance" CATEGORY = "LLM/Prompt Enhancement" - DESCRIPTION = "Enhances and translates a Spanish prompt using a pre-loaded LLM model from LoadLLMModel node." + DESCRIPTION = "Enhances and translates a prompt using a pre-loaded LLM model from LoadLLMModel node." def enhance(self, llm_model, prompt, enhancement_level, max_tokens, temperature, seed): if not prompt.strip(): diff --git a/pyproject.toml b/pyproject.toml index 73aaa19..0513229 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-prompt-translator-enhancer" description = "Translates and enhances prompts from multiple languages to English using local GGUF models." -version = "0.0.1" +version = "1.0.0" license = { file = "LICENSE" } dependencies = ["llama-cpp-python>=0.2.0"] @@ -11,4 +11,4 @@ Repository = "https://github.com/ezetojo/ComfyUI-PromptTranslatorEnhancer" [tool.comfy] PublisherId = "ezetojo" DisplayName = "Prompt Translator & Enhancer" -Icon = "" +Icon = "https://raw.githubusercontent.com/ezetojo/ComfyUI-PromptTranslatorEnhancer/refs/heads/master/assets/prompt_translator_enhancer.png" \ No newline at end of file