commit 9736f62bd2a6b7ff9d75e9ad4fdaa626571b4a69 Author: jeremie louvaert Date: Fri Mar 20 14:51:27 2026 +0100 Initial release: production-hardened API optimization nodes for ComfyUI Four custom nodes for cost tracking, deterministic caching, and lazy execution bypass when using external APIs (Kling, Magnific, RunPod, etc.). Key production features: - Full-content tensor hashing (no lossy approximations) - Decimal arithmetic for financial precision - File locking for concurrent access safety - Atomic writes to prevent corruption - Cache TTL support - Transaction audit logging - Device-portable tensor serialization Co-Authored-By: Claude Opus 4.6 (1M context) diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..97962a0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +*.pyc +*.pyo +*.egg-info/ +dist/ +build/ +*.lock +.venv/ +.env diff --git a/README.md b/README.md new file mode 100644 index 0000000..de2e91e --- /dev/null +++ b/README.md @@ -0,0 +1,84 @@ +# ComfyUI API Optimizer + +A suite of production-grade custom nodes for ComfyUI designed for workflows that rely on external remote APIs (Kling 3.0, Magnific, Banana.dev, RunPod, etc.). + +When your compute moves to the cloud, your bottlenecks shift from VRAM limitations to **API Costs, Latency, and Serialization**. This custom node pack solves these problems natively inside ComfyUI. + +## Included Nodes + +### 1. API Cost & Quota Tracker + +Acts as a circuit-breaker for your wallet. Pass your prompt or image through this node before it hits your API node. + +- **Budget Enforcement:** Set a `$ Budget Limit` and `$ Cost Per Run`. A persistent ledger tracks all charges. If the next run would exceed the budget, execution halts *before* the API is charged. +- **Precise Arithmetic:** Uses `decimal.Decimal` internally — no floating-point drift on large batch runs. +- **Transaction Audit Log:** Every charge is appended to `api_transactions.jsonl` with timestamps for full traceability. +- **Safe Resets:** Resetting the budget archives the previous ledger instead of discarding it. +- **Concurrent-Safe:** File locking prevents corruption when multiple ComfyUI instances share the same output directory. + +### 2. The Deterministic Hash Vault Suite (3 Nodes) + +ComfyUI's native caching often breaks with external API nodes (dynamic timestamps, non-deterministic seeds). The Hash Vault is an aggressive disk-caching layer that strictly hashes your prompt, parameters, and input tensors. + +- **Hash Vault (Check Cache):** Hashes your inputs and checks local disk for a cached result. +- **Lazy API Switch:** Uses ComfyUI's `{"lazy": True}` evaluation engine. On a cache hit, this switch **physically prevents** the upstream API node from executing — saving money and time. +- **Hash Vault (Save Result):** Writes new API outputs to the vault for future cache hits. + +#### Key Features + +- **Full-Content Tensor Hashing:** Hashes the complete byte representation of tensors (including dtype and shape metadata) — no lossy approximations. +- **Recursive Hashing:** Correctly handles nested data structures (dicts, lists, tuples) common in ComfyUI latents and conditioning. +- **Cache TTL:** Optional time-to-live for cache entries. Expired entries are automatically removed and treated as cache misses. Set to `0` for entries that never expire. +- **Device-Portable Caching:** All tensors are saved to CPU and loaded with `map_location="cpu"`, so cache files work regardless of GPU configuration. +- **Atomic Writes:** Cache files are written to a temp file first, then atomically replaced — preventing corruption from interrupted writes. +- **Concurrent-Safe:** File locking on every cache read/write operation. + +## How to Use the Hash Vault + +To properly bypass an API node, sandwich it with the vault nodes: + +1. Connect your Prompt/Image to **Check Cache**. +2. Connect the `is_cached` output to the **Lazy API Switch**. +3. Connect your Prompt/Image to your actual API Node. +4. Connect the output of your API Node to **Save Result** (using the `hash_key` from step 1). +5. Connect both the `cached_data` (from step 1) and the `api_data` (from step 4) to the **Lazy API Switch**. + +``` + ┌─────────────┐ + ┌───────────►│ API Node ├──► 💾 Save Result ──┐ + │ └─────────────┘ │ + Prompt/Image─┤ │ + │ ┌─────────────┐ ▼ + └───────────►│ 🔍 Check ├──────────────► 🔀 Lazy Switch ──► Output + │ Cache │ is_cached ▲ + └──┬──────────┘ │ + │ cached_data ────────────────────┘ +``` + +## Output Files + +All data is stored under your ComfyUI output directory: + +| Path | Description | +|------|-------------| +| `output/api_metrics/api_costs.json` | Current cost ledger (per-provider totals) | +| `output/api_metrics/api_transactions.jsonl` | Append-only audit log with timestamps | +| `output/api_metrics/api_costs_archive_*.json` | Archived ledgers from budget resets | +| `output/hash_vault/*.pt` | Cached API outputs (PyTorch format) | + +## Installation + +Clone this repository into your `ComfyUI/custom_nodes/` directory: + +```bash +cd ComfyUI/custom_nodes/ +git clone https://github.com/YOUR_USERNAME/ComfyUI-API-Optimizer.git +pip install -r ComfyUI-API-Optimizer/requirements.txt +``` + +Restart ComfyUI. + +### Dependencies + +- **PyTorch** — already present in any ComfyUI installation +- **filelock** — typically already installed as a transitive dependency of PyTorch/HuggingFace. If not, `pip install filelock`. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..eb1fd92 --- /dev/null +++ b/__init__.py @@ -0,0 +1,5 @@ +from .api_optimizer_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ =['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + +print("\033[34m[ComfyUI-API-Optimizer] \033[92mLoaded Cost Tracker & Deterministic Hash Vault\033[0m") \ No newline at end of file diff --git a/api_optimizer_nodes.py b/api_optimizer_nodes.py new file mode 100644 index 0000000..cd6b506 --- /dev/null +++ b/api_optimizer_nodes.py @@ -0,0 +1,339 @@ +import os +import json +import hashlib +import time +import torch +import folder_paths +from decimal import Decimal +from datetime import datetime, timezone + +try: + from filelock import FileLock +except ImportError: + # filelock is typically available in ComfyUI environments (transitive dep of torch). + # This no-op fallback ensures nodes still work without it, just without concurrency safety. + class FileLock: + def __init__(self, lock_file, timeout=-1): + pass + def __enter__(self): + return self + def __exit__(self, *args): + pass + +# ------------------------------------------------------------------------ +# UTILITIES +# ------------------------------------------------------------------------ +class AnyType(str): + """Universal type that passes ComfyUI's type-checking by never reporting inequality.""" + def __ne__(self, __value: object) -> bool: + return False + +any_type = AnyType("*") + +# ------------------------------------------------------------------------ +# NODE 1: API Cost & Quota Tracker +# ------------------------------------------------------------------------ +class APICostTracker: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "passthrough": (any_type,), + "api_provider": ("STRING", {"default": "Kling 3.0"}), + "cost_per_run": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 1000.0, "step": 0.001}), + "budget_limit": ("FLOAT", {"default": 10.0, "min": 0.0, "max": 10000.0, "step": 0.1}), + "reset_budget": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = (any_type, "STRING") + RETURN_NAMES = ("passthrough", "cost_summary") + FUNCTION = "track_cost" + CATEGORY = "API Optimization" + + def track_cost(self, passthrough, api_provider, cost_per_run, budget_limit, reset_budget): + # Convert floats to Decimal at the boundary for precise financial arithmetic + cost = Decimal(str(cost_per_run)) + limit = Decimal(str(budget_limit)) + + if cost > limit: + raise ValueError( + f"[API Cost Tracker] cost_per_run (${cost}) exceeds budget_limit (${limit}). " + f"Fix your node configuration." + ) + + log_dir = os.path.join(folder_paths.get_output_directory(), "api_metrics") + os.makedirs(log_dir, exist_ok=True) + log_file = os.path.join(log_dir, "api_costs.json") + lock_file = log_file + ".lock" + tx_log_file = os.path.join(log_dir, "api_transactions.jsonl") + + with FileLock(lock_file, timeout=10): + costs = {} + + if os.path.exists(log_file) and not reset_budget: + try: + with open(log_file, "r") as f: + raw = json.load(f) + costs = {k: Decimal(str(v)) for k, v in raw.items()} + except json.JSONDecodeError as e: + print(f"⚠️ [API Cost Tracker] Corrupted ledger, backing up and resetting. Error: {e}") + backup_path = log_file + f".corrupt.{int(time.time())}" + os.replace(log_file, backup_path) + costs = {} + except (OSError, IOError) as e: + print(f"⚠️ [API Cost Tracker] Failed to read ledger: {e}") + costs = {} + + if reset_budget: + if costs: + archive_path = os.path.join(log_dir, f"api_costs_archive_{int(time.time())}.json") + with open(archive_path, "w") as f: + json.dump({k: str(v) for k, v in costs.items()}, f, indent=4) + print(f"📋 [API Cost Tracker] Budget reset. Previous ledger archived.") + costs = {} + + current_total = sum(costs.values(), Decimal("0")) + + # CIRCUIT BREAKER: Stop execution before the API is charged + if current_total + cost > limit: + remaining = limit - current_total + raise ValueError( + f"\n[🛑 API Cost Tracker] BUDGET EXCEEDED!\n" + f"Budget Limit: ${limit}\n" + f"Total Spent: ${current_total}\n" + f"This Run: ${cost}\n" + f"Remaining: ${remaining}\n" + f"Halting execution to prevent unauthorized API charges." + ) + + # Update ledger + if api_provider not in costs: + costs[api_provider] = Decimal("0") + costs[api_provider] += cost + + # Atomic write: temp file then replace + tmp_file = log_file + ".tmp" + with open(tmp_file, "w") as f: + json.dump({k: str(v) for k, v in costs.items()}, f, indent=4) + os.replace(tmp_file, log_file) + + # Append to transaction audit log + new_total = sum(costs.values(), Decimal("0")) + tx = { + "timestamp": datetime.now(timezone.utc).isoformat(), + "provider": api_provider, + "cost": str(cost), + "running_total": str(new_total), + } + with open(tx_log_file, "a") as f: + f.write(json.dumps(tx) + "\n") + + remaining = limit - new_total + summary = f"Total: ${new_total} | Remaining: ${remaining}" + print(f"💰 [API Cost Tracker] Billed ${cost} to {api_provider}. {summary}") + + return (passthrough, summary) + +# ------------------------------------------------------------------------ +# NODE 2: Deterministic Hash Vault (Check Cache) +# ------------------------------------------------------------------------ +class DeterministicHashVault: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "payload_string": ("STRING", {"forceInput": True, "tooltip": "Connect your prompt or JSON params here"}), + }, + "optional": { + "any_input": (any_type, {"tooltip": "Optional: Connect init image or latent to factor into hash"}), + "cache_ttl_hours": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 8760.0, "step": 1.0, + "tooltip": "Cache time-to-live in hours. 0 = never expires"}), + } + } + + RETURN_TYPES = (any_type, "INT", "STRING") + RETURN_NAMES = ("cached_data", "is_cached", "hash_key") + FUNCTION = "check_vault" + CATEGORY = "API Optimization" + + def _hash_tensor(self, hash_obj, tensor): + """Hash full tensor content deterministically including dtype and shape metadata.""" + t = tensor.contiguous().cpu() + hash_obj.update(f"__tensor:{t.dtype}:{list(t.shape)}:".encode("utf-8")) + try: + hash_obj.update(t.numpy().tobytes()) + except Exception: + # Fallback for non-standard tensor types that can't convert to numpy + hash_obj.update(bytes(t.untyped_storage())) + + def _hash_value(self, hash_obj, value): + """Recursively hash any value — handles tensors, dicts, lists, and primitives.""" + if isinstance(value, torch.Tensor): + self._hash_tensor(hash_obj, value) + elif isinstance(value, dict): + hash_obj.update(b"__dict:") + for key in sorted(value.keys()): + hash_obj.update(str(key).encode("utf-8")) + self._hash_value(hash_obj, value[key]) + elif isinstance(value, (list, tuple)): + hash_obj.update(b"__seq:") + for item in value: + self._hash_value(hash_obj, item) + else: + hash_obj.update(str(value).encode("utf-8")) + + def check_vault(self, payload_string, any_input=None, cache_ttl_hours=0.0): + hash_obj = hashlib.sha256() + hash_obj.update(str(payload_string).encode("utf-8")) + + if any_input is not None: + self._hash_value(hash_obj, any_input) + + hash_key = hash_obj.hexdigest() + + vault_dir = os.path.join(folder_paths.get_output_directory(), "hash_vault") + os.makedirs(vault_dir, exist_ok=True) + file_path = os.path.join(vault_dir, f"{hash_key}.pt") + lock_path = file_path + ".lock" + + with FileLock(lock_path, timeout=10): + if os.path.exists(file_path): + # Check TTL if set + if cache_ttl_hours > 0: + file_age_hours = (time.time() - os.path.getmtime(file_path)) / 3600 + if file_age_hours > cache_ttl_hours: + print( + f"⏰ [API Vault] Cache expired for {hash_key[:8]} " + f"(age: {file_age_hours:.1f}h > TTL: {cache_ttl_hours:.1f}h)" + ) + os.remove(file_path) + return (None, 0, hash_key) + + try: + # weights_only=False is required because cache may contain dicts/lists + # alongside tensors. Safe here — we only load files our own save node wrote. + cached_data = torch.load(file_path, map_location="cpu", weights_only=False) + print(f"🟢 [API Vault] Cache Hit! Hash: {hash_key[:8]}") + return (cached_data, 1, hash_key) + except Exception as e: + print(f"⚠️ [API Vault] Corrupted cache for {hash_key[:8]}, removing. Error: {e}") + try: + os.remove(file_path) + except OSError: + pass + + print(f"🔴 [API Vault] Cache Miss! Hash: {hash_key[:8]}") + return (None, 0, hash_key) + +# ------------------------------------------------------------------------ +# NODE 3: Hash Vault (Save API Result) +# ------------------------------------------------------------------------ +class HashVaultSave: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "hash_key": ("STRING", {"forceInput": True}), + "api_output": (any_type,), + } + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("api_output",) + FUNCTION = "save_to_vault" + CATEGORY = "API Optimization" + + # NOTE: OUTPUT_NODE must be False. We only want this to execute if the Switch demands it! + + def _to_cpu(self, data): + """Recursively move all tensors to CPU for portable serialization.""" + if isinstance(data, torch.Tensor): + return data.cpu() + elif isinstance(data, dict): + return {k: self._to_cpu(v) for k, v in data.items()} + elif isinstance(data, list): + return [self._to_cpu(item) for item in data] + elif isinstance(data, tuple): + return tuple(self._to_cpu(item) for item in data) + return data + + def save_to_vault(self, hash_key, api_output): + vault_dir = os.path.join(folder_paths.get_output_directory(), "hash_vault") + os.makedirs(vault_dir, exist_ok=True) + file_path = os.path.join(vault_dir, f"{hash_key}.pt") + lock_path = file_path + ".lock" + tmp_path = file_path + ".tmp" + + # Move to CPU for portable serialization + cpu_data = self._to_cpu(api_output) + + with FileLock(lock_path, timeout=10): + # Atomic write: save to temp then replace + torch.save(cpu_data, tmp_path) + os.replace(tmp_path, file_path) + + print(f"💾 [API Vault] Saved API output to Vault: {hash_key[:8]}") + return (api_output,) + +# ------------------------------------------------------------------------ +# NODE 4: Lazy API Switch (The Bypass Engine) +# ------------------------------------------------------------------------ +class LazyAPISwitch: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "is_cached": ("INT",), + }, + "optional": { + # {"lazy": True} is the magic engine feature. ComfyUI will NOT execute + # upstream nodes connected to these sockets unless check_lazy_status demands it. + "cached_data": (any_type, {"lazy": True}), + "api_data": (any_type, {"lazy": True}), + } + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("final_output",) + FUNCTION = "switch" + CATEGORY = "API Optimization" + + def check_lazy_status(self, is_cached, cached_data=None, api_data=None): + """Tell ComfyUI which inputs to actually evaluate. + + If cache hit → demand only cached_data, leaving the API branch completely idle. + If cache miss → demand only api_data, triggering the API execution branch. + """ + if is_cached == 1: + if cached_data is None: + return ["cached_data"] + else: + if api_data is None: + return ["api_data"] + return [] + + def switch(self, is_cached, cached_data=None, api_data=None): + if is_cached == 1: + print("🔀 [Lazy Switch] Cache Hit routed! API execution completely bypassed.") + return (cached_data,) + else: + print("🔀 [Lazy Switch] Cache Miss routed! API was executed.") + return (api_data,) + +# ------------------------------------------------------------------------ +# MAPPINGS: Registering the nodes with ComfyUI +# ------------------------------------------------------------------------ +NODE_CLASS_MAPPINGS = { + "APICostTracker": APICostTracker, + "DeterministicHashVault": DeterministicHashVault, + "HashVaultSave": HashVaultSave, + "LazyAPISwitch": LazyAPISwitch, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "APICostTracker": "💰 API Cost & Quota Tracker", + "DeterministicHashVault": "🔍 Hash Vault (Check Cache)", + "HashVaultSave": "💾 Hash Vault (Save API Result)", + "LazyAPISwitch": "🔀 Lazy API Switch (Bypass)", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..83c2e35 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +filelock