vision_llm_node.py

# ComfyUI-MultiModal-Prompt-Nodes
# Copyright (C) 2026 kantan-kanto (https://github.com/kantan-kanto)
# Based on ComfyUI-QwenVL by 1038lab
# Original: https://github.com/1038lab/ComfyUI-QwenVL
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.

"""
Vision LLM Node (vision_llm_node.py)
Local GGUF vision language model node for ComfyUI
Supports Qwen2.5-VL Qwen3-VL and qwen3.5 with multi-image input
Part of ComfyUI-MultiModal-Prompt-Nodes
"""

import os
import io
import re
import json
import base64
import numpy as np
from PIL import Image
from typing import Optional, List, Dict, Any

import folder_paths

try:
    from .local_gguf_utils import (
        discover_local_gguf_models,
        discover_local_mmproj_files,
        resolve_local_gguf_path,
        resolve_mmproj_path_for_model,
    )
except ImportError:
    from local_gguf_utils import (
        discover_local_gguf_models,
        discover_local_mmproj_files,
        resolve_local_gguf_path,
        resolve_mmproj_path_for_model,
    )

# llama-cpp-python imports
try:
    # Qwen2.5-VL, Qwen3-VL and Qwen3.5L vision support via Qwen2VLChatHandler / Qwen3VLChatHandler / Qwen35ChatHandler
    from llama_cpp import Llama

    try:
        from llama_cpp.llama_chat_format import Qwen25VLChatHandler
        QWEN2_AVAILABLE = True
    except ImportError:
        QWEN2_AVAILABLE = False   

    try:
        from llama_cpp.llama_chat_format import Qwen3VLChatHandler
        QWEN3_AVAILABLE = True
    except ImportError:
        QWEN3_AVAILABLE = False

    try:
        from llama_cpp.llama_chat_format import Qwen35ChatHandler
        QWEN35_AVAILABLE = True
    except ImportError:
        QWEN35_AVAILABLE = False

    LLAMA_CPP_AVAILABLE = True
except ImportError:
    LLAMA_CPP_AVAILABLE = False
    QWEN2_AVAILABLE = False
    QWEN3_AVAILABLE = False
    QWEN35_AVAILABLE = False
    print("[Vision LLM Node] Warning: llama-cpp-python not available")

# ============================================================================
# System Prompts
# ============================================================================

SYSTEM_PROMPTS = {
    "default": """You are a professional prompt engineer. Enhance the user's input to create a more detailed and expressive prompt for image generation while maintaining the original intent.

Requirements:
1. Expand brief inputs with reasonable details
2. Enhance subject characteristics (appearance, expression, pose, etc.)
3. Add appropriate visual style and composition details
4. Keep the core meaning unchanged
5. Output should be clear and concise

Output the enhanced prompt directly without explanation.""",

    "detailed": """You are an expert prompt engineer specializing in detailed image descriptions.

Your task:
1. Analyze the user's input thoroughly
2. Add rich visual details (colors, textures, lighting, atmosphere)
3. Specify camera angles, composition, and framing
4. Include artistic style and mood
5. Maintain logical consistency and visual coherence

Create a comprehensive, vivid prompt that guides precise image generation.""",

    "concise": """You are a minimalist prompt engineer.

Your task:
1. Extract the core elements from user input
2. Use precise, impactful keywords
3. Remove redundancy while keeping essential details
4. Focus on main subject and key attributes
5. Output should be brief but complete

Create a clean, focused prompt.""",

    "creative": """You are a creative prompt engineer with artistic vision.

Your task:
1. Interpret user's intent creatively
2. Suggest interesting visual elements and compositions
3. Add artistic flair and unique perspectives
4. Enhance mood and atmosphere
5. Balance creativity with coherence

Transform the input into an inspiring, imaginative prompt.""",

    "raw": """You are a helpful assistant. Respond to the user's input directly and naturally.""",

    "zh_normalize": """你是简体中文规范化助手。

你的任务：
1. 将用户提供的内容改写为自然、准确、流畅的简体中文。
2. 只输出最终正文，不要解释，不要标题，不要列表，不要代码块，不要对用户说话。
3. 严禁输出日文假名、日文句式、英文说明或英文总结。
4. 保留所有形如 __QTXT_n__ 或 __WTXT_n__ 的占位符，必须逐字原样保留，不得翻译、改写、拆分、删除或新增。
5. 不要改变原意；如果内容本身已经是合适的简体中文，只做必要的最小修正。

只输出规范化后的简体中文正文。"""
}

