949 lines
41 KiB
Python
949 lines
41 KiB
Python
"""
|
|
ComfyUI custom node: ClaudeCliVisionText
|
|
|
|
Calls the local `claude` CLI (`claude -p`) for a vision+text or text-only
|
|
completion. No model-loader input.
|
|
|
|
Design follows the project's answer set (claude_vision.md and the CLI
|
|
verification/decision answers). Key behaviors:
|
|
|
|
- Image attachment uses `@/abs/path.png` embedded in the prompt text
|
|
(verification required per target CLI version). Flip ATTACH_FLAG to
|
|
"--image" or similar if the CLI supports a native flag.
|
|
- Optional IMAGE input; unwired → text-only completion.
|
|
- `send_all_frames` toggle — default sends frame 0 only; on, attaches
|
|
every frame in the batch (for cross-frame reasoning).
|
|
- `--system-prompt` replaces Claude Code's default system prompt.
|
|
- Session handling: if `reuse_session` is on, tries `--resume` first; on
|
|
"unknown session" or "already in use" errors, falls back to
|
|
`--session-id <new-uuid>` with a fresh UUID.
|
|
- Model widget accepts shorthand (haiku/sonnet/opus) — forward-compatible
|
|
as Anthropic releases new versions. "default" omits `--model`.
|
|
- Response parsed from the final envelope's `.result` field (stream-json
|
|
format is a list of objects; last has type:"result").
|
|
- Auth-failure detection: non-zero exit with "not logged in" /
|
|
"authentication_failed" in stderr raises a ComfyUI-visible error.
|
|
- Optional `api_key` widget enables headless mode: sets
|
|
`ANTHROPIC_API_KEY` in subprocess env and passes `--bare` (disables
|
|
session persistence; node warns if reuse_session is also on).
|
|
- Cross-platform: pathlib.Path for all paths, shutil.which for CLI
|
|
resolution, subprocess with list args (no shell), session dir under
|
|
`<ComfyUI>/output/claude-sessions/` via folder_paths when available.
|
|
- Debug output: routed through logging.getLogger("comfyui.claude-cli")
|
|
for the console, plus a per-session `<session-id>.debug.log` file in
|
|
the session dir on CLI failures.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
try:
|
|
import folder_paths # Available when running inside ComfyUI.
|
|
COMFYUI_OUTPUT_DIR = Path(folder_paths.get_output_directory())
|
|
except Exception:
|
|
COMFYUI_OUTPUT_DIR = Path.cwd() / "output"
|
|
|
|
|
|
CLAUDE_SESSIONS_DIR = COMFYUI_OUTPUT_DIR / "claude-sessions"
|
|
VISION_MAX_LONG_EDGE = 1568
|
|
DEFAULT_TIMEOUT_SECONDS = 120
|
|
ATTACH_FLAG = "" # "" → embed @path in prompt; e.g. "--image" to use a native flag
|
|
|
|
AUTH_ERROR_SUBSTRINGS = (
|
|
"not logged in",
|
|
"authentication_failed",
|
|
"authentication failed",
|
|
"please login",
|
|
"run `claude login`",
|
|
"claude login",
|
|
)
|
|
|
|
MODEL_DROPDOWN_OPTIONS = [
|
|
"haiku",
|
|
"sonnet",
|
|
"opus",
|
|
"claude-haiku-4-5",
|
|
"claude-sonnet-4-6",
|
|
"claude-opus-4-7",
|
|
]
|
|
|
|
# Each `claude -p` invocation cold-starts an agent environment (~30k tokens to
|
|
# initialize system prompt, tools, context). To keep per-frame cost sane across
|
|
# a batch, STRONGLY prefer reusing a single session_id across many calls — the
|
|
# context is cached and subsequent calls are dramatically cheaper.
|
|
#
|
|
# Agent dropdown. `--agent <name>` selects a sub-agent preset. `general-purpose`
|
|
# is the safe fallback if a named preset isn't installed on the target machine.
|
|
# (unsure) whether every entry below is available on all CLI versions — the
|
|
# node passes the string through verbatim and surfaces the CLI's error if the
|
|
# agent is unknown.
|
|
AGENT_DROPDOWN_OPTIONS = [
|
|
"general-purpose", # generic fallback — always available
|
|
"Explore", # fast codebase exploration / search
|
|
"Plan", # planning / architecture
|
|
"code-reviewer", # code review
|
|
"debugger", # debugging errors, test failures
|
|
"data-scientist", # data analysis / SQL
|
|
]
|
|
|
|
logger = logging.getLogger("comfyui.claude-cli")
|
|
|
|
|
|
def ensure_sessions_dir_exists() -> None:
|
|
CLAUDE_SESSIONS_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def write_session_debug_log(session_id: str, payload: dict) -> None:
|
|
try:
|
|
ensure_sessions_dir_exists()
|
|
log_path = CLAUDE_SESSIONS_DIR / f"{session_id}.debug.log"
|
|
with log_path.open("a", encoding="utf-8") as log_file:
|
|
log_file.write(f"--- {time.strftime('%Y-%m-%d %H:%M:%S')} ---\n")
|
|
log_file.write(json.dumps(payload, default=str, indent=2))
|
|
log_file.write("\n")
|
|
except OSError as write_error:
|
|
logger.debug("failed to write session debug log: %s", write_error)
|
|
|
|
|
|
def resize_image_array_to_long_edge(image_array: np.ndarray, max_long_edge: int) -> np.ndarray:
|
|
height, width = image_array.shape[:2]
|
|
long_edge = max(height, width)
|
|
if long_edge <= max_long_edge:
|
|
return image_array
|
|
scale = max_long_edge / float(long_edge)
|
|
new_width = int(round(width * scale))
|
|
new_height = int(round(height * scale))
|
|
pil_image = Image.fromarray(image_array)
|
|
pil_image = pil_image.resize((new_width, new_height), Image.LANCZOS)
|
|
return np.asarray(pil_image)
|
|
|
|
|
|
def write_image_tensor_frame_to_png(image_tensor, frame_index: int) -> Path:
|
|
array = image_tensor
|
|
if hasattr(array, "cpu"):
|
|
array = array.cpu().numpy()
|
|
array = np.asarray(array)
|
|
if array.ndim == 4:
|
|
array = array[frame_index]
|
|
array = np.clip(array * 255.0, 0, 255).astype(np.uint8)
|
|
array = resize_image_array_to_long_edge(array, VISION_MAX_LONG_EDGE)
|
|
ensure_sessions_dir_exists()
|
|
output_png_path = CLAUDE_SESSIONS_DIR / f"frame_{uuid.uuid4().hex}.png"
|
|
Image.fromarray(array).save(str(output_png_path), format="PNG")
|
|
return output_png_path
|
|
|
|
|
|
def count_frames_in_image_tensor(image_tensor) -> int:
|
|
array = image_tensor
|
|
if hasattr(array, "cpu"):
|
|
array = array.cpu().numpy()
|
|
array = np.asarray(array)
|
|
if array.ndim == 4:
|
|
return int(array.shape[0])
|
|
return 1
|
|
|
|
|
|
def find_session_jsonl_path(session_id: str) -> Path | None:
|
|
projects_root = Path.home() / ".claude" / "projects"
|
|
if not projects_root.is_dir():
|
|
return None
|
|
target_filename = f"{session_id}.jsonl"
|
|
for dirpath, _dirnames, filenames in os.walk(projects_root):
|
|
if target_filename in filenames:
|
|
return Path(dirpath) / target_filename
|
|
return None
|
|
|
|
|
|
def extract_usage_dict_from_line_object(line_object) -> dict | None:
|
|
"""Walks common Claude Code JSONL shapes for a `usage` block with token counts."""
|
|
if not isinstance(line_object, dict):
|
|
return None
|
|
if isinstance(line_object.get("usage"), dict):
|
|
return line_object["usage"]
|
|
message_value = line_object.get("message")
|
|
if isinstance(message_value, dict) and isinstance(message_value.get("usage"), dict):
|
|
return message_value["usage"]
|
|
return None
|
|
|
|
|
|
def compute_session_cumulative_usage(session_id: str) -> dict:
|
|
"""
|
|
Sums usage across every turn in ~/.claude/projects/**/<session_id>.jsonl.
|
|
Returns a dict with input_tokens, output_tokens, cache_creation_input_tokens,
|
|
cache_read_input_tokens, total_cost_usd, turn_count, and jsonl_path.
|
|
Missing fields are treated as zero; unparseable lines are skipped.
|
|
"""
|
|
jsonl_path = find_session_jsonl_path(session_id)
|
|
cumulative_usage = {
|
|
"input_tokens": 0,
|
|
"output_tokens": 0,
|
|
"cache_creation_input_tokens": 0,
|
|
"cache_read_input_tokens": 0,
|
|
"total_cost_usd": 0.0,
|
|
"turn_count": 0,
|
|
"jsonl_path": str(jsonl_path) if jsonl_path else "",
|
|
}
|
|
if jsonl_path is None or not jsonl_path.is_file():
|
|
return cumulative_usage
|
|
|
|
try:
|
|
with jsonl_path.open("r", encoding="utf-8") as jsonl_file:
|
|
for raw_line in jsonl_file:
|
|
stripped_line = raw_line.strip()
|
|
if not stripped_line:
|
|
continue
|
|
try:
|
|
line_object = json.loads(stripped_line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
usage_dict = extract_usage_dict_from_line_object(line_object)
|
|
if usage_dict is None:
|
|
continue
|
|
cumulative_usage["input_tokens"] += int(usage_dict.get("input_tokens", 0) or 0)
|
|
cumulative_usage["output_tokens"] += int(usage_dict.get("output_tokens", 0) or 0)
|
|
cumulative_usage["cache_creation_input_tokens"] += int(usage_dict.get("cache_creation_input_tokens", 0) or 0)
|
|
cumulative_usage["cache_read_input_tokens"] += int(usage_dict.get("cache_read_input_tokens", 0) or 0)
|
|
cumulative_usage["turn_count"] += 1
|
|
line_cost = line_object.get("total_cost_usd") or line_object.get("cost_usd")
|
|
if isinstance(line_cost, (int, float)):
|
|
cumulative_usage["total_cost_usd"] += float(line_cost)
|
|
except OSError as read_error:
|
|
logger.debug("failed to read session jsonl %s: %s", jsonl_path, read_error)
|
|
|
|
return cumulative_usage
|
|
|
|
|
|
def format_session_usage_summary(session_id: str, cumulative_usage: dict) -> str:
|
|
return (
|
|
f"session {session_id}\n"
|
|
f" turns: {cumulative_usage['turn_count']}\n"
|
|
f" input_tokens: {cumulative_usage['input_tokens']:,}\n"
|
|
f" output_tokens: {cumulative_usage['output_tokens']:,}\n"
|
|
f" cache_creation_input_tokens: {cumulative_usage['cache_creation_input_tokens']:,}\n"
|
|
f" cache_read_input_tokens: {cumulative_usage['cache_read_input_tokens']:,}\n"
|
|
f" total_cost_usd: ${cumulative_usage['total_cost_usd']:.4f}\n"
|
|
f" jsonl: {cumulative_usage['jsonl_path'] or '(not found — session may not have been written yet)'}"
|
|
)
|
|
|
|
|
|
def session_jsonl_exists(session_id: str) -> bool:
|
|
projects_root = Path.home() / ".claude" / "projects"
|
|
if not projects_root.is_dir():
|
|
return False
|
|
target_filename = f"{session_id}.jsonl"
|
|
for _dirpath, _dirnames, filenames in os.walk(projects_root):
|
|
if target_filename in filenames:
|
|
return True
|
|
return False
|
|
|
|
|
|
def build_prompt_with_embedded_image_paths(user_prompt_text: str, attached_image_paths: list[Path]) -> str:
|
|
if not attached_image_paths:
|
|
return user_prompt_text
|
|
resolved_paths = [str(path.resolve()) for path in attached_image_paths]
|
|
if len(resolved_paths) == 1:
|
|
return f"@{resolved_paths[0]}\n\n{user_prompt_text}"
|
|
numbered_lines = "\n".join(
|
|
f"Image {i + 1}: @{path}" for i, path in enumerate(resolved_paths)
|
|
)
|
|
return f"{numbered_lines}\n\n{user_prompt_text}"
|
|
|
|
|
|
def build_claude_command(
|
|
session_id: str,
|
|
prompt_text: str,
|
|
system_prompt_text: str,
|
|
model_shorthand: str | None,
|
|
agent_name: str | None,
|
|
attached_image_paths: list[Path],
|
|
use_resume: bool,
|
|
use_bare: bool,
|
|
) -> list[str]:
|
|
if ATTACH_FLAG:
|
|
prompt_to_send = prompt_text
|
|
else:
|
|
prompt_to_send = build_prompt_with_embedded_image_paths(prompt_text, attached_image_paths)
|
|
|
|
command_parts: list[str] = ["claude", "-p", prompt_to_send, "--output-format", "json"]
|
|
if system_prompt_text.strip():
|
|
command_parts.extend(["--system-prompt", system_prompt_text.strip()])
|
|
if model_shorthand is not None:
|
|
command_parts.extend(["--model", model_shorthand])
|
|
if agent_name is not None:
|
|
command_parts.extend(["--agent", agent_name])
|
|
if use_bare:
|
|
command_parts.append("--bare")
|
|
else:
|
|
if use_resume:
|
|
command_parts.extend(["--resume", session_id])
|
|
else:
|
|
command_parts.extend(["--session-id", session_id])
|
|
if ATTACH_FLAG:
|
|
for image_path in attached_image_paths:
|
|
command_parts.extend([ATTACH_FLAG, str(image_path.resolve())])
|
|
return command_parts
|
|
|
|
|
|
def invoke_claude_cli_once(
|
|
command_parts: list[str], timeout_seconds: int, subprocess_env: dict | None
|
|
) -> subprocess.CompletedProcess:
|
|
ensure_sessions_dir_exists()
|
|
return subprocess.run(
|
|
command_parts,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout_seconds,
|
|
cwd=str(CLAUDE_SESSIONS_DIR),
|
|
env=subprocess_env,
|
|
)
|
|
|
|
|
|
def parse_result_text_from_envelope(stdout_text: str) -> tuple[str, bool]:
|
|
"""Returns (result_text, is_error_flag). Falls back to raw stdout on parse failure."""
|
|
try:
|
|
parsed = json.loads(stdout_text)
|
|
except json.JSONDecodeError:
|
|
return (stdout_text, False)
|
|
|
|
if isinstance(parsed, list) and parsed:
|
|
last = parsed[-1]
|
|
if isinstance(last, dict):
|
|
text_value = last.get("result") or last.get("response") or last.get("text") or ""
|
|
return (str(text_value), bool(last.get("is_error", False)))
|
|
if isinstance(parsed, dict):
|
|
text_value = parsed.get("result") or parsed.get("response") or parsed.get("text") or ""
|
|
return (str(text_value) or stdout_text, bool(parsed.get("is_error", False)))
|
|
return (stdout_text, False)
|
|
|
|
|
|
def looks_like_auth_failure(stderr_text: str, stdout_text: str) -> bool:
|
|
combined_lowercase = (stderr_text + "\n" + stdout_text).lower()
|
|
return any(marker in combined_lowercase for marker in AUTH_ERROR_SUBSTRINGS)
|
|
|
|
|
|
def call_claude_cli_with_fallbacks(
|
|
session_id: str,
|
|
prompt_text: str,
|
|
system_prompt_text: str,
|
|
model_shorthand: str | None,
|
|
agent_name: str | None,
|
|
attached_image_paths: list[Path],
|
|
prefer_resume: bool,
|
|
use_bare: bool,
|
|
timeout_seconds: int,
|
|
allow_retry_once: bool,
|
|
subprocess_env: dict | None,
|
|
) -> tuple[str, str, int, str]:
|
|
"""Returns (response_text, stderr_text, return_code, effective_session_id)."""
|
|
effective_session_id = session_id
|
|
use_resume_first = prefer_resume and session_jsonl_exists(session_id) and not use_bare
|
|
|
|
command_parts = build_claude_command(
|
|
effective_session_id, prompt_text, system_prompt_text, model_shorthand, agent_name,
|
|
attached_image_paths, use_resume=use_resume_first, use_bare=use_bare,
|
|
)
|
|
|
|
try:
|
|
completed = invoke_claude_cli_once(command_parts, timeout_seconds, subprocess_env)
|
|
except subprocess.TimeoutExpired as timeout_error:
|
|
write_session_debug_log(effective_session_id, {"stage": "timeout_initial", "command": command_parts, "err": str(timeout_error)})
|
|
return ("", f"[timeout after {timeout_seconds}s]", 124, effective_session_id)
|
|
|
|
stdout_text, stderr_text = completed.stdout, completed.stderr
|
|
|
|
if not use_bare and completed.returncode != 0:
|
|
combined_lower = (stderr_text + stdout_text).lower()
|
|
if "already in use" in combined_lower and not use_resume_first:
|
|
command_parts = build_claude_command(
|
|
effective_session_id, prompt_text, system_prompt_text, model_shorthand, agent_name,
|
|
attached_image_paths, use_resume=True, use_bare=False,
|
|
)
|
|
try:
|
|
completed = invoke_claude_cli_once(command_parts, timeout_seconds, subprocess_env)
|
|
stdout_text, stderr_text = completed.stdout, completed.stderr
|
|
except subprocess.TimeoutExpired as timeout_error:
|
|
write_session_debug_log(effective_session_id, {"stage": "timeout_resume", "err": str(timeout_error)})
|
|
return ("", f"[timeout after {timeout_seconds}s on resume]", 124, effective_session_id)
|
|
elif ("unknown session" in combined_lower or "no such session" in combined_lower) and use_resume_first:
|
|
effective_session_id = str(uuid.uuid4())
|
|
command_parts = build_claude_command(
|
|
effective_session_id, prompt_text, system_prompt_text, model_shorthand, agent_name,
|
|
attached_image_paths, use_resume=False, use_bare=False,
|
|
)
|
|
try:
|
|
completed = invoke_claude_cli_once(command_parts, timeout_seconds, subprocess_env)
|
|
stdout_text, stderr_text = completed.stdout, completed.stderr
|
|
except subprocess.TimeoutExpired as timeout_error:
|
|
write_session_debug_log(effective_session_id, {"stage": "timeout_new_session", "err": str(timeout_error)})
|
|
return ("", f"[timeout after {timeout_seconds}s on new session]", 124, effective_session_id)
|
|
|
|
if completed.returncode != 0 and allow_retry_once:
|
|
time.sleep(1.5)
|
|
try:
|
|
completed = invoke_claude_cli_once(command_parts, timeout_seconds, subprocess_env)
|
|
stdout_text, stderr_text = completed.stdout, completed.stderr
|
|
except subprocess.TimeoutExpired as timeout_error:
|
|
write_session_debug_log(effective_session_id, {"stage": "timeout_retry", "err": str(timeout_error)})
|
|
return ("", f"[timeout after {timeout_seconds}s on retry]", 124, effective_session_id)
|
|
|
|
if completed.returncode != 0:
|
|
write_session_debug_log(
|
|
effective_session_id,
|
|
{"stage": "cli_error", "rc": completed.returncode, "command": command_parts,
|
|
"stdout": stdout_text, "stderr": stderr_text},
|
|
)
|
|
logger.warning("claude CLI failed rc=%s: %s", completed.returncode, (stderr_text or stdout_text).strip()[:500])
|
|
return ("", (stderr_text.strip() or stdout_text.strip()), completed.returncode, effective_session_id)
|
|
|
|
result_text, envelope_is_error = parse_result_text_from_envelope(stdout_text)
|
|
if envelope_is_error:
|
|
write_session_debug_log(
|
|
effective_session_id,
|
|
{"stage": "envelope_error", "command": command_parts, "stdout": stdout_text, "stderr": stderr_text},
|
|
)
|
|
return (result_text, stderr_text.strip() or "[envelope is_error=true]", 1, effective_session_id)
|
|
|
|
return (result_text, stderr_text.strip(), 0, effective_session_id)
|
|
|
|
|
|
class ClaudeCliVisionText:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"system_prompt": ("STRING", {"multiline": True, "default": "You are an assistant aiding in completing user request. Be succinct and accurate."}),
|
|
"user_prompt": ("STRING", {"multiline": True, "default": "Describe the image."}),
|
|
"model": (MODEL_DROPDOWN_OPTIONS, {"default": "haiku"}),
|
|
"agent": (AGENT_DROPDOWN_OPTIONS, {"default": "general-purpose"}),
|
|
"reuse_session": ("BOOLEAN", {"default": True}),
|
|
"session_id": ("STRING", {"default": "", "tooltip": "Paste a UUID here and keep reuse_session ON to avoid the ~30k-token cold-start per call. A blank id auto-generates one — copy it from the output and paste back to reuse across runs."}),
|
|
"send_all_frames": ("BOOLEAN", {"default": False}),
|
|
"timeout_seconds": ("INT", {"default": DEFAULT_TIMEOUT_SECONDS, "min": 10, "max": 900}),
|
|
"on_error": (["empty", "raise", "retry_once"], {"default": "retry_once"}),
|
|
"is_enable": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
"api_key": ("STRING", {"default": "", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("response", "session_id", "stderr", "session_usage")
|
|
FUNCTION = "run"
|
|
CATEGORY = "claude-cli"
|
|
|
|
def run(self, system_prompt, user_prompt, model, agent, reuse_session, session_id,
|
|
send_all_frames, timeout_seconds, on_error,
|
|
is_enable, images=None, api_key=""):
|
|
if not is_enable:
|
|
return ("", session_id, "disabled", "")
|
|
|
|
if shutil.which("claude") is None:
|
|
error_message = "claude CLI not found on PATH (install Claude Code or adjust PATH)"
|
|
logger.error(error_message)
|
|
if on_error == "raise":
|
|
raise RuntimeError(error_message)
|
|
return ("", session_id, error_message, "")
|
|
|
|
normalized_api_key = normalize_optional_string_widget_value(api_key)
|
|
use_bare = bool(normalized_api_key)
|
|
if use_bare and reuse_session:
|
|
logger.warning("api_key provided → using --bare; session reuse is disabled in bare mode")
|
|
|
|
effective_session_id = coerce_to_valid_uuid_string(session_id, allow_reuse=(reuse_session and not use_bare))
|
|
|
|
model_shorthand = model # always pass-through; dropdown has no "default" option
|
|
agent_name = None if not agent or agent == "general-purpose" else agent
|
|
# general-purpose is the CLI's implicit default; omit the flag for it so
|
|
# older CLI versions without `--agent` still work.
|
|
|
|
subprocess_env: dict | None = None
|
|
if use_bare:
|
|
subprocess_env = os.environ.copy()
|
|
subprocess_env["ANTHROPIC_API_KEY"] = normalized_api_key
|
|
|
|
attached_image_paths: list[Path] = []
|
|
if images is not None:
|
|
try:
|
|
frame_count = count_frames_in_image_tensor(images)
|
|
indices_to_send = range(frame_count) if send_all_frames else [0]
|
|
for frame_index in indices_to_send:
|
|
attached_image_paths.append(write_image_tensor_frame_to_png(images, frame_index))
|
|
except Exception as image_write_error:
|
|
delete_paths_quietly(attached_image_paths)
|
|
error_message = f"failed to write image: {image_write_error}"
|
|
logger.exception("image write failed")
|
|
if on_error == "raise":
|
|
raise
|
|
return ("", effective_session_id, error_message, "")
|
|
|
|
try:
|
|
response_text, stderr_text, return_code, effective_session_id = call_claude_cli_with_fallbacks(
|
|
session_id=effective_session_id,
|
|
prompt_text=user_prompt,
|
|
system_prompt_text=system_prompt,
|
|
model_shorthand=model_shorthand,
|
|
agent_name=agent_name,
|
|
attached_image_paths=attached_image_paths,
|
|
prefer_resume=(reuse_session and not use_bare),
|
|
use_bare=use_bare,
|
|
timeout_seconds=int(timeout_seconds),
|
|
allow_retry_once=(on_error == "retry_once"),
|
|
subprocess_env=subprocess_env,
|
|
)
|
|
finally:
|
|
delete_paths_quietly(attached_image_paths)
|
|
|
|
if return_code != 0:
|
|
if looks_like_auth_failure(stderr_text, response_text):
|
|
auth_message = (
|
|
"Claude CLI is not authenticated. Run `claude login` in your terminal, "
|
|
"then retry. (Or set an api_key on this node for headless mode.)"
|
|
)
|
|
logger.error(auth_message)
|
|
raise RuntimeError(auth_message)
|
|
if on_error == "raise":
|
|
raise RuntimeError(f"claude CLI failed rc={return_code}: {stderr_text}")
|
|
|
|
cumulative_usage = compute_session_cumulative_usage(effective_session_id)
|
|
session_usage_summary = format_session_usage_summary(effective_session_id, cumulative_usage)
|
|
|
|
return (response_text, effective_session_id, stderr_text, session_usage_summary)
|
|
|
|
|
|
def normalize_optional_string_widget_value(raw_widget_value) -> str:
|
|
"""
|
|
Treats None, non-string values, empty strings, whitespace, and the literal
|
|
round-trip artifacts 'null'/'undefined'/'None' (case-insensitive) as blank.
|
|
ComfyUI + the frontend sometimes persist these as serialized strings when
|
|
an optional widget has never been touched or when JSON round-trips through
|
|
a null. Returns the stripped string or '' for any of the above.
|
|
"""
|
|
if raw_widget_value is None:
|
|
return ""
|
|
if not isinstance(raw_widget_value, str):
|
|
return ""
|
|
stripped_value = raw_widget_value.strip()
|
|
if not stripped_value:
|
|
return ""
|
|
if stripped_value.lower() in ("null", "undefined", "none"):
|
|
return ""
|
|
return stripped_value
|
|
|
|
|
|
def coerce_to_valid_uuid_string(raw_session_id: str, allow_reuse: bool) -> str:
|
|
"""
|
|
Claude CLI requires --session-id / --resume values be valid UUIDs.
|
|
- If reuse is allowed AND the widget value parses as a UUID, return its
|
|
canonical lowercase hyphenated form.
|
|
- If reuse is allowed but the value is non-empty and NOT a UUID, warn
|
|
and regenerate a fresh UUID (so user isn't blocked by a label).
|
|
- Otherwise (blank widget or reuse disabled), generate a fresh UUID.
|
|
"""
|
|
stripped_session_id = normalize_optional_string_widget_value(raw_session_id)
|
|
if stripped_session_id:
|
|
try:
|
|
return str(uuid.UUID(stripped_session_id))
|
|
except (ValueError, AttributeError, TypeError):
|
|
raise ValueError(
|
|
f"session_id must be a valid UUID (got {stripped_session_id!r}). "
|
|
"Click the 'Generate Random UUID' button on the node, or clear the field to auto-generate."
|
|
)
|
|
return str(uuid.uuid4())
|
|
|
|
|
|
def delete_paths_quietly(paths: list[Path]) -> None:
|
|
for path in paths:
|
|
try:
|
|
if path.exists():
|
|
path.unlink()
|
|
except OSError as delete_error:
|
|
logger.debug("cleanup failed for %s: %s", path, delete_error)
|
|
|
|
|
|
def parse_index_expression_to_int_list(index_expression: str, batch_length: int) -> list[int]:
|
|
"""
|
|
Parses a VHS-style indexes string into a list of ints. Supports:
|
|
- single ints: "3" → [3]
|
|
- ranges: "2:5" → [2,3,4]
|
|
- ranges with step: "0:10:2" → [0,2,4,6,8]
|
|
- negatives: "-1" → [batch_length-1]
|
|
- comma-separated combinations: "0,2,5:8"
|
|
Does NOT validate bounds — caller filters/raises as needed.
|
|
"""
|
|
if not index_expression or not index_expression.strip():
|
|
return []
|
|
result_indices: list[int] = []
|
|
for group_text in index_expression.split(","):
|
|
group_text = group_text.strip()
|
|
if not group_text:
|
|
continue
|
|
if ":" in group_text:
|
|
range_parts = [part.strip() for part in group_text.split(":", 2)]
|
|
start_value = int(range_parts[0]) if range_parts[0] else 0
|
|
end_value = int(range_parts[1]) if range_parts[1] else batch_length
|
|
step_value = int(range_parts[2]) if (len(range_parts) > 2 and range_parts[2]) else 1
|
|
if start_value < 0:
|
|
start_value = batch_length + start_value
|
|
if end_value < 0:
|
|
end_value = batch_length + end_value
|
|
result_indices.extend(range(start_value, end_value, step_value))
|
|
else:
|
|
single_value = int(group_text)
|
|
if single_value < 0:
|
|
single_value = batch_length + single_value
|
|
result_indices.append(single_value)
|
|
return result_indices
|
|
|
|
|
|
class SelectImagesBoundsSafe:
|
|
"""
|
|
Like VHS SelectImages but with `err_if_bad_index`. When False, out-of-bounds
|
|
indices are silently filtered out instead of crashing the tensor indexing.
|
|
|
|
Fixes the VHS bug where err_if_missing=False still lets OOB indices through
|
|
to the tensor index step → IndexError.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"indexes": ("STRING", {"default": "0"}),
|
|
"err_if_bad_index": ("BOOLEAN", {"default": False}),
|
|
"err_if_empty": ("BOOLEAN", {"default": False}),
|
|
"fallback_on_empty": (["first_frame", "all_frames", "blank"], {"default": "first_frame"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT")
|
|
RETURN_NAMES = ("image", "selected_count", "skipped_count")
|
|
FUNCTION = "run"
|
|
CATEGORY = "utility"
|
|
|
|
def run(self, image, indexes, err_if_bad_index, err_if_empty, fallback_on_empty):
|
|
batch_length = int(image.shape[0])
|
|
try:
|
|
requested_index_list = parse_index_expression_to_int_list(indexes, batch_length)
|
|
except ValueError as parse_error:
|
|
raise ValueError(f"SelectImagesBoundsSafe: could not parse indexes {indexes!r}: {parse_error}")
|
|
|
|
valid_index_list: list[int] = []
|
|
skipped_index_list: list[int] = []
|
|
for candidate_index in requested_index_list:
|
|
if 0 <= candidate_index < batch_length:
|
|
valid_index_list.append(candidate_index)
|
|
else:
|
|
skipped_index_list.append(candidate_index)
|
|
|
|
if skipped_index_list and err_if_bad_index:
|
|
raise IndexError(
|
|
f"SelectImagesBoundsSafe: indexes {skipped_index_list} out of range for batch of size {batch_length}. "
|
|
"Set err_if_bad_index=False to skip them silently."
|
|
)
|
|
|
|
if not valid_index_list:
|
|
if err_if_empty:
|
|
raise ValueError(
|
|
f"SelectImagesBoundsSafe: no valid indexes in {indexes!r} for batch of size {batch_length}."
|
|
)
|
|
if fallback_on_empty == "first_frame":
|
|
return (image[0:1], 0, len(skipped_index_list))
|
|
if fallback_on_empty == "all_frames":
|
|
return (image, batch_length, len(skipped_index_list))
|
|
blank_frame_tensor = image[0:1] * 0
|
|
return (blank_frame_tensor, 0, len(skipped_index_list))
|
|
|
|
selected_image_tensor = image[valid_index_list]
|
|
return (selected_image_tensor, len(valid_index_list), len(skipped_index_list))
|
|
|
|
|
|
class ImageBatchCurrentIndex:
|
|
"""
|
|
Drives ComfyUI's per-item iteration while emitting the matching index.
|
|
|
|
Input: an IMAGE batch (tensor shape [N,H,W,C]) OR a single image OR a list
|
|
of batches/images (via INPUT_IS_LIST).
|
|
Output: two parallel lists — images and indices. Downstream nodes that take
|
|
a single IMAGE and a single INT will be called N times, once per (frame, i)
|
|
pair. Use `index_offset` to get "current + 1" by adding a parallel instance
|
|
with offset=1 (or just add 1 downstream with a math node).
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"index_offset": ("INT", {"default": 0, "min": -1000000, "max": 1000000}),
|
|
}
|
|
}
|
|
|
|
INPUT_IS_LIST = True
|
|
RETURN_TYPES = ("IMAGE", "INT")
|
|
RETURN_NAMES = ("image", "index")
|
|
OUTPUT_IS_LIST = (True, True)
|
|
FUNCTION = "run"
|
|
CATEGORY = "utility"
|
|
|
|
def run(self, images, index_offset):
|
|
starting_index = int(index_offset[0]) if isinstance(index_offset, list) and index_offset else 0
|
|
flattened_image_list: list = []
|
|
for image_tensor in images:
|
|
if hasattr(image_tensor, "shape") and len(image_tensor.shape) == 4:
|
|
batch_count = int(image_tensor.shape[0])
|
|
for frame_index_in_batch in range(batch_count):
|
|
flattened_image_list.append(image_tensor[frame_index_in_batch:frame_index_in_batch + 1])
|
|
else:
|
|
flattened_image_list.append(image_tensor)
|
|
indices = list(range(starting_index, starting_index + len(flattened_image_list)))
|
|
return (flattened_image_list, indices)
|
|
|
|
|
|
def compute_indexed_grid_groups(batch_length: int, stride: int) -> list[list[int | None]]:
|
|
"""
|
|
Returns length-4 groups of frame indices. Each group entry is either an int
|
|
frame index or None (meaning "fill this cell with black"). Never repeats
|
|
real indices — black cells appear only when batch_length < 4.
|
|
|
|
- N >= 4: normal stride-based sliding window plus a backfilled final
|
|
group `[N-4 .. N-1]` if the last window didn't already end at N-1.
|
|
- N < 4: single group `[0, 1, ..., N-1, None, ..., None]`.
|
|
"""
|
|
if batch_length <= 0:
|
|
return []
|
|
if stride < 1:
|
|
stride = 1
|
|
if batch_length < 4:
|
|
return [list(range(batch_length)) + [None] * (4 - batch_length)]
|
|
groups: list[list[int | None]] = []
|
|
start_frame_index = 0
|
|
while start_frame_index + 4 <= batch_length:
|
|
groups.append([start_frame_index + offset for offset in range(4)])
|
|
start_frame_index += stride
|
|
last_covered_frame_index = groups[-1][-1]
|
|
if isinstance(last_covered_frame_index, int) and last_covered_frame_index < batch_length - 1:
|
|
groups.append([batch_length - 4 + offset for offset in range(4)])
|
|
return groups
|
|
|
|
|
|
def prepend_label_bar_to_image_tensor(
|
|
cell_image_tensor: torch.Tensor,
|
|
label_text: str,
|
|
bar_height_px: int,
|
|
font_size_px: int,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Prepends a black bar of height `bar_height_px` above `cell_image_tensor`,
|
|
with centered white label text. If `label_text` is empty, the bar is drawn
|
|
anyway (so grid geometry stays consistent across cells).
|
|
|
|
`cell_image_tensor` shape: (1, H, W, C) float 0..1. Returns (1, H + bar, W, C).
|
|
"""
|
|
_batch_size, _image_height, image_width, channel_count = cell_image_tensor.shape
|
|
label_pil = Image.new("RGB", (image_width, bar_height_px), (0, 0, 0))
|
|
if label_text:
|
|
draw = ImageDraw.Draw(label_pil)
|
|
try:
|
|
font = ImageFont.load_default(size=font_size_px) # Pillow >= 10
|
|
except TypeError:
|
|
font = ImageFont.load_default()
|
|
text_bbox = draw.textbbox((0, 0), label_text, font=font)
|
|
text_width_px = text_bbox[2] - text_bbox[0]
|
|
text_height_px = text_bbox[3] - text_bbox[1]
|
|
text_origin_x = max(0, (image_width - text_width_px) // 2)
|
|
text_origin_y = max(0, (bar_height_px - text_height_px) // 2 - text_bbox[1])
|
|
draw.text((text_origin_x, text_origin_y), label_text, fill=(255, 255, 255), font=font)
|
|
label_array = np.asarray(label_pil, dtype=np.float32) / 255.0
|
|
label_tensor = torch.from_numpy(label_array).unsqueeze(0).to(cell_image_tensor.dtype)
|
|
label_tensor = label_tensor.to(cell_image_tensor.device)
|
|
if channel_count == 4:
|
|
alpha_channel = torch.ones(
|
|
(1, bar_height_px, image_width, 1),
|
|
dtype=cell_image_tensor.dtype, device=cell_image_tensor.device,
|
|
)
|
|
label_tensor = torch.cat([label_tensor, alpha_channel], dim=-1)
|
|
return torch.cat([label_tensor, cell_image_tensor], dim=1)
|
|
|
|
|
|
def build_labeled_cell_tensor(
|
|
images_batch_tensor: torch.Tensor,
|
|
cell_entry: int | None,
|
|
label_prefix: str,
|
|
bar_height_px: int,
|
|
font_size_px: int,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Produces a single (1, H+bar, W, C) tensor for one grid cell. Real frame →
|
|
labeled with its batch index. None → black cell + blank black bar.
|
|
"""
|
|
if cell_entry is None:
|
|
black_cell_tensor = torch.zeros_like(images_batch_tensor[0:1])
|
|
return prepend_label_bar_to_image_tensor(black_cell_tensor, "", bar_height_px, font_size_px)
|
|
label_text = f"{label_prefix}{cell_entry}"
|
|
return prepend_label_bar_to_image_tensor(
|
|
images_batch_tensor[cell_entry:cell_entry + 1], label_text, bar_height_px, font_size_px
|
|
)
|
|
|
|
|
|
class SplitImageBatchToListOfImageBatches:
|
|
"""
|
|
Splits an IMAGE batch into a LIST of smaller IMAGE batches (sub-batches).
|
|
The scheduler then iterates downstream nodes once per sub-batch, with each
|
|
call receiving a multi-frame tensor of up to `sub_batch_size` frames. Also
|
|
emits the original-batch start/end indices for each sub-batch so downstream
|
|
can label or track which frames were in that call.
|
|
|
|
Example: batch of 8 frames, sub_batch_size=3, last_sub_batch_policy=keep_remainder
|
|
sub_batches[0] = images[0:3], start_index=0, end_index=2
|
|
sub_batches[1] = images[3:6], start_index=3, end_index=5
|
|
sub_batches[2] = images[6:8], start_index=6, end_index=7
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"sub_batch_size": ("INT", {"default": 2, "min": 1, "max": 1024}),
|
|
"stride": ("INT", {"default": 0, "min": 0, "max": 1024, "tooltip": "0 = non-overlapping (same as sub_batch_size). 1 = sliding window. Any other value = custom stride."}),
|
|
"last_sub_batch_policy": (["keep_remainder", "drop_remainder", "pad_with_last"], {"default": "keep_remainder"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT", "INT")
|
|
RETURN_NAMES = ("sub_batches", "start_index", "end_index", "sub_batch_count")
|
|
OUTPUT_IS_LIST = (True, True, True, False)
|
|
FUNCTION = "run"
|
|
CATEGORY = "utility"
|
|
|
|
def run(self, images, sub_batch_size, stride, last_sub_batch_policy):
|
|
if images.ndim != 4:
|
|
raise ValueError(
|
|
f"SplitImageBatchToListOfImageBatches: expected IMAGE tensor shape (N,H,W,C), got {tuple(images.shape)}"
|
|
)
|
|
batch_length = int(images.shape[0])
|
|
effective_sub_batch_size = max(1, int(sub_batch_size))
|
|
effective_stride = effective_sub_batch_size if int(stride) <= 0 else int(stride)
|
|
|
|
if batch_length == 0:
|
|
return ([], [], [], 0)
|
|
|
|
sub_batch_tensor_list: list = []
|
|
start_index_list: list[int] = []
|
|
end_index_list: list[int] = []
|
|
|
|
start_position = 0
|
|
while start_position < batch_length:
|
|
end_position_exclusive = start_position + effective_sub_batch_size
|
|
if end_position_exclusive > batch_length:
|
|
remainder_frame_count = batch_length - start_position
|
|
if last_sub_batch_policy == "drop_remainder" and remainder_frame_count < effective_sub_batch_size:
|
|
break
|
|
if last_sub_batch_policy == "pad_with_last":
|
|
last_frame_tensor = images[batch_length - 1:batch_length]
|
|
missing_frame_count = effective_sub_batch_size - remainder_frame_count
|
|
padded_sub_batch_tensor = torch.cat(
|
|
[images[start_position:batch_length]] + [last_frame_tensor] * missing_frame_count,
|
|
dim=0,
|
|
)
|
|
sub_batch_tensor_list.append(padded_sub_batch_tensor)
|
|
start_index_list.append(start_position)
|
|
end_index_list.append(batch_length - 1)
|
|
else: # keep_remainder
|
|
sub_batch_tensor_list.append(images[start_position:batch_length])
|
|
start_index_list.append(start_position)
|
|
end_index_list.append(batch_length - 1)
|
|
break
|
|
sub_batch_tensor_list.append(images[start_position:end_position_exclusive])
|
|
start_index_list.append(start_position)
|
|
end_index_list.append(end_position_exclusive - 1)
|
|
start_position += effective_stride
|
|
|
|
return (sub_batch_tensor_list, start_index_list, end_index_list, len(sub_batch_tensor_list))
|
|
|
|
|
|
class IndexedGrid2x2FromBatch:
|
|
"""
|
|
Batch → batch of 2x2 grid images with per-cell index labels.
|
|
|
|
Example (N=8, stride=3): produces 3 grid images from frames
|
|
(0,1,2,3), (3,4,5,6), (4,5,6,7). The final grid is backfilled to always
|
|
contain 4 real frames when N >= 4. When N < 4, a single grid is produced
|
|
with the missing cells filled with black (real frames never repeat).
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"stride": ("INT", {"default": 3, "min": 1, "max": 4}),
|
|
"label_prefix": ("STRING", {"default": ""}),
|
|
"label_font_size": ("INT", {"default": 16, "min": 6, "max": 128}),
|
|
"label_bar_height": ("INT", {"default": 24, "min": 0, "max": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "STRING")
|
|
RETURN_NAMES = ("grids", "grid_count", "index_groups_debug")
|
|
FUNCTION = "run"
|
|
CATEGORY = "utility"
|
|
|
|
def run(self, images, stride, label_prefix, label_font_size, label_bar_height):
|
|
if images.ndim != 4:
|
|
raise ValueError(
|
|
f"IndexedGrid2x2FromBatch: expected IMAGE tensor shape (N,H,W,C), got {tuple(images.shape)}"
|
|
)
|
|
batch_length = int(images.shape[0])
|
|
index_groups = compute_indexed_grid_groups(batch_length, int(stride))
|
|
if not index_groups:
|
|
return (images[0:0], 0, "")
|
|
|
|
grid_image_tensor_list: list[torch.Tensor] = []
|
|
for group_of_four in index_groups:
|
|
top_left_cell = build_labeled_cell_tensor(images, group_of_four[0], label_prefix, label_bar_height, label_font_size)
|
|
top_right_cell = build_labeled_cell_tensor(images, group_of_four[1], label_prefix, label_bar_height, label_font_size)
|
|
bottom_left_cell = build_labeled_cell_tensor(images, group_of_four[2], label_prefix, label_bar_height, label_font_size)
|
|
bottom_right_cell = build_labeled_cell_tensor(images, group_of_four[3], label_prefix, label_bar_height, label_font_size)
|
|
top_row_tensor = torch.cat([top_left_cell, top_right_cell], dim=2)
|
|
bottom_row_tensor = torch.cat([bottom_left_cell, bottom_right_cell], dim=2)
|
|
grid_image_tensor_list.append(torch.cat([top_row_tensor, bottom_row_tensor], dim=1))
|
|
|
|
grids_batch_tensor = torch.cat(grid_image_tensor_list, dim=0)
|
|
debug_string = " | ".join(
|
|
",".join("b" if entry is None else str(entry) for entry in group) for group in index_groups
|
|
)
|
|
return (grids_batch_tensor, len(index_groups), debug_string)
|
|
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ClaudeCliVisionText": ClaudeCliVisionText,
|
|
"ImageBatchCurrentIndex": ImageBatchCurrentIndex,
|
|
"SelectImagesBoundsSafe": SelectImagesBoundsSafe,
|
|
"IndexedGrid2x2FromBatch": IndexedGrid2x2FromBatch,
|
|
"SplitImageBatchToListOfImageBatches": SplitImageBatchToListOfImageBatches,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ClaudeCliVisionText": "Claude CLI (Vision+Text)",
|
|
"ImageBatchCurrentIndex": "Image Batch Current Index",
|
|
"SelectImagesBoundsSafe": "Select Images (Bounds Safe)",
|
|
"IndexedGrid2x2FromBatch": "Indexed 2x2 Grid from Batch",
|
|
"SplitImageBatchToListOfImageBatches": "Split Image Batch to List of Image Batches",
|
|
}
|