280 lines
9.3 KiB
Python
280 lines
9.3 KiB
Python
"""
|
|
VLM节点基类
|
|
为VLM类节点(图像反推、视频反推等)提供专用基础能力
|
|
"""
|
|
|
|
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节点基类
|
|
|
|
提供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_id = service.get('id', '')
|
|
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
|
|
|
|
def _get_provider_config(
|
|
self,
|
|
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
|
|
|
|
def _image_to_base64(self, 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)
|
|
|
|
def _compute_image_hash(self, image_tensor: Optional[torch.Tensor]) -> str:
|
|
"""
|
|
计算图像tensor的哈希值(用于IS_CHANGED)
|
|
|
|
参数:
|
|
image_tensor: 图像tensor或None
|
|
|
|
返回:
|
|
MD5哈希值的十六进制字符串
|
|
"""
|
|
return compute_image_hash(image_tensor)
|
|
|
|
def _run_async_task(
|
|
self,
|
|
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 = self._run_async_task(
|
|
VisionService.analyze_image,
|
|
"zhipu",
|
|
image_data=data,
|
|
request_id=req_id,
|
|
...
|
|
)
|
|
"""
|
|
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"{self.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()
|
|
|
|
# 使用基类的线程中断检测方法
|
|
self._run_thread_with_interrupt(
|
|
thread_target,
|
|
(),
|
|
task_name="异步任务"
|
|
)
|
|
|
|
def _run_vision_task(
|
|
self,
|
|
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: 当检测到用户中断时
|
|
|
|
示例:
|
|
result = self._run_vision_task(
|
|
VisionService.analyze_image,
|
|
"zhipu",
|
|
image_data=data,
|
|
request_id=req_id,
|
|
prompt_content=prompt,
|
|
custom_provider="zhipu",
|
|
custom_provider_config=config
|
|
)
|
|
"""
|
|
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"{self.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()
|
|
|
|
# 创建取消事件
|
|
import threading
|
|
cancel_event = threading.Event()
|
|
|
|
# 使用基类的中断检测执行
|
|
return self._execute_with_interrupt(thread_task, (), cancel_event)
|
|
|