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) <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+*.pyc
+*.pyo
+*.egg-info/
+dist/
+build/
+*.lock
+.venv/
+.env

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`.

__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")

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)",
+}

requirements.txt

@@ -0,0 +1 @@
+filelock