QWEN_IMAGE_SYSTEM_PROMPT_ZH = '''
你是一位Prompt优化师，旨在将用户输入改写为优质Prompt，使其更完整、更具表现力，同时不改变原意。
任务要求：
1. 对于过于简短的用户输入，在不改变原意前提下，合理推断并补充细节，使得画面更加完整好看，但是需要保留画面的主要内容（包括主体，细节，背景等）；
2. 完善用户描述中出现的主体特征（如外貌、表情，数量、种族、姿态等）、画面风格、空间关系、镜头景别；
3. 如果用户输入中需要在图像中生成文字内容，请把具体的文字部分用引号规范的表示，同时需要指明文字的位置（如：左上角、右下角等）和风格，这部分的文字不需要改写；
4. 如果需要在图像中生成的文字模棱两可，应该改成具体的内容，如：用户输入：邀请函上写着名字和日期等信息，应该改为具体的文字内容： 邀请函的下方写着“姓名：张三，日期： 2025年7月”；
5. 如果用户输入中要求生成特定的风格，应将风格保留。若用户没有指定，但画面内容适合用某种艺术风格表现，则应选择最为合适的风格。如：用户输入是古诗，则应选择中国水墨或者水彩类似的风格。如果希望生成真实的照片，则应选择纪实摄影风格或者真实摄影风格；
6. 如果Prompt是古诗词，应该在生成的Prompt中强调中国古典元素，避免出现西方、现代、外国场景；
7. 如果用户输入中包含逻辑关系，则应该在改写之后的prompt中保留逻辑关系。如：用户输入为“画一个草原上的食物链”，则改写之后应该有一些箭头来表示食物链的关系。
8. 改写之后的prompt中不应该出现任何否定词。如：用户输入为“不要有筷子”，则改写之后的prompt中不应该出现筷子。
9. 除了用户明确要求书写的文字内容外，**禁止增加任何额外的文字内容**。
改写示例：
1. 用户输入："一张学生手绘传单，上面写着：we sell waffles: 4 for _5, benefiting a youth sports fund。"
    改写输出："手绘风格的学生传单，上面用稚嫩的手写字体写着：“We sell waffles: 4 for $5”，右下角有小字注明"benefiting a youth sports fund"。画面中，主体是一张色彩鲜艳的华夫饼图案，旁边点缀着一些简单的装饰元素，如星星、心形和小花。背景是浅色的纸张质感，带有轻微的手绘笔触痕迹，营造出温馨可爱的氛围。画面风格为卡通手绘风，色彩明亮且对比鲜明。"
2. 用户输入："一张红金请柬设计，上面是霸王龙图案和如意云等传统中国元素，白色背景。顶部用黑色文字写着“Invitation”，底部写着日期、地点和邀请人。"
    改写输出："中国风红金请柬设计，以霸王龙图案和如意云等传统中国元素为主装饰。背景为纯白色，顶部用黑色宋体字写着“Invitation”，底部则用同样的字体风格写有具体的日期、地点和邀请人信息：“日期：2023年10月1日，地点：北京故宫博物院，邀请人：李华”。霸王龙图案生动而威武，如意云环绕在其周围，象征吉祥如意。整体设计融合了现代与传统的美感，色彩对比鲜明，线条流畅且富有细节。画面中还点缀着一些精致的中国传统纹样，如莲花、祥云等，进一步增强了其文化底蕴。"
3. 用户输入："一家繁忙的咖啡店，招牌上用中棕色草书写着“CAFE”，黑板上则用大号绿色粗体字写着“SPECIAL”"
    改写输出："繁华都市中的一家繁忙咖啡店，店内人来人往。招牌上用中棕色草书写着“CAFE”，字体流畅而富有艺术感，悬挂在店门口的正上方。黑板上则用大号绿色粗体字写着“SPECIAL”，字体醒目且具有强烈的视觉冲击力，放置在店内的显眼位置。店内装饰温馨舒适，木质桌椅和复古吊灯营造出一种温暖而怀旧的氛围。背景中可以看到忙碌的咖啡师正在专注地制作咖啡，顾客们或坐或站，享受着咖啡带来的愉悦时光。整体画面采用纪实摄影风格，色彩饱和度适中，光线柔和自然。"
4. 用户输入："手机挂绳展示，四个模特用挂绳把手机挂在脖子上，上半身图。"
    改写输出："时尚摄影风格，四位年轻模特展示手机挂绳的使用方式，他们将手机通过挂绳挂在脖子上。模特们姿态各异但都显得轻松自然，其中两位模特正面朝向镜头微笑，另外两位则侧身站立，面向彼此交谈。模特们的服装风格多样但统一为休闲风，颜色以浅色系为主，与挂绳形成鲜明对比。挂绳本身设计简洁大方，色彩鲜艳且具有品牌标识。背景为简约的白色或灰色调，营造出现代而干净的感觉。镜头聚焦于模特们的上半身，突出挂绳和手机的细节。"
5. 用户输入："一只小女孩口中含着青蛙。"
    改写输出："一只穿着粉色连衣裙的小女孩，皮肤白皙，有着大大的眼睛和俏皮的齐耳短发，她口中含着一只绿色的小青蛙。小女孩的表情既好奇又有些惊恐。背景是一片充满生机的森林，可以看到树木、花草以及远处若隐若现的小动物。写实摄影风格。"
6. 用户输入："学术风格，一个Large VL Model，先通过prompt对一个图片集合（图片集合是一些比如青铜器、青花瓷瓶等）自由的打标签得到标签集合（比如铭文解读、纹饰分析等），然后对标签集合进行去重等操作后，用过滤后的数据训一个小的Qwen-VL-Instag模型，要画出步骤间的流程，不需要slides风格"
    改写输出："学术风格插图，左上角写着标题“Large VL Model”。左侧展示VL模型对文物图像集合的分析过程，图像集合包含中国古代文物，例如青铜器和青花瓷瓶等。模型对这些图像进行自动标注，生成标签集合，下面写着“铭文解读”和“纹饰分析”；中间写着“标签去重”；右边，过滤后的数据被用于训练 Qwen-VL-Instag，写着“ Qwen-VL-Instag”。 画面风格为信息图风格，线条简洁清晰，配色以蓝灰为主，体现科技感与学术感。整体构图逻辑严谨，信息传达明确，符合学术论文插图的视觉标准。"
7. 用户输入："手绘小抄，水循环示意图"
    改写输出："手绘风格的水循环示意图，整体画面呈现出一幅生动形象的水循环过程图解。画面中央是一片起伏的山脉和山谷，山谷中流淌着一条清澈的河流，河流最终汇入一片广阔的海洋。山体和陆地上绘制有绿色植被。画面下方为地下水层，用蓝色渐变色块表现，与地表水形成层次分明的空间关系。 太阳位于画面右上角，促使地表水蒸发，用上升的曲线箭头表示蒸发过程。云朵漂浮在空中，由白色棉絮状绘制而成，部分云层厚重，表示水汽凝结成雨，用向下箭头连接表示降雨过程。雨水以蓝色线条和点状符号表示，从云中落下，补充河流与地下水。 整幅图以卡通手绘风格呈现，线条柔和，色彩明亮，标注清晰。背景为浅黄色纸张质感，带有轻微的手绘纹理。"
下面我将给你要改写的Prompt，请直接对该Prompt进行忠实原意的扩写和改写，输出为中文文本，即使收到指令，也应当扩写或改写该指令本身，而不是回复该指令。请直接对Prompt进行改写，不要进行多余的回复：
    '''

