265 lines
8.7 KiB
Python
265 lines
8.7 KiB
Python
"""
|
||
VLM 节点基类
|
||
为 VLM 类节点（图像反推、视频反推等）提供专用基础能力
|
||
|
||
V3 迁移说明：
|
||
所有原实例方法（self）均已转换为类方法（cls），以适配 V3 execute 必须为 @classmethod 的要求。
|
||
"""
|
||
|
||
import asyncio
|
||
import threading
|
||
from typing import Any, Callable, Dict, Optional
|
||
|
||
import torch
|
||
from comfy.model_management import InterruptProcessingException
|
||
|
||
from .base_node import BaseNode
|
||
from ...utils.image import tensor_to_base64, compute_image_hash
|
||
from ...utils.common import format_api_error
|
||
|
||
|
||
class VLMNodeBase(BaseNode):
|
||
"""
|
||
VLM 节点基类（V3 Mixin 版本）
|
||
|
||
提供 VLM 节点专用功能：
|
||
- VLM Provider 配置获取
|
||
- 图像处理工具方法
|
||
- 图像哈希计算
|
||
- 统一的异步任务执行与中断处理
|
||
- 动态服务/模型列表生成
|
||
"""
|
||
|
||
@staticmethod
|
||
def get_vlm_service_options():
|
||
"""
|
||
获取所有可用的 VLM 服务/模型选项列表
|
||
|
||
返回格式: ["智谱/glm-4v-flash", "Ollama/llava:Q6_K", ...]
|
||
所有服务显示为"服务名/模型名"格式
|
||
|
||
返回：
|
||
List[str]: 服务/模型选项列表
|
||
"""
|
||
from ...config_manager import config_manager
|
||
|
||
options = []
|
||
services = config_manager.get_all_services()
|
||
|
||
for service in services:
|
||
service_name = service.get('name', '')
|
||
service_type = service.get('type', '')
|
||
|
||
# 百度翻译没有 vlm_models，跳过
|
||
if service_type == 'baidu':
|
||
continue
|
||
|
||
# 遍历 vlm_models
|
||
vlm_models = service.get('vlm_models', [])
|
||
for model in vlm_models:
|
||
model_name = model.get('name', '')
|
||
if model_name:
|
||
# 格式: "服务名/模型名"
|
||
options.append(f"{service_name}/{model_name}")
|
||
|
||
# 如果没有任何选项，返回默认值避免 ComfyUI 报错
|
||
if not options:
|
||
options = ["智谱"]
|
||
|
||
return options
|
||
|
||
@staticmethod
|
||
def parse_service_model(service_model_str: str):
|
||
"""
|
||
解析"服务名/模型名"格式的字符串
|
||
|
||
参数：
|
||
service_model_str: 服务/模型字符串，例如 "智谱/glm-4v-flash"
|
||
|
||
返回：
|
||
Tuple[str, Optional[str]]: (service_id, model_name)
|
||
- service_id: 服务 ID（如 'zhipu', 'ollama'）
|
||
- model_name: 模型名称，如果没有则为 None
|
||
"""
|
||
from ...config_manager import config_manager
|
||
|
||
# 分割字符串
|
||
if '/' in service_model_str:
|
||
service_name, model_name = service_model_str.split('/', 1)
|
||
else:
|
||
service_name = service_model_str
|
||
model_name = None
|
||
|
||
# 查找对应的 service_id
|
||
services = config_manager.get_all_services()
|
||
for service in services:
|
||
if service.get('name') == service_name:
|
||
return service.get('id'), model_name
|
||
|
||
# 未找到，返回 None
|
||
return None, None
|
||
|
||
@classmethod
|
||
def _get_provider_config(
|
||
cls,
|
||
config_manager: Any,
|
||
provider: str
|
||
) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
获取指定 VLM Provider 的配置
|
||
|
||
参数：
|
||
config_manager: 配置管理器实例
|
||
provider: Provider 标识符（如 'zhipu', 'ollama' 等）
|
||
|
||
返回：
|
||
Provider 配置字典，如果未找到则返回 None
|
||
"""
|
||
vision_config = config_manager.get_vision_config()
|
||
|
||
if 'providers' in vision_config and provider in vision_config['providers']:
|
||
return vision_config['providers'][provider]
|
||
|
||
return None
|
||
|
||
@classmethod
|
||
def _image_to_base64(cls, image_tensor: torch.Tensor, quality: int = 95) -> str:
|
||
"""
|
||
将图像 tensor 转换为 base64 编码
|
||
|
||
参数：
|
||
image_tensor: 图像 tensor
|
||
quality: JPEG 压缩质量（1-100）
|
||
|
||
返回：
|
||
base64 编码的 data URL
|
||
"""
|
||
return tensor_to_base64(image_tensor, quality)
|
||
|
||
@classmethod
|
||
def _compute_image_hash(cls, image_tensor: Optional[torch.Tensor]) -> str:
|
||
"""
|
||
计算图像 tensor 的哈希值（用于 fingerprint_inputs）
|
||
|
||
参数：
|
||
image_tensor: 图像 tensor 或 None
|
||
|
||
返回：
|
||
MD5 哈希值的十六进制字符串
|
||
"""
|
||
return compute_image_hash(image_tensor)
|
||
|
||
@classmethod
|
||
def _run_async_task(
|
||
cls,
|
||
async_func: Callable,
|
||
provider: str,
|
||
*args,
|
||
**kwargs
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
在独立线程中运行异步任务，并处理中断
|
||
|
||
这个方法封装了：
|
||
1. 异步事件循环的创建和清理
|
||
2. 中断异常的捕获和处理
|
||
3. 普通异常的格式化
|
||
|
||
参数：
|
||
async_func: 异步函数（coroutine function）
|
||
provider: Provider 名称（用于错误格式化）
|
||
*args, **kwargs: 传递给 async_func 的参数
|
||
|
||
返回：
|
||
{"success": bool, "data": dict, "error": str} 格式的结果
|
||
"""
|
||
result_container = {}
|
||
|
||
def thread_target():
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
# 执行异步任务
|
||
result = loop.run_until_complete(async_func(*args, **kwargs))
|
||
result_container['result'] = result
|
||
except (asyncio.CancelledError, KeyboardInterrupt):
|
||
# 捕获中断异常，不格式化错误
|
||
print(f"{cls.LOG_PREFIX} 检测到异步任务被取消")
|
||
result_container['result'] = {"success": False, "error": "任务被中断"}
|
||
except Exception as e:
|
||
# 捕获其他异常，格式化错误信息
|
||
error_message = format_api_error(e, provider)
|
||
result_container['result'] = {"success": False, "error": error_message}
|
||
finally:
|
||
loop.close()
|
||
|
||
# 使用基类的线程中断检测方法
|
||
cls._run_thread_with_interrupt(
|
||
thread_target,
|
||
(),
|
||
task_name="异步任务"
|
||
)
|
||
|
||
@classmethod
|
||
def _run_vision_task(
|
||
cls,
|
||
vision_service_func: Callable,
|
||
provider: str,
|
||
*args,
|
||
**kwargs
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
执行视觉任务（图像/视频分析）的统一方法
|
||
|
||
封装了：
|
||
1. 异步任务的线程执行
|
||
2. 中断异常的捕获和处理
|
||
3. API 异常的格式化
|
||
|
||
参数：
|
||
vision_service_func: VisionService 的异步方法（如 analyze_image, analyze_images）
|
||
provider: Provider 名称（用于错误格式化，如 "zhipu"）
|
||
*args, **kwargs: 传递给 vision_service_func 的参数
|
||
|
||
返回：
|
||
{"success": bool, "data": dict, "error": str} 格式的结果
|
||
|
||
异常：
|
||
InterruptProcessingException: 当检测到用户中断时
|
||
"""
|
||
def thread_task(result_container, cancel_event):
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
# 将 cancel_event 传递给服务函数
|
||
kwargs['cancel_event'] = cancel_event
|
||
result = loop.run_until_complete(vision_service_func(*args, **kwargs))
|
||
result_container['result'] = result
|
||
except (asyncio.CancelledError, KeyboardInterrupt):
|
||
# 捕获中断异常，不格式化错误
|
||
print(f"{cls.LOG_PREFIX} 检测到异步任务被取消")
|
||
result_container['result'] = {"success": False, "error": "任务被中断"}
|
||
except Exception as e:
|
||
# 捕获其他异常，格式化错误信息
|
||
error_message = format_api_error(e, provider)
|
||
result_container['result'] = {"success": False, "error": error_message}
|
||
finally:
|
||
# 清理所有未完成的任务，消除 "Task was destroyed but it is pending" 警告
|
||
try:
|
||
pending = asyncio.all_tasks(loop)
|
||
if pending:
|
||
for task in pending:
|
||
task.cancel()
|
||
# 等待任务取消完成，忽略取消错误
|
||
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
loop.close()
|
||
|
||
# 创建取消事件
|
||
cancel_event = threading.Event()
|
||
|
||
# 使用基类的中断检测执行
|
||
return cls._execute_with_interrupt(thread_task, (), cancel_event)
|