From 1d005acfe17bcd02ee7d899600c87c2945751f68 Mon Sep 17 00:00:00 2001 From: yawiii <77927734@qq.com> Date: Sat, 25 Apr 2026 10:38:46 +0800 Subject: [PATCH] =?UTF-8?q?=E7=BB=A7=E7=BB=AD=E4=BC=98=E5=8C=96=E8=AF=B7?= =?UTF-8?q?=E6=B1=82=E6=B5=81=E7=A8=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 3 +- node/base/base_node.py | 24 ++++++++++++++++ node/expand_node.py | 4 ++- node/image_caption_node.py | 12 ++++++-- node/translate_node.py | 4 ++- node/video_caption_node.py | 58 +++++++++++++++++++++++++++----------- services/ollama_native.py | 27 ++++++++++++++---- services/openai_base.py | 11 ++++++-- test_load.py | 8 ------ 9 files changed, 112 insertions(+), 39 deletions(-) delete mode 100644 test_load.py diff --git a/.gitignore b/.gitignore index d50e83a..f3295ac 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,5 @@ __pycache__/ .DS_Store skills-lock.json js/version.js -doc/ \ No newline at end of file +doc/ +AGENTS.md \ No newline at end of file diff --git a/node/base/base_node.py b/node/base/base_node.py index 822207d..cf3bc63 100644 --- a/node/base/base_node.py +++ b/node/base/base_node.py @@ -196,6 +196,30 @@ class BaseNode: config_copy['auto_unload'] = auto_unload return config_copy + @classmethod + def _service_requires_api_key(cls, service: Optional[Dict[str, Any]]) -> bool: + """ + Decide whether a node should block execution before calling the service + because the API key is missing. + """ + if not service: + return True + + service_type = str(service.get('type', '') or '').strip().lower() + if service_type == 'ollama': + return False + + for key in ('requires_api_key', 'api_key_required', 'auth_required'): + if service.get(key) is False: + return False + + base_url = str(service.get('base_url', '') or '').strip().lower() + local_hosts = ('localhost', '127.0.0.1', '0.0.0.0', '[::1]', '::1') + if any(host in base_url for host in local_hosts): + return False + + return True + @classmethod def _get_prompt_template( cls, diff --git a/node/expand_node.py b/node/expand_node.py index e451950..1efa491 100644 --- a/node/expand_node.py +++ b/node/expand_node.py @@ -294,7 +294,9 @@ class PromptExpand(LLMNodeBase, io.ComfyNode): log_prepare(TASK_EXPAND, request_id, SOURCE_NODE, service_display_name, model_display, rule_name, {"长度": len(combined_text)}) # 检查 API 密钥和模型 - if not provider_config.get('api_key', '') or not provider_config.get('model', ''): + if not provider_config.get('model', ''): + raise ValueError(f"Please configure model for {llm_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 {llm_service}") # 执行扩写(异步线程 + 可中断) diff --git a/node/image_caption_node.py b/node/image_caption_node.py index e9d037a..678ddc3 100644 --- a/node/image_caption_node.py +++ b/node/image_caption_node.py @@ -325,7 +325,9 @@ class ImageCaptionNode(VLMNodeBase, io.ComfyNode): if service.get('type') == 'ollama': provider_config['auto_unload'] = ollama_auto_unload - if not provider_config.get('api_key', '') or not provider_config.get('model', ''): + 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}") # ------------------------------------------------------------------ @@ -342,6 +344,10 @@ class ImageCaptionNode(VLMNodeBase, io.ComfyNode): 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( @@ -350,7 +356,7 @@ class ImageCaptionNode(VLMNodeBase, io.ComfyNode): rule_name=rule_name, service_id=service_id, service=service, - provider_config=provider_config, + provider_config=frame_provider_config, request_id=frame_request_id, frame_index=i ) @@ -381,4 +387,4 @@ class ImageCaptionNode(VLMNodeBase, io.ComfyNode): 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}") \ No newline at end of file + raise RuntimeError(f"Analysis error: {error_msg}") diff --git a/node/translate_node.py b/node/translate_node.py index 915244e..c3ddc1f 100644 --- a/node/translate_node.py +++ b/node/translate_node.py @@ -210,7 +210,9 @@ class PromptTranslate(LLMNodeBase, io.ComfyNode): api_key = provider_config.get('api_key', '') model = provider_config.get('model', '') - if not api_key or not model: + if not model: + return request_id, {"success": False, "error": f"Please configure model for {service_display_name}"} + if cls._service_requires_api_key(service) and not api_key: return request_id, {"success": False, "error": f"Please configure API key and model for {service_display_name}"} # 执行翻译(异步线程 + 可中断) diff --git a/node/video_caption_node.py b/node/video_caption_node.py index 7303d05..f75e78f 100644 --- a/node/video_caption_node.py +++ b/node/video_caption_node.py @@ -406,13 +406,37 @@ class VideoCaptionNode(VLMNodeBase, io.ComfyNode): # 提取选中帧张量 [K, H, W, C] preview_tensor = tensor[selected] + request_frame_count = preview_tensor.shape[0] + if request_frame_count <= 4: + max_edge = 1024 + jpeg_quality = 85 + elif request_frame_count <= 8: + max_edge = 896 + jpeg_quality = 80 + else: + max_edge = 768 + jpeg_quality = 75 + # 转换为 base64 列表 base64_list = [] + resample_filter = getattr( + getattr(PILImage, "Resampling", PILImage), + "LANCZOS", + PILImage.BICUBIC, + ) for i in range(preview_tensor.shape[0]): frame_np = (preview_tensor[i].cpu().numpy() * 255).clip(0, 255).astype(np.uint8) img = PILImage.fromarray(frame_np) + longest_edge = max(img.size) + if longest_edge > max_edge: + scale = max_edge / longest_edge + new_size = ( + max(1, int(img.size[0] * scale)), + max(1, int(img.size[1] * scale)), + ) + img = img.resize(new_size, resample_filter) buf = BytesIO() - img.save(buf, format="JPEG", quality=85) + img.save(buf, format="JPEG", quality=jpeg_quality, optimize=True) encoded = base64.b64encode(buf.getvalue()).decode('utf-8') base64_list.append(f"data:image/jpeg;base64,{encoded}") @@ -485,20 +509,6 @@ class VideoCaptionNode(VLMNodeBase, io.ComfyNode): system_message = {"role": "system", "content": "请详细描述这段视频的内容"} rule_name = "Default Rule" - # 注入帧数元信息,帮助模型逐帧分析 - actual_frame_count = len(frames_data) - system_text = ( - system_message.get('content', '') - if isinstance(system_message, dict) - else str(system_message) - ) - frame_info_prefix = ( - f"[重要提示:本次共提供了 {actual_frame_count} 帧图像," - f"请务必逐帧分析,确保输出的描述数量与帧数一致。]\n\n" - ) - system_text = frame_info_prefix + system_text - prompt_to_send = f"{system_text}\n\n{user_prompt}".strip() if user_prompt else system_text - # ------------------------------------------------------------------ # 4. 解析服务与模型 # ------------------------------------------------------------------ @@ -554,6 +564,20 @@ class VideoCaptionNode(VLMNodeBase, io.ComfyNode): f"已忽略:{truncated_count}帧 | 模型:{model_full_name}" ) + # 注入帧数元信息,帮助模型逐帧分析 + actual_frame_count = len(frames_data) + system_text = ( + system_message.get('content', '') + if isinstance(system_message, dict) + else str(system_message) + ) + frame_info_prefix = ( + f"[重要提示:本次共提供了 {actual_frame_count} 帧图像," + f"请务必逐帧分析,确保输出的描述数量与帧数一致。]\n\n" + ) + system_text = frame_info_prefix + system_text + prompt_to_send = f"{system_text}\n\n{user_prompt}".strip() if user_prompt else system_text + # ------------------------------------------------------------------ # 6. 准备日志与 API 调用 # ------------------------------------------------------------------ @@ -573,7 +597,9 @@ class VideoCaptionNode(VLMNodeBase, io.ComfyNode): {"抽帧模式": sampling_mode, "帧数": len(frames_data)} ) - if not provider_config.get('api_key', '') or not provider_config.get('model', ''): + 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}") result = cls._run_vision_task( diff --git a/services/ollama_native.py b/services/ollama_native.py index c073e26..1ec046a 100644 --- a/services/ollama_native.py +++ b/services/ollama_native.py @@ -5,6 +5,7 @@ from typing import Any, Callable, Dict, Optional import httpx +from .core import HTTPClientPool from .ollama_utils import ( build_ollama_final_answer_retry_payload, collect_ollama_message_parts, @@ -31,11 +32,15 @@ class OllamaNativeAdapter: ) -> Dict[str, Any]: start_time = time.perf_counter() - async with httpx.AsyncClient( - timeout=httpx.Timeout(timeout, connect=10.0, read=timeout), + client = HTTPClientPool.get_client( + provider="ollama_native", + base_url=native_base, + timeout=timeout, + verify_ssl=True, limits=httpx.Limits(max_keepalive_connections=5, max_connections=10), - trust_env=False, - ) as client: + ) + + try: async def request_once(current_payload: Dict[str, Any]) -> Dict[str, Any]: full_content = "" reasoning_content = "" @@ -88,10 +93,14 @@ class OllamaNativeAdapter: content, reasoning = collect_ollama_message_parts(message) if reasoning: reasoning_content += reasoning + progress_count = len(full_content) + len(reasoning_content) + pbar.set_generating(progress_count) + pbar.update(progress_count) if content: full_content += content - pbar.set_generating(len(full_content)) - pbar.update(len(full_content)) + progress_count = len(full_content) + len(reasoning_content) + pbar.set_generating(progress_count) + pbar.update(progress_count) if stream_callback: stream_callback(content) @@ -165,12 +174,16 @@ class OllamaNativeAdapter: if should_retry_without_ollama_think(request_payload, result): request_payload = request_payload.copy() request_payload.pop("think", None) + if not monitor_task.done(): + monitor_task.cancel() req_task = asyncio.create_task(request_once(request_payload)) monitor_task = asyncio.create_task(monitor_interrupts(req_task)) result = await req_task if is_ollama_reasoning_only_response(result): retry_payload = build_ollama_final_answer_retry_payload(payload) + if not monitor_task.done(): + monitor_task.cancel() req_task = asyncio.create_task(request_once(retry_payload)) monitor_task = asyncio.create_task(monitor_interrupts(req_task)) result = await req_task @@ -188,3 +201,5 @@ class OllamaNativeAdapter: finally: if not monitor_task.done(): monitor_task.cancel() + finally: + pass diff --git a/services/openai_base.py b/services/openai_base.py index 756fad0..338d2dd 100644 --- a/services/openai_base.py +++ b/services/openai_base.py @@ -316,12 +316,17 @@ class OpenAICompatibleService(BaseAPIService): delta.get('thinking_process', '') or # 备选 '' ) - if reasoning: reasoning_content += reasoning + if reasoning: + reasoning_content += reasoning + progress_count = len(full_content) + len(reasoning_content) + pbar.set_generating(progress_count) + pbar.update(progress_count) if content: full_content += content if stream_callback: stream_callback(content) - pbar.set_generating(len(full_content)) - pbar.update(len(full_content)) + progress_count = len(full_content) + len(reasoning_content) + pbar.set_generating(progress_count) + pbar.update(progress_count) except json.JSONDecodeError: continue # 非 JSON 行(如注释、空行),正常跳过 except asyncio.CancelledError: diff --git a/test_load.py b/test_load.py deleted file mode 100644 index ed0c9d0..0000000 --- a/test_load.py +++ /dev/null @@ -1,8 +0,0 @@ -import sys -sys.path.insert(0, "/Users/laiyuewei/Desktop/ComfyUI.dev/ComfyUI") -try: - from custom_nodes.ComfyUI_Prompt_Assistant import __init__ as pa - print("Extension loaded successfully") -except Exception as e: - import traceback - traceback.print_exc()