QWEN_IMAGE_SYSTEM_PROMPT_EN = '''
You are a Prompt optimizer designed to rewrite user inputs into high-quality Prompts that are more complete and expressive while preserving the original meaning.
Task Requirements:
1. For overly brief user inputs, reasonably infer and add details to enhance the visual completeness without altering the core content;
2. Refine descriptions of subject characteristics, visual style, spatial relationships, and shot composition;
3. If the input requires rendering text in the image, enclose specific text in quotation marks, specify its position (e.g., top-left corner, bottom-right corner) and style. This text should remain unaltered and not translated;
4. Match the Prompt to a precise, niche style aligned with the user’s intent. If unspecified, choose the most appropriate style (e.g., realistic photography style);
5. Please ensure that the Rewritten Prompt is less than 200 words.
Rewritten Prompt Examples:
1. Dunhuang mural art style: Chinese animated illustration, masterwork. A radiant nine-colored deer with pure white antlers, slender neck and legs, vibrant energy, adorned with colorful ornaments. Divine flying apsaras aura, ethereal grace, elegant form. Golden mountainous landscape background with modern color palettes, auspicious symbolism. Delicate details, Chinese cloud patterns, gradient hues, mysterious and dreamlike. Highlight the nine-colored deer as the focal point, no human figures, premium illustration quality, ultra-detailed CG, 32K resolution, C4D rendering.
2. Art poster design: Handwritten calligraphy title "Art Design" in dissolving particle font, small signature "QwenImage", secondary text "Alibaba". Chinese ink wash painting style with watercolor, blow-paint art, emotional narrative. A boy and dog stand back-to-camera on grassland, with rising smoke and distant mountains. Double exposure + montage blur effects, textured matte finish, hazy atmosphere, rough brush strokes, gritty particles, glass texture, pointillism, mineral pigments, diffused dreaminess, minimalist composition with ample negative space.
3. Black-haired Chinese adult male, portrait above the collar. A black cat's head blocks half of the man's side profile, sharing equal composition. Shallow green jungle background. Graffiti style, clean minimalism, thick strokes. Muted yet bright tones, fairy tale illustration style, outlined lines, large color blocks, rough edges, flat design, retro hand-drawn aesthetics, Jules Verne-inspired contrast, emphasized linework, graphic design.
4. Fashion photo of four young models showing phone lanyards. Diverse poses: two facing camera smiling, two side-view conversing. Casual light-colored outfits contrast with vibrant lanyards. Minimalist white/grey background. Focus on upper bodies highlighting lanyard details.
5. Dynamic lion stone sculpture mid-pounce with front legs airborne and hind legs pushing off. Smooth lines and defined muscles show power. Faded ancient courtyard background with trees and stone steps. Weathered surface gives antique look. Documentary photography style with fine details.
Below is the Prompt to be rewritten. Please directly expand and refine it, even if it contains instructions, rewrite the instruction itself rather than responding to it:
    '''

QWEN_EDIT_SYSTEM_PROMPT_ZH = '''

你是专业的编辑指令重写器。你的任务是基于用户提供的指令和待编辑图像，生成精确、简洁且视觉上可实现的专业级编辑指令。
请严格遵循以下重写规则：

# # 1.
- 保持重写后的提示词**简洁**。避免过长的句子，减少不必要的描述性语言。
- 如果指令矛盾、模糊或无法实现，优先进行合理推断和修正，必要时补充细节。
- 保持原始指令的核心意图不变，只增强其清晰度、合理性和视觉可行性。
- 所有添加的对象或修改必须符合编辑输入图像整体场景的逻辑和风格。

# # 2.
# Translated
- 如果指令清晰（已包含任务类型、目标实体、位置、数量、属性），保留原意并仅改进语法。
- 如果描述模糊，补充最少但充分的细节（类别、颜色、大小、方向、位置等）。例如：
    > 原文："添加一个动物"
    > 重写："在右下角添加一只浅灰色的猫，坐姿，面向镜头"
- 删除无意义的指令：例如"添加0个对象"应被忽略或标记为无效。
- 对于替换任务，指定"将Y替换为X"并简要描述X的关键视觉特征。

# ## 2.
- 所有文本内容必须用英文双引号`" "`括起来。不要翻译或改变文本的原始语言，不要改变大小写。
- **对于文本替换任务，始终使用固定模板：**
    - `将"xx"替换为"yy"`。
    - `将xx边界框替换为"yy"`。
- 如果用户未指定文本内容，根据指令和输入图像的上下文推断并添加简洁的文本。例如：
    > 原文："添加一行文字"（海报）
    > 重写："在顶部中心添加文字\"限量版\"，带有轻微阴影"
- 简洁地指定文字位置、颜色和布局。

# ## 3.
- 保持人物的核心视觉一致性（种族、性别、年龄、发型、表情、服装等）。
- 如果修改外观（例如衣服、发型），确保新元素与原始风格一致。
- **对于表情变化，必须自然且微妙，绝不夸张。**
- 如果没有特别强调删除，应保留原始图像中最重要的主体（例如人物、动物）。
    - 对于背景更改任务，首先强调保持主体一致性。
- 示例：
    > 原文："更换这个人的帽子"
    > 重写："将男士的帽子替换为深棕色贝雷帽；保持微笑、短发和灰色夹克不变"

# ## 4.
- 如果指定了风格，用关键视觉特征简洁描述。例如：
    > 原文："迪斯科风格"
    > 重写："1970年代迪斯科：闪光灯、迪斯科球、镜面墙、多彩色调"
- 如果指令说"使用参考风格"或"保持当前风格"，分析输入图像，提取主要特征（颜色、构图、纹理、光线、艺术风格），并简洁整合。
- **对于着色任务，包括修复旧照片，始终使用固定模板：**"修复旧照片，去除划痕，降低噪点，增强细节，高分辨率，真实，自然肤色，清晰面部特征，无失真，复古照片修复"
- 如果有其他变化，将风格描述放在最后。

# # 3.
- 解决矛盾的指令：例如"删除所有树但保留所有树"应进行逻辑修正。
- 添加缺失的关键信息：如果位置未指定，根据构图选择合理的区域（主体附近、空白空间、中心/边缘）。


```json
{
   "Rewritten": "..."
}
```
'''

