391 lines
16 KiB
Python
391 lines
16 KiB
Python
"""
|
||
图像反推（图片反推提示词）节点 - V3 版本
|
||
|
||
V3 迁移说明：
|
||
- 继承 VLMNodeBase + io.ComfyNode
|
||
- INPUT_TYPES → define_schema()
|
||
- IS_CHANGED → fingerprint_inputs()
|
||
- def analyze_image(self, ...) → @classmethod execute(cls, ...)
|
||
|
||
功能说明：
|
||
- 支持单张图像输入
|
||
- 支持 IMAGE batch 输入：逐帧独立调用 VLM，输出合并文本 + 字符串列表
|
||
- 两个输出端口：
|
||
· caption_text (STRING) - 所有帧结果用 "\\n---\\n" 合并的完整文本
|
||
· caption_list (STRING List) - 每帧独立结果的字符串列表（is_output_list=True）
|
||
"""
|
||
|
||
from comfy.model_management import InterruptProcessingException
|
||
from comfy_api.latest import io
|
||
|
||
from ..services.vlm import VisionService
|
||
from ..utils.common import (
|
||
format_api_error, format_model_with_thinking, generate_request_id,
|
||
log_prepare, log_error, TASK_IMAGE_CAPTION, SOURCE_NODE
|
||
)
|
||
from ..services.thinking_control import build_thinking_suppression
|
||
from .base import VLMNodeBase
|
||
|
||
|
||
class ImageCaptionNode(VLMNodeBase, io.ComfyNode):
|
||
"""图像反推节点（V3），支持单张/批量图像输入，输出合并文本与字符串列表"""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
# 动态获取服务列表
|
||
service_options = cls.get_vlm_service_options()
|
||
default_service = service_options[0] if service_options else "智谱"
|
||
|
||
# 获取模板
|
||
from ..config_manager import config_manager
|
||
system_prompts = config_manager.get_system_prompts()
|
||
|
||
vision_prompts = {}
|
||
active_vision_id = None
|
||
if system_prompts:
|
||
vision_prompts = system_prompts.get('vision_prompts', {}) or {}
|
||
active_vision_id = system_prompts.get('active_prompts', {}).get('vision')
|
||
|
||
prompt_template_options = []
|
||
id_to_display_name = {}
|
||
for key, value in vision_prompts.items():
|
||
show_in = value.get('showIn', ["frontend", "node"])
|
||
if 'node' not in show_in:
|
||
continue
|
||
name = value.get('name', key)
|
||
category = value.get('category', '')
|
||
display_name = f"{category}/{name}" if category else name
|
||
id_to_display_name[key] = display_name
|
||
prompt_template_options.append(display_name)
|
||
|
||
default_template_name = prompt_template_options[0] if prompt_template_options else "反推-自然语言"
|
||
if active_vision_id and active_vision_id in id_to_display_name:
|
||
default_template_name = id_to_display_name[active_vision_id]
|
||
|
||
if not prompt_template_options:
|
||
prompt_template_options = ["反推-自然语言"]
|
||
|
||
return io.Schema(
|
||
node_id="ImageCaptionNode",
|
||
display_name="✨Image Caption (VLM)",
|
||
category="✨Prompt Assistant",
|
||
description="Extract text prompt from image using Vision-Language Models",
|
||
inputs=[
|
||
io.Image.Input(
|
||
"image",
|
||
tooltip="The image to analyze. Supports single image or IMAGE batch (processes each frame independently)"
|
||
),
|
||
io.Combo.Input(
|
||
"rule",
|
||
options=prompt_template_options,
|
||
default=default_template_name,
|
||
tooltip="Choose a preset rule for analysis"
|
||
),
|
||
io.Boolean.Input(
|
||
"custom_rule",
|
||
default=False,
|
||
label_on="Enable",
|
||
label_off="Disable",
|
||
tooltip="Enable custom rule input"
|
||
),
|
||
io.String.Input(
|
||
"custom_rule_content",
|
||
multiline=True,
|
||
default="",
|
||
tooltip="Custom rule content, only used when Custom Rule is enabled"
|
||
),
|
||
io.String.Input(
|
||
"user_prompt",
|
||
multiline=True,
|
||
default="",
|
||
tooltip="Enter additional prompts here, sent with the rule"
|
||
),
|
||
io.Combo.Input(
|
||
"vlm_service",
|
||
options=service_options,
|
||
default=default_service,
|
||
tooltip="Select VLM Service"
|
||
),
|
||
io.Boolean.Input(
|
||
"ollama_auto_unload",
|
||
default=True,
|
||
label_on="Enable",
|
||
label_off="Disable",
|
||
tooltip="Auto unload Ollama model after generation"
|
||
),
|
||
io.Int.Input(
|
||
"seed",
|
||
default=0,
|
||
min=0,
|
||
max=0xffffffffffffffff,
|
||
control_after_generate=True,
|
||
tooltip="Controls randomness. Set to non-fixed mode to force re-execution"
|
||
),
|
||
],
|
||
outputs=[
|
||
# 合并文本（单张直接输出；batch 时用 "\n---\n" 拼接）
|
||
io.String.Output("caption_text"),
|
||
# 每帧独立结果列表（单张时为单元素列表）
|
||
io.String.Output("caption_list", is_output_list=True),
|
||
],
|
||
hidden=[io.Hidden.unique_id],
|
||
)
|
||
|
||
# -------------------------------------------------------------------------
|
||
# fingerprint_inputs（替代 IS_CHANGED）
|
||
# -------------------------------------------------------------------------
|
||
|
||
@classmethod
|
||
def fingerprint_inputs(
|
||
cls,
|
||
image=None, rule=None, custom_rule=None, custom_rule_content=None,
|
||
user_prompt=None, vlm_service=None, ollama_auto_unload=None, seed=None
|
||
):
|
||
"""替代 V1 IS_CHANGED"""
|
||
import hashlib
|
||
temp_rule_hash = hashlib.md5((custom_rule_content or "").encode('utf-8')).hexdigest()
|
||
user_hint_hash = hashlib.md5((user_prompt or "").encode('utf-8')).hexdigest()
|
||
img_hash = cls._compute_image_hash(image)
|
||
|
||
return hash((
|
||
img_hash,
|
||
rule,
|
||
bool(custom_rule),
|
||
temp_rule_hash,
|
||
user_hint_hash,
|
||
vlm_service,
|
||
bool(ollama_auto_unload),
|
||
seed
|
||
))
|
||
|
||
# -------------------------------------------------------------------------
|
||
# 内部辅助：分析单张图像
|
||
# -------------------------------------------------------------------------
|
||
|
||
@classmethod
|
||
def _analyze_single_image(
|
||
cls,
|
||
image_data: str,
|
||
prompt_to_send: str,
|
||
rule_name: str,
|
||
service_id: str,
|
||
service: dict,
|
||
provider_config: dict,
|
||
request_id: str,
|
||
frame_index=None
|
||
) -> str:
|
||
"""
|
||
对单张图像调用 VLM 进行分析，返回文本结果。
|
||
|
||
参数：
|
||
image_data : base64 编码的图像 data URL
|
||
prompt_to_send: 拼接好的提示词（system + user_prompt）
|
||
rule_name : 当前规则名称（用于日志）
|
||
service_id : 服务 ID
|
||
service : 服务配置字典
|
||
provider_config: API 调用参数字典
|
||
request_id : 本次请求 ID
|
||
frame_index : batch 模式下当前帧序号（None 表示单张模式）
|
||
|
||
返回：
|
||
分析结果文本字符串
|
||
"""
|
||
frame_label = f" [Frame {frame_index + 1}]" if frame_index is not None else ""
|
||
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_IMAGE_CAPTION, request_id, SOURCE_NODE,
|
||
service_display_name, model_display, rule_name + frame_label
|
||
)
|
||
|
||
result = cls._run_vision_task(
|
||
VisionService.analyze_image,
|
||
service_id,
|
||
image_data=image_data,
|
||
prompt_content=prompt_to_send,
|
||
request_id=request_id,
|
||
custom_provider=service_id,
|
||
custom_provider_config=provider_config,
|
||
source=SOURCE_NODE
|
||
)
|
||
|
||
if result and result.get('success'):
|
||
data = result.get('data', {})
|
||
# V3 重构后返回键名为 description，兼容旧版 caption
|
||
caption_text = data.get('description', data.get('caption', '')).strip()
|
||
if not caption_text:
|
||
error_msg = 'API returned empty result'
|
||
log_error(TASK_IMAGE_CAPTION, request_id, error_msg, source=SOURCE_NODE)
|
||
raise RuntimeError(f"Analysis failed: {error_msg}")
|
||
return caption_text
|
||
else:
|
||
error_msg = (
|
||
result.get('error', 'Unknown error') if result
|
||
else 'No result returned'
|
||
)
|
||
if error_msg == "任务被中断":
|
||
raise InterruptProcessingException()
|
||
log_error(TASK_IMAGE_CAPTION, request_id, error_msg, source=SOURCE_NODE)
|
||
raise RuntimeError(f"Analysis failed: {error_msg}")
|
||
|
||
# -------------------------------------------------------------------------
|
||
# execute（V3 主执行方法）
|
||
# -------------------------------------------------------------------------
|
||
|
||
@classmethod
|
||
def execute(
|
||
cls,
|
||
image, rule, custom_rule, custom_rule_content,
|
||
user_prompt, vlm_service, ollama_auto_unload, seed=None
|
||
):
|
||
unique_id = cls.hidden.unique_id
|
||
request_id = None
|
||
|
||
try:
|
||
if image is None:
|
||
raise ValueError("No image provided. Please connect an image to the 'image' input.")
|
||
|
||
# ------------------------------------------------------------------
|
||
# 1. 构建提示词
|
||
# ------------------------------------------------------------------
|
||
rule_name = "Custom Rule" if (custom_rule and custom_rule_content) else rule
|
||
system_message = None
|
||
|
||
if custom_rule and custom_rule_content:
|
||
system_message = {"role": "system", "content": custom_rule_content}
|
||
else:
|
||
from ..config_manager import config_manager
|
||
system_prompts = config_manager.get_system_prompts()
|
||
vision_prompts = system_prompts.get('vision_prompts', {}) if system_prompts else {}
|
||
|
||
template_found = False
|
||
for key, value in vision_prompts.items():
|
||
name = value.get('name', key)
|
||
category = value.get('category', '')
|
||
display_name = f"{category}/{name}" if category else name
|
||
if display_name == rule or value.get('name') == rule or key == rule:
|
||
system_message = {
|
||
"role": value.get('role', 'system'),
|
||
"content": value.get('content', '')
|
||
}
|
||
template_found = True
|
||
break
|
||
|
||
if not template_found or not system_message or not system_message.get('content'):
|
||
system_message = {"role": "system", "content": "请详细描述这张图片的内容"}
|
||
rule_name = "Default Rule"
|
||
|
||
system_text = (
|
||
system_message.get('content', '')
|
||
if isinstance(system_message, dict)
|
||
else str(system_message)
|
||
)
|
||
prompt_to_send = f"{system_text}\n\n{user_prompt}".strip() if user_prompt else system_text
|
||
|
||
# ------------------------------------------------------------------
|
||
# 2. 解析服务与模型
|
||
# ------------------------------------------------------------------
|
||
service_id, model_name = cls.parse_service_model(vlm_service)
|
||
if not service_id:
|
||
raise ValueError(f"Invalid service selection: {vlm_service}")
|
||
|
||
from ..config_manager import config_manager
|
||
service = config_manager.get_service(service_id)
|
||
if not service:
|
||
raise ValueError(f"Service config not found: {vlm_service}")
|
||
|
||
vlm_models = service.get('vlm_models', [])
|
||
target_model = None
|
||
if model_name:
|
||
target_model = next((m for m in vlm_models if m.get('name') == model_name), None)
|
||
if not target_model:
|
||
target_model = next(
|
||
(m for m in vlm_models if m.get('is_default')),
|
||
vlm_models[0] if vlm_models else None
|
||
)
|
||
if not target_model:
|
||
raise ValueError(f"Service {vlm_service} has no available vision 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),
|
||
}
|
||
if service.get('type') == 'ollama':
|
||
provider_config['auto_unload'] = ollama_auto_unload
|
||
|
||
if not provider_config.get('model', ''):
|
||
raise ValueError(f"Please configure model for {vlm_service}")
|
||
if cls._service_requires_api_key(service) and not provider_config.get('api_key', ''):
|
||
raise ValueError(f"Please configure API key and model for {vlm_service}")
|
||
|
||
# ------------------------------------------------------------------
|
||
# 3. 分析图像：支持 batch（逐帧独立调用）
|
||
# ------------------------------------------------------------------
|
||
request_id = generate_request_id("vlm", None, unique_id)
|
||
|
||
# 检查是否为多帧 batch [N, H, W, C] 且 N > 1
|
||
if len(image.shape) == 4 and image.shape[0] > 1:
|
||
# Batch 模式：逐帧独立调用 VLM
|
||
batch_size = image.shape[0]
|
||
results: list[str] = []
|
||
|
||
for i in range(batch_size):
|
||
single_frame = image[i:i + 1] # 保持 4D 形状 [1, H, W, C]
|
||
image_data = cls._image_to_base64(single_frame)
|
||
frame_provider_config = provider_config
|
||
if service.get('type') == 'ollama':
|
||
frame_provider_config = provider_config.copy()
|
||
frame_provider_config['auto_unload'] = bool(ollama_auto_unload) and i == batch_size - 1
|
||
# 为每帧生成独立的 request_id，便于日志追踪
|
||
frame_request_id = generate_request_id("vlm", None, f"{unique_id}_f{i}")
|
||
description = cls._analyze_single_image(
|
||
image_data=image_data,
|
||
prompt_to_send=prompt_to_send,
|
||
rule_name=rule_name,
|
||
service_id=service_id,
|
||
service=service,
|
||
provider_config=frame_provider_config,
|
||
request_id=frame_request_id,
|
||
frame_index=i
|
||
)
|
||
results.append(description)
|
||
|
||
# 合并文本（与旧版保持一致）
|
||
combined_text = "\n---\n".join(results)
|
||
return io.NodeOutput(combined_text, results)
|
||
|
||
else:
|
||
# 单张模式
|
||
image_data = cls._image_to_base64(image)
|
||
description = cls._analyze_single_image(
|
||
image_data=image_data,
|
||
prompt_to_send=prompt_to_send,
|
||
rule_name=rule_name,
|
||
service_id=service_id,
|
||
service=service,
|
||
provider_config=provider_config,
|
||
request_id=request_id,
|
||
frame_index=None
|
||
)
|
||
# 单张时 caption_list 为单元素列表
|
||
return io.NodeOutput(description, [description])
|
||
|
||
except InterruptProcessingException:
|
||
raise
|
||
except Exception as e:
|
||
error_msg = format_api_error(e, vlm_service)
|
||
log_error(TASK_IMAGE_CAPTION, request_id, error_msg, source=SOURCE_NODE)
|
||
raise RuntimeError(f"Analysis error: {error_msg}")
|