nodes/yawiii-ComfyUI-Prompt-Assistant
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

继续优化请求流程

Browse Source
This commit is contained in:
yawiii
2026-04-25 10:38:46 +08:00
parent 78639c337b
commit 1d005acfe1
9 changed files with 112 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2 -1
View File
@@ -10,4 +10,5 @@ __pycache__/
.DS_Store
skills-lock.json
js/version.js
doc/
doc/
AGENTS.md
node/base/base_node.py
+24
View File
@@ -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,
node/expand_node.py
+3 -1
View File
@@ -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}")
# 执行扩写（异步线程 + 可中断）
node/image_caption_node.py
+9 -3
View File
@@ -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}")
raise RuntimeError(f"Analysis error: {error_msg}")
node/translate_node.py
+3 -1
View File
@@ -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}"}
# 执行翻译（异步线程 + 可中断）
node/video_caption_node.py
+42 -16
View File
@@ -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(
services/ollama_native.py
+21 -6
View File
@@ -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
services/openai_base.py
+8 -3
View File
@@ -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:
test_load.py
-8
View File
@@ -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()