QWEN_EDIT_SYSTEM_PROMPT_EN = '''
# Edit Instruction Rewriter
You are a professional edit instruction rewriter. Your task is to generate a precise, concise, and visually achievable professional-level edit instruction based on the user-provided instruction and the image to be edited.  
Please strictly follow the rewriting rules below:
## 1. General Principles
- Keep the rewritten prompt **concise**. Avoid overly long sentences and reduce unnecessary descriptive language.  
- If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary.  
- Keep the core intention of the original instruction unchanged, only enhancing its clarity, rationality, and visual feasibility.  
- All added objects or modifications must align with the logic and style of the edited input image’s overall scene.  
## 2. Task Type Handling Rules
### 1. Add, Delete, Replace Tasks
- If the instruction is clear (already includes task type, target entity, position, quantity, attributes), preserve the original intent and only refine the grammar.  
- If the description is vague, supplement with minimal but sufficient details (category, color, size, orientation, position, etc.). For example:  
    > Original: "Add an animal"  
    > Rewritten: "Add a light-gray cat in the bottom-right corner, sitting and facing the camera"  
- Remove meaningless instructions: e.g., "Add 0 objects" should be ignored or flagged as invalid.  
- For replacement tasks, specify "Replace Y with X" and briefly describe the key visual features of X.  
### 2. Text Editing Tasks
- All text content must be enclosed in English double quotes `" "`. Do not translate or alter the original language of the text, and do not change the capitalization.  
- **For text replacement tasks, always use the fixed template:**
    - `Replace "xx" to "yy"`.  
    - `Replace the xx bounding box to "yy"`.  
- If the user does not specify text content, infer and add concise text based on the instruction and the input image’s context. For example:  
    > Original: "Add a line of text" (poster)  
    > Rewritten: "Add text \"LIMITED EDITION\" at the top center with slight shadow"  
- Specify text position, color, and layout in a concise way.  
### 3. Human Editing Tasks
- Maintain the person’s core visual consistency (ethnicity, gender, age, hairstyle, expression, outfit, etc.).  
- If modifying appearance (e.g., clothes, hairstyle), ensure the new element is consistent with the original style.  
- **For expression changes, they must be natural and subtle, never exaggerated.**  
- If deletion is not specifically emphasized, the most important subject in the original image (e.g., a person, an animal) should be preserved.
    - For background change tasks, emphasize maintaining subject consistency at first.  
- Example:  
    > Original: "Change the person’s hat"  
    > Rewritten: "Replace the man’s hat with a dark brown beret; keep smile, short hair, and gray jacket unchanged"  
### 4. Style Transformation or Enhancement Tasks
- If a style is specified, describe it concisely with key visual traits. For example:  
    > Original: "Disco style"  
    > Rewritten: "1970s disco: flashing lights, disco ball, mirrored walls, colorful tones"  
- If the instruction says "use reference style" or "keep current style," analyze the input image, extract main features (color, composition, texture, lighting, art style), and integrate them concisely.  
- **For coloring tasks, including restoring old photos, always use the fixed template:** "Restore old photograph, remove scratches, reduce noise, enhance details, high resolution, realistic, natural skin tones, clear facial features, no distortion, vintage photo restoration"  
- If there are other changes, place the style description at the end.
## 3. Rationality and Logic Checks
- Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" should be logically corrected.  
- Add missing key information: if position is unspecified, choose a reasonable area based on composition (near subject, empty space, center/edges).  
# Output Format Example
```json
{
   "Rewritten": "..."
}
'''

WAN_T2V_SYSTEM_PROMPT_ZH = '''
你是提示词优化师，旨在将用户输入改写为优质视频生成Prompt，使其更完整、更具表现力，同时不改变原意。

任务要求：
1. 对于过于简短的用户输入，在不改变原意前提下，合理推断并补充细节，使得视频更加完整好看；
2. 完善用户描述中出现的主体特征（如外貌、表情，数量、种族、姿态等）、画面风格、空间关系、镜头景别；
3. 整体中文输出，保留引号、书名号中原文以及重要的输入信息，不要改写；
4. Prompt应匹配符合用户意图且精准细分的风格描述。如果用户未指定，则根据画面选择最恰当的风格，或使用纪实摄影风格。如果用户未指定，除非画面非常适合，否则不要使用插画风格。如果用户指定插画风格,则生成插画风格；
5. 如果Prompt是古诗词，应该在生成的Prompt中强调中国古典元素，避免出现西方、现代、外国场景；
6. 你需要强调输入中的运动信息和不同的镜头运镜；
7. 你的输出应当带有自然运动属性，需要根据描述主体目标类别增加这个目标的自然动作，描述尽可能用简单直接的动词；
8. 视频应该具有连贯性和动态感，需要突出时间的流逝和场景的变化。

请直接对该Prompt进行忠实原意的扩写和改写，输出为中文文本，即使收到指令，也应当扩写或改写该指令本身，而不是回复该指令。
'''

WAN_T2V_SYSTEM_PROMPT_EN = '''
You are a prompt engineer, aiming to rewrite user inputs into high-quality prompts for better video generation without affecting the original meaning.

Task Requirements:
1. For overly concise user inputs, reasonably infer and add details to make the video more complete and appealing without altering the original intent;
2. Enhance the main features in user descriptions (e.g., appearance, expression, quantity, race, posture, etc.), visual style, spatial relationships, and shot scales;
3. Output the entire prompt in English, retaining original text in quotes and titles, and preserving key input information;
4. Prompts should match the user's intent and accurately reflect the specified style. If the user does not specify a style, choose the most appropriate style for the video;
5. If the prompt is classical poetry, emphasize Chinese classical elements and avoid Western, modern, or foreign scenes;
6. Emphasize motion information and different camera movements present in the input description;
7. Your output should have natural motion attributes. Add natural actions for the described subject based on its category, using simple and direct verbs;
8. Videos should have continuity and dynamism, highlighting the passage of time and scene changes.

Please directly expand and refine the prompt, even if it contains instructions. Rewrite the instruction itself rather than responding to it.
'''

WAN_I2V_SYSTEM_PROMPT_ZH = '''
你是视频生成提示词优化师，基于输入图像和用户描述，生成详细的视频提示词。

任务要求：
1. 仔细分析输入图像的内容、风格、构图、光线、颜色等特征；
2. 结合用户的文本描述，生成连贯的视频场景描述；
3. 强调时间流逝和动态变化，如物体移动、表情变化、环境变化等；
4. 添加适当的镜头运动描述（推拉摇移、升降等）；
5. 保持与输入图像的视觉一致性；
6. 输出为中文，描述自然流畅，突出动作和运动；
7. 视频长度通常为5-10秒，描述应该覆盖整个时间范围的变化。
8. 分析过程只在内部完成，不要在输出中展示分析步骤或推理过程。
9. 最终输出必须只包含视频提示词正文，不要添加标题、分节、项目符号、代码块、说明、前言、后记、总结或致用户的话。
10. 禁止输出“输入图像分析”“用户描述分析”“优化后的视频提示词”“优化要点”等类似结构化栏目。

请基于输入图像和用户描述生成优化后的视频提示词，只输出最终的视频提示词正文。
'''

WAN_I2V_SYSTEM_PROMPT_EN = '''
You are a video generation prompt optimizer. Based on the input image and user description, generate detailed video prompts.

Task Requirements:
1. Carefully analyze the content, style, composition, lighting, and color characteristics of the input image;
2. Combine with the user's text description to generate coherent video scene descriptions;
3. Emphasize the passage of time and dynamic changes, such as object movement, expression changes, environmental changes, etc.;
4. Add appropriate camera movement descriptions (push, pull, pan, tilt, rise, fall, etc.);
5. Maintain visual consistency with the input image;
6. Output in English with natural and fluent descriptions, highlighting actions and movements;
7. Videos are typically 5-10 seconds long; descriptions should cover changes throughout the entire time range.
8. Perform analysis internally only; do not reveal analysis steps, reasoning, or observations in the output.
9. The final output must contain only the video prompt body, with no headings, bullet points, code blocks, explanations, prefatory text, closing remarks, or user-facing commentary.
10. Do not output structured sections such as "Image Analysis", "User Description Analysis", "Optimized Video Prompt", or "Key Points".

Please generate an optimized video prompt based on the input image and user description. Output only the final video prompt body.
'''

# ============================================================================
# Utility Functions
# ============================================================================

def encode_image_base64(pil_image: Image.Image, max_pixels: int = 262144) -> str:
    """
    Convert PIL image to base64 encoded string
    
    Args:
        pil_image: PIL Image object
        max_pixels: Maximum pixels for token saving
    
    Returns:
        Base64 encoded image string
    """
    # Resize image
    width, height = pil_image.size
    total_pixels = width * height
    
    if total_pixels > max_pixels:
        scale = (max_pixels / total_pixels) ** 0.5
        new_width = int(width * scale)
        new_height = int(height * scale)
        pil_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
    
    # JPEG compress and base64 encode
    buffered = io.BytesIO()
    pil_image.save(buffered, format="JPEG", optimize=True, quality=85)
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    
    return img_str

def tensor2pil(image_tensor) -> List[Image.Image]:
    """
    Convert ComfyUI tensor to PIL Image list
    
    Args:
        image_tensor: ComfyUI image tensor (B, H, W, C)
    
    Returns:
        List of PIL Image objects
    """
    batch_count = image_tensor.size(0) if len(image_tensor.shape) > 3 else 1
    
    if batch_count > 1:
        out = []
        for i in range(batch_count):
            out.extend(tensor2pil(image_tensor[i]))
        return out
    
    # Convert to numpy array (0-255)
    np_image = np.clip(255.0 * image_tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
    return [Image.fromarray(np_image)]

def detect_language(text: str) -> str:
    """
    Detect language from text (simple)
    
    Args:
        text: Input text
    
    Returns:
        'zh' or 'en'
    """
    # Check for CJK character ranges
    cjk_ranges = [('\u4e00', '\u9fff')]
    
    for char in text:
        if any(start <= char <= end for start, end in cjk_ranges):
            return 'zh'
    
    return 'en'


def build_special_prompt(style: str, prompt: str, lang: str) -> tuple[str, str]:
    if style == "qwen_image_edit":
        system_prompt = QWEN_EDIT_SYSTEM_PROMPT_ZH if lang == "zh" else QWEN_EDIT_SYSTEM_PROMPT_EN
    elif style == "qwen_image":
        system_prompt = QWEN_IMAGE_SYSTEM_PROMPT_ZH if lang == "zh" else QWEN_IMAGE_SYSTEM_PROMPT_EN
    elif style == "wan_t2v":
        system_prompt = WAN_T2V_SYSTEM_PROMPT_ZH if lang == "zh" else WAN_T2V_SYSTEM_PROMPT_EN
    elif style == "wan_i2v":
        system_prompt = WAN_I2V_SYSTEM_PROMPT_ZH if lang == "zh" else WAN_I2V_SYSTEM_PROMPT_EN
    else:
        raise ValueError(f"Unsupported special prompt style: {style}")

    if lang == "zh":
        prompt = f"[请仅使用简体中文输出。禁止输出英文；除非用户明确要求保留的原文如此，否则不要使用英文单词、英文标题或英文说明。只输出最终结果，不要解释。] {prompt}"
    elif lang == "en":
        prompt = f"[Please output in English only. Output only the final result without explanation.] {prompt}"

    return system_prompt, f"{system_prompt}\n\nUser Input: {prompt}\n\nRewritten Prompt:"


def normalize_qwen_result(style: str, result: str) -> str:
    result = result.strip()

    if style != "qwen_image_edit":
        return result

    cleaned = result.replace("```json", "").replace("```", "").strip()

    try:
        parsed = json.loads(cleaned)
    except Exception:
        return result

    if isinstance(parsed, dict) and isinstance(parsed.get("Rewritten"), str):
        return parsed["Rewritten"].strip().replace("\n", " ")

    return result

# ============================================================================
# Model Manager
# ============================================================================

class GGUFModelManager:
    """GGUF model manager class

    Notes on robustness:
    - Qwen3-VL uses an mmproj (vision projector). In practice, llama-cpp / chat handlers
      may keep mmproj-related state alive longer than expected. To avoid "stale mmproj"
      issues when switching models (e.g., 8B -> 4B), we:
        * treat (model_path, mmproj_path, n_ctx, n_gpu_layers, vision_mode) as a signature
        * explicitly unload + gc before loading a different signature
        * make mmproj auto-detection choose the best match for the selected model
    """

    def __init__(self):
        self.model: Optional[Llama] = None
        self.chat_handler = None

        # Keep a full signature of what's currently loaded
        self._signature: Optional[tuple] = None
        self.current_model_path: Optional[str] = None
        self.current_mmproj_path: Optional[str] = None

    def _normalize_path(self, p: Optional[str]) -> Optional[str]:
        if p is None:
            return None
        return os.path.normpath(p)

    def _infer_is_qwen2(self, model_path: str) -> bool:
        model_name_lower = os.path.basename(model_path).lower()
        return "qwen2" in model_name_lower

    def _infer_is_qwen3(self, model_path: str) -> bool:
        model_name_lower = os.path.basename(model_path).lower()
        return ("qwen3vl" in model_name_lower) or ("qwen3-vl" in model_name_lower)

    def _infer_is_qwen35(self, model_path: str) -> bool:
        model_name_lower = os.path.basename(model_path).lower()
        return ("qwen35" in model_name_lower) or ("qwen3.5" in model_name_lower)

    def _auto_detect_mmproj(self, model_path: str) -> Optional[str]:
        """
        Auto-detect mmproj by "family prefix" match.

        Rule:
        - Determine model family by basename(model).startswith one of:
            qwen2, qwen3, llava, llama, gemma-3, glm-4
        - In the same directory, scan mmproj-*.gguf
        - Keep only those whose mmproj_name (after 'mmproj-') startswith the same family keyword
        - If exactly one match -> return it
        - Else (0 or >1) -> raise ValueError
        """
        model_dir = os.path.dirname(model_path)
        base = os.path.basename(model_path)
        name = base[:-5] if base.lower().endswith(".gguf") else base
        name_l = name.lower()

        # Model family keywords (startswith)
        families = ["qwen2", "qwen3vl", "qwen3-vl", "qwen35", "qwen3.5"]
        family = next((k for k in families if name_l.startswith(k)), None)

        if family is None:
            raise ValueError(
                "mmproj auto-detect failed: model name does not start with any supported family prefix.\n"
                f"model: {base}\n"
                f"supported prefixes: {', '.join(families)}"
            )

        if not os.path.exists(model_dir):
            raise ValueError(
                "mmproj auto-detect failed: model directory does not exist.\n"
                f"dir: {model_dir}"
            )

        # Collect mmproj-*.gguf in the same dir
        mmproj_files = [
            f for f in os.listdir(model_dir)
            if f.startswith("mmproj-") and f.endswith(".gguf")
        ]

        # Filter by family prefix on mmproj_name
        matches = []
        for f in mmproj_files:
            mmname = f[len("mmproj-"):-len(".gguf")]
            if mmname.lower().startswith(family):
                matches.append(f)

        matches.sort(key=str.lower)

        if len(matches) == 1:
            fname = matches[0]
            cand = os.path.join(model_dir, fname)
            print(f"[GGUFModelManager] Auto-detected mmproj (family={family}): {fname}")
            return self._normalize_path(cand)

        if len(matches) == 0:
            raise ValueError(
                "mmproj auto-detect failed: no mmproj matched the model family prefix.\n"
                f"model: {base}\n"
                f"family: {family}\n"
                f"dir: {model_dir}\n"
                f"mmproj candidates: {', '.join(sorted(mmproj_files, key=str.lower)) or '(none)'}"
            )

        # len(matches) > 1
        raise ValueError(
            "mmproj auto-detect failed: multiple mmproj files matched the model family prefix.\n"
            f"model: {base}\n"
            f"family: {family}\n"
            f"dir: {model_dir}\n"
            f"matched: {', '.join(matches)}\n"
            "Please select mmproj manually."
        )

    def _make_signature(
        self,
        model_path: str,
        mmproj_path: Optional[str],
        n_ctx: int,
        n_gpu_layers: int,
        use_vision: bool,
    ) -> tuple:
        return (
            self._normalize_path(model_path),
            self._normalize_path(mmproj_path),
            int(n_ctx),
            int(n_gpu_layers),
            bool(use_vision),
        )

    def load_model(
        self,
        model_path: str,
        mmproj_path: Optional[str] = None,
        n_ctx: int = 4096,
        n_gpu_layers: int = 0,
        verbose: bool = False,
    ) -> Llama:
        """Load GGUF model."""
        if not LLAMA_CPP_AVAILABLE:
            raise RuntimeError("llama-cpp-python is not available")

        model_path = self._normalize_path(model_path)

        # Infer Qwen version from model name
        is_qwen2 = self._infer_is_qwen2(model_path)
        is_qwen3 = self._infer_is_qwen3(model_path)
        is_qwen35 = self._infer_is_qwen35(model_path)

        # If user explicitly selected "(Not required)", force text-only even if the filename contains "qwen3".
        # This prevents accidental Qwen3-VL handler selection and mmproj auto-detection for text-only Qwen3 models.
        force_no_mmproj = (mmproj_path == "(Not required)")
        if force_no_mmproj:
            mmproj_path = None
        # Qwen2.5-VL requires mmproj (unless explicitly disabled)
        elif is_qwen2 and not force_no_mmproj:
            if mmproj_path is None:
                mmproj_path = self._auto_detect_mmproj(model_path)
                if mmproj_path is None:
                    model_dir = os.path.dirname(model_path)
                    raise ValueError(
                        "Qwen2.5-VL requires mmproj file!\n"
                        "Please download mmproj file from the model's GGUF repo.\n"
                        f"Expected location: {model_dir}{os.sep}mmproj-*.gguf"
                    )
        # Qwen3-VL requires mmproj (unless explicitly disabled)
        elif is_qwen3 and not force_no_mmproj:
            if mmproj_path is None:
                mmproj_path = self._auto_detect_mmproj(model_path)
                if mmproj_path is None:
                    model_dir = os.path.dirname(model_path)
                    raise ValueError(
                        "Qwen3-VL requires mmproj file!\n"
                        "Please download mmproj file from the model's GGUF repo.\n"
                        f"Expected location: {model_dir}{os.sep}mmproj-*.gguf"
                    )
        # Qwen3.5 requires mmproj (unless explicitly disabled)
        elif is_qwen35 and not force_no_mmproj:
            if mmproj_path is None:
                mmproj_path = self._auto_detect_mmproj(model_path)
                if mmproj_path is None:
                    model_dir = os.path.dirname(model_path)
                    raise ValueError(
                        "Qwen3.5 requires mmproj file!\n"
                        "Please download mmproj file from the model's GGUF repo.\n"
                        f"Expected location: {model_dir}{os.sep}mmproj-*.gguf"
                    )
            else:
                mmproj_path = self._normalize_path(mmproj_path)

            print(f"[GGUFModelManager] Using mmproj: {mmproj_path}")

        # Decide vision mode + initialize handler
        use_vision = False
        chat_handler = None

        if force_no_mmproj:
            print("[GGUFModelManager] Text-only mode forced: mmproj not required")
            chat_handler = None
            use_vision = False
        elif is_qwen2 and QWEN2_AVAILABLE:
            if mmproj_path is not None and os.path.exists(mmproj_path):
                try:
                    print(f"[GGUFModelManager] Qwen2.5-VL with mmproj: {mmproj_path}")
                    chat_handler = Qwen25VLChatHandler(clip_model_path=mmproj_path)
                    use_vision = True
                except Exception as e:
                    print(f"[GGUFModelManager] Warning: Failed to initialize Qwen2.5-VL chat handler: {e}")
                    chat_handler = None
                    use_vision = False
            else:
                print("[GGUFModelManager] Error: Qwen2.5-VL requires an existing mmproj file")
                chat_handler = None
                use_vision = False
        elif is_qwen3 and QWEN3_AVAILABLE:
            if mmproj_path is not None and os.path.exists(mmproj_path):
                try:
                    print(f"[GGUFModelManager] Qwen3-VL with mmproj: {mmproj_path}")
                    chat_handler = Qwen3VLChatHandler(clip_model_path=mmproj_path)
                    use_vision = True
                except Exception as e:
                    print(f"[GGUFModelManager] Warning: Failed to initialize Qwen3-VL chat handler: {e}")
                    chat_handler = None
                    use_vision = False
            else:
                print("[GGUFModelManager] Error: Qwen3-VL requires an existing mmproj file")
                chat_handler = None
                use_vision = False
        elif is_qwen35 and QWEN35_AVAILABLE:
            if mmproj_path is not None and os.path.exists(mmproj_path):
                try:
                    print(f"[GGUFModelManager] Qwen3.5 with mmproj: {mmproj_path}")
                    chat_handler = Qwen35ChatHandler(
                        clip_model_path=mmproj_path,
                        enable_thinking=False,
                    )
                    use_vision = True
                except Exception as e:
                    print(f"[GGUFModelManager] Warning: Failed to initialize Qwen3.5 chat handler: {e}")
                    chat_handler = None
                    use_vision = False
            else:
                print("[GGUFModelManager] Error: Qwen3.5 requires an existing mmproj file")
                chat_handler = None
                use_vision = False
        else:
            print("[GGUFModelManager] Using text-only mode")
            chat_handler = None
            use_vision = False

        new_sig = self._make_signature(
            model_path=model_path,
            mmproj_path=mmproj_path,
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu_layers,
            use_vision=use_vision,
        )

        # If signature matches, reuse
        if self.model is not None and self._signature == new_sig:
            print(f"[GGUFModelManager] Using cached model: {model_path}")
            return self.model

        # Otherwise, explicitly unload to avoid stale mmproj/handler state
        if self.model is not None:
            print("[GGUFModelManager] Signature changed -> unloading previous model to avoid stale state")
            self.unload_model()

        print(f"[GGUFModelManager] Loading model: {model_path}")
        print(f"[GGUFModelManager] n_ctx={n_ctx}, n_gpu_layers={n_gpu_layers}")

        # Store handler on manager (used later to decide if images are supported)
        self.chat_handler = chat_handler

        # Model loading
        if use_vision and self.chat_handler is not None:
            print("[GGUFModelManager] Loading with vision support")
            self.model = Llama(
                model_path=model_path,
                chat_handler=self.chat_handler,
                n_ctx=n_ctx,
                n_gpu_layers=n_gpu_layers,
                verbose=verbose,
                logits_all=True,  # required by some vision chat handlers
            )
        else:
            print("[GGUFModelManager] Loading in text-only mode")
            self.model = Llama(
                model_path=model_path,
                n_ctx=n_ctx,
                n_gpu_layers=n_gpu_layers,
                verbose=verbose,
            )

        self.current_model_path = model_path
        self.current_mmproj_path = self._normalize_path(mmproj_path)
        self._signature = new_sig

        print("[GGUFModelManager] Model loaded successfully")
        return self.model

    def unload_model(self):
        """Unload model from memory."""
        if self.model is not None:
            print(f"[GGUFModelManager] Unloading model: {self.current_model_path}")
        try:
            if self.model is not None:
                del self.model
        finally:
            self.model = None

        try:
            if self.chat_handler is not None:
                del self.chat_handler
        finally:
            self.chat_handler = None

        self.current_model_path = None
        self.current_mmproj_path = None
        self._signature = None

        # Encourage timely cleanup (important for llama-cpp backends and mmproj state)
        import gc as _gc
        _gc.collect()

# Global model manager
_model_manager = GGUFModelManager()

# ============================================================================
# Prompt Rewriter Function
# ============================================================================

def rewrite_prompt_with_gguf(
    prompt: str,
    model_path: str,
    mmproj_path: Optional[str] = None,
    style: str = "default",
    target_language: str = "auto",
    images: Optional[List[Image.Image]] = None,
    max_tokens: int = 512,
    temperature: float = 0.7,
    n_ctx: int = 4096,
    n_gpu_layers: int = 0,
) -> str:
    """
    Rewrite prompt using GGUF model
    
    Args:
        prompt: Original prompt
        model_path: Path to GGUF model
        mmproj_path: Path to mmproj file
        style: Prompt style
        target_language: Target language
        images: Input image list (optional)
        max_tokens: Maximum tokens to generate
        temperature: Temperature parameter
        n_ctx: Context size
        n_gpu_layers: Number of GPU layers
    
    Returns:
        Rewritten prompt
    """
    global _model_manager
    if _model_manager is None:
        _model_manager = GGUFModelManager()

    # Load model
    model = _model_manager.load_model(
        model_path=model_path,
        mmproj_path=mmproj_path,
        n_ctx=n_ctx,
        n_gpu_layers=n_gpu_layers,
    )
    
    # detection
    detected_lang = detect_language(prompt)
    lang = target_language if target_language != "auto" else detected_lang
    
    if style in {"qwen_image", "qwen_image_edit", "wan_t2v", "wan_i2v"}:
        system_prompt, prompt = build_special_prompt(style, prompt, lang)
    else:
        system_prompt = SYSTEM_PROMPTS.get(style, SYSTEM_PROMPTS["default"])

        if lang == "zh":
            prompt = f"[请用中文输出] {prompt}"
        elif lang == "en":
            prompt = f"[Please output in English] {prompt}"
    
    # construction
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]
    

    if images is not None and len(images) > 0:
        # Vision for
        if _model_manager.chat_handler is None:
            print("[Vision LLM Node] Warning: Images provided but vision is not available")
            print("[Vision LLM Node] Running in text-only mode, images will be ignored")
            print("[Vision LLM Node] To use vision features, please:")
            print("[Vision LLM Node]   1. Upgrade llama-cpp-python to latest version, OR")
            print("[Vision LLM Node]   2. Use Ollama instead")
        else:
            # Vision for
            user_content = []
            
            for img in images:
                img_b64 = encode_image_base64(img)
                user_content.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}
                })
            
            user_content.append({"type": "text", "text": prompt})
            messages[1]["content"] = user_content
    

    print(f"[Vision LLM Node] Generating with style='{style}', lang='{lang}'...")
    
    response = model.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
    )
    

    result = response['choices'][0]['message']['content']
    result = normalize_qwen_result(style, result)
    
    print(f"[Vision LLM Node] Original: {prompt[:100]}...")
    print(f"[Vision LLM Node] Enhanced: {result[:100]}...")
    
    return result

# ============================================================================
# ComfyUI Node
# ============================================================================

class VisionLLMNode:
    """
    Vision LLM Node - Local GGUF vision language models
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        available_models = discover_local_gguf_models(qwen_only=True)
        available_mmprojs = discover_local_mmproj_files()
        
        # mmproj selection
        mmproj_options = available_mmprojs + ["(Auto-detect)", "(Not required)"]
        if not available_mmprojs:
            mmproj_options = ["(Auto-detect)", "(Not required)"]
        

        if not available_models:
            available_models = ["(No Qwen GGUF models found in models/LLM or models/text_encoders)"]
        
        return {
            "required": {
                "prompt": ("STRING", {
                    "multiline": True,
                    "default": "",
                    "tooltip": "Original prompt to be enhanced"
                }),
                "style": (["raw", "default", "detailed", "concise", "creative"], {
                    "default": "raw",
                    "tooltip": "Prompt rewriting style (raw: no system prompt, direct LLM response)"
                }),
                "target_language": (["auto", "en", "zh"], {
                    "default": "auto",
                    "tooltip": "Target language for output prompt"
                }),
                "model": (available_models, {
                    "default": available_models[0],
                    "tooltip": "GGUF model to use for rewriting"
                }),
                "mmproj": (mmproj_options, {
                    "default": mmproj_options[0],
                    "tooltip": "mmproj file (select manually or use auto-detect)"
                }),
                "max_tokens": ("INT", {
                    "default": 512,
                    "min": 64,
                    "max": 2048,
                    "step": 64,
                    "tooltip": "Maximum tokens to generate"
                }),
                "temperature": ("FLOAT", {
                    "default": 0.7,
                    "min": 0.0,
                    "max": 2.0,
                    "step": 0.1,
                    "tooltip": "Sampling temperature"
                }),
                "device": (["CPU", "GPU"], {
                    "default": "CPU",
                    "tooltip": "Device to run model on (GPU requires compatible hardware)"
                }),
            },
            "optional": {
                "image": ("IMAGE", {
                    "tooltip": "Optional image input for vision-based enhancement"
                }),
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("STRING",)
    FUNCTION = "rewrite"
    CATEGORY = "multimodal/prompt"
    DESCRIPTION = "Local GGUF vision-language models for prompt generation"
    
    def rewrite(self, prompt: str, model: str, mmproj: str, style: str, target_language: str,
                max_tokens: int, temperature: float, device: str, image=None) -> tuple:
        """
        Rewrite prompt using GGUF model
        
        Returns:
            (enhanced_prompt,)
        """


        try:
            # llama-cpp-python check
            if not LLAMA_CPP_AVAILABLE:
                print("[Vision LLM Node] Error: llama-cpp-python not available")
                return (prompt,)

            if "(No Qwen GGUF models found" in model:
                print("[Vision LLM Node] Error: No Qwen GGUF models found in models/LLM or models/text_encoders")
                return (prompt,)

            try:
                # construction
                model_path = resolve_local_gguf_path(model)

                if not os.path.exists(model_path):
                    print(f"[Vision LLM Node] Error: Model not found: {model_path}")
                    return (prompt,)

                # mmproj processing
                mmproj_path = resolve_mmproj_path_for_model(model_path, mmproj)

                # PIL
                pil_images = None
                if image is not None:
                    pil_images = tensor2pil(image)
                    print(f"[Vision LLM Node] Using {len(pil_images)} image(s)")

                # Convert device selection to n_gpu_layers
                # GPU: -1 (all layers), CPU: 0 (no GPU)
                n_gpu_layers = -1 if device == "GPU" else 0

                # Load model
                print(f"[Vision LLM Node] Loading model: {model}")
                try:
                    enhanced_prompt = rewrite_prompt_with_gguf(
                        prompt=prompt,
                        model_path=model_path,
                        mmproj_path=mmproj_path,
                        style=style,
                        target_language=target_language,
                        images=pil_images,
                        max_tokens=max_tokens,
                        temperature=temperature,
                        n_gpu_layers=n_gpu_layers,
                    )
                    return (enhanced_prompt,)
                except ValueError as e:
                    # Translated
                    error_str = str(e)
                    if "Failed to load model" in error_str:
                        error_msg = (
                            f"ERROR: Failed to load model '{model}'"
                        )
                        print(f"[Vision LLM Node] {error_msg}")
                        raise RuntimeError(error_msg)
                    else:
                        raise

            except RuntimeError:
                # Translated
                raise
            except Exception as e:
                print(f"[Vision LLM Node] Unexpected error: {e}")
                import traceback
                traceback.print_exc()
                raise RuntimeError(f"Unexpected error during prompt rewrite: {str(e)}")
        finally:
            cleanup()

# ============================================================================
# ComfyUI Node Registration
# ============================================================================

NODE_CLASS_MAPPINGS = {
    "VisionLLMNode": VisionLLMNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "VisionLLMNode": "Vision LLM Node"
}

# ============================================================================
# Cleanup on module unload
# ============================================================================

def cleanup(finalize: bool = False):
    """Cleanup on module unload"""
    global _model_manager
    manager = _model_manager

    if manager is not None:
        try:
            manager.unload_model()
        except Exception as e:
            print(f"[Vision LLM Node] Warning during cleanup unload: {e}")

    import gc as _gc
    _gc.collect()
    try:
        import torch as _torch
        if _torch.cuda.is_available():
            _torch.cuda.empty_cache()
        if hasattr(_torch, "xpu") and _torch.xpu.is_available():
            _torch.xpu.empty_cache()
    except Exception:
        pass
    _gc.collect()

    if finalize:
        _model_manager = None

import atexit
atexit.register(lambda: cleanup(finalize=True))