Limbicnation-ComfyUI-PromptGenerator/docs/VRAM_OPTIMIZATION_PLAN.md

VRAM-Aware Ollama Lifecycle — Execution Plan

Project: ComfyUI-PromptGenerator VRAM Optimization
Date: 2026-06-23
RTX 4090 | 24 GB VRAM | PyTorch 2.9.0+cu129 | ComfyUI 0.25.0
Status: IMPLEMENTED — All changes passing (79/79 tests)

---

Executive Summary

The ComfyUI PromptRefinerNode (and sibling nodes: PromptGeneratorNode, NegativePromptNode, PromptDualStreamRefinerNode) calls local LLMs via Ollama for text generation. After each node completes, Ollama's default behavior retains the loaded model in GPU VRAM for 5 minutes (keep_alive=5m). When the ComfyUI pipeline subsequently attempts to load VRAM-heavy diffusion models (e.g., LTXAVTEModel_ text encoder at 11.2 GB staged), CUDA OOM occurs:

RuntimeError: VRAM grow failed: 2013757440 bytes

Root Cause: Ollama model VRAM residency is invisible to PyTorch's memory allocator. After 3 sequential Ollama calls (qwen3:8b → qwen3-4b-deforum-prompt:v7 → qwen3:8b), approximately 4–7 GB of VRAM remains occupied by the Ollama runtime. The diffusion pipeline's text encoder then requests ~1.9 GB for weight casting, which fails because the available 17.5 GB free minus the invisible Ollama residency leaves insufficient contiguous VRAM.

Solution: Three-layer VRAM release strategy:

1. keep_alive="0s" on every ollama.generate() call — forces immediate model eviction
2. torch.cuda.empty_cache() in a background thread — releases PyTorch allocator fragments
3. Explicit unload_model() fallback — safety net for subprocess or error paths

Key Change: Every node now passes keep_alive="0s" to Ollama and runs async VRAM cleanup in a finally block, ensuring cleanup even on exceptions.

---

Architectural VRAM Bottleneck Analysis

Current State (Before Fix)

Timeline (from ComfyUI error logs):
03:38:21 — PromptRefiner calls qwen3:8b          → Ollama loads ~4.7 GB into VRAM
03:38:27 — PromptGenerator calls qwen3-4b-deforum → Ollama loads ~2.6 GB (qwen3:8b evicted, swapped)
03:38:30 — PromptRefiner calls qwen3:8b again     → Ollama swaps models again
03:38:36 — NegativePrompt calls qwen3:8b          → Ollama keeps model resident
03:38:52 — NegativePrompt completes               → Ollama model STAYS in VRAM (keep_alive=5m default)
03:38:54 — ComfyUI loads LTXAVTEModel_ (11.2 GB)  → Tries to allocate 1.9 GB for weight cast
03:38:55 — ❌ VRAM grow failed: 2013757440 bytes   → CUDA OOM!

VRAM accounting at failure:

Component	VRAM Usage	PyTorch-Visible?
Ollama runtime (qwen3:8b)	~4.7 GB	No — managed by Ollama's Go runtime
ComfyUI VAE/CLIP staging	~6.0 GB	Yes — tracked by comfy_aimdo
LTXAVTEModel_ weight cast buffer	1.9 GB (requested)	Yes — attempted allocation
Total needed	~12.6 GB
Available	~17.5 GB free	But Ollama-held VRAM fragmented

Target State (After Fix)

Timeline (with autounload):
03:38:21 — PromptRefiner calls qwen3:8b with keep_alive="0s"
03:38:27 — Model evicted from VRAM immediately, torch.cuda.empty_cache() runs async
03:38:27 — PromptGenerator calls qwen3-4b-deforum with keep_alive="0s"
03:38:30 — Model evicted, CUDA cache cleared
03:38:30 — PromptRefiner calls qwen3:8b with keep_alive="0s"
03:38:35 — Model evicted, CUDA cache cleared
03:38:36 — NegativePrompt calls qwen3:8b with keep_alive="0s"
03:38:52 — Model evicted, CUDA cache cleared
03:38:54 — ComfyUI loads LTXAVTEModel_ (11.2 GB) → Full 23.5 GB available
03:38:55 — ✅ Weight cast succeeds — 1.9 GB allocated cleanly

VRAM accounting after fix:

Component	VRAM Usage	PyTorch-Visible?
Ollama runtime	0 GB — model evicted	N/A
ComfyUI VAE/CLIP staging	~6.0 GB	Yes
LTXAVTEModel_ weight cast	1.9 GB (succeeds)	Yes
Total needed	~7.9 GB
Available	~23.5 GB	Full GPU available

---

Refactoring Blueprint for prompt_refiner_node.py

Structural Changes

1. OllamaClient (nodes/adapters/ollama_client.py)

New methods added:

Method	Purpose	Threading
generate_streaming(keep_alive=...)	New keep_alive kwarg forwarded to ollama.generate()	Synchronous (blocking)
unload_model(model)	Sends ollama.generate(keep_alive=0) to evict model	Synchronous
release_vram()	gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect()	Synchronous
cleanup(model, unload, release_cuda)	Combines unload + release_vram	Synchronous
cleanup_async(model, ...)	Runs cleanup() in a daemon thread	Non-blocking

Key API payload — Ollama keep_alive parameter:

# Streaming generation with immediate unload
stream = ollama.generate(
    model="qwen3:8b",
    prompt="...",
    stream=True,
    options={"temperature": 0.7, "top_p": 0.9},
    keep_alive="0s",  # ← Evict model from VRAM after this request
)

# Explicit unload (safety net)
ollama.generate(
    model="qwen3:8b",
    prompt="",
    keep_alive=0,           # ← 0 and "0s" are equivalent
    options={"num_predict": 1},
)

Ollama REST API equivalent (for direct HTTP calls):

POST http://127.0.0.1:11434/api/generate
{
    "model": "qwen3:8b",
    "prompt": "...",
    "stream": false,
    "keep_alive": "0s",
    "options": {
        "temperature": 0.7,
        "top_p": 0.9
    }
}

2. PromptRefinerNode (nodes/prompt_refiner_node.py)

Changes:

• Added keep_alive="0s" to client.generate_streaming() call
• Added unload_model boolean input (default: True) — user-configurable via ComfyUI UI
• Wrapped entire refine() body in try/finally — async VRAM cleanup always runs
• Cleanup uses cleanup_async(unload=False, release_cuda=True) — keep_alive="0s" already handles Ollama-side; only PyTorch CUDA cache needs explicit clearing

3. PromptGeneratorNode (nodes/prompt_generator_node.py)

Changes:

• Added keep_alive="0s" to client.generate_streaming() call
• Wrapped generate() body in try/finally with cleanup_async()
• Subprocess fallback path also covered by the finally block

4. NegativePromptNode (nodes/negative_prompt_node.py)

Changes:

• Added keep_alive="0s" to client.generate_streaming() call
• Wrapped generate_negative() body in try/finally with cleanup_async()

5. PromptDualStreamRefinerNode (nodes/prompt_dual_stream_refiner_node.py)

Changes:

• Added keep_alive="0s" to client.generate_streaming() call
• Wrapped refine() body in try/finally with cleanup_async()

---

Ollama Autounload Implementation Guide

keep_alive Parameter Behavior

Value	Behavior	Use Case
"0s" or 0	Unload immediately after request	Shared GPU (recommended)
"5m" (default)	Keep loaded 5 minutes	Single-model GPU
"1h"	Keep loaded 1 hour	Dedicated LLM server
-1	Never unload	Persistent serving

Exact API Payloads

Python (ollama package)

# In generate_streaming() — every streaming call now includes keep_alive
stream = ollama.generate(
    model=model,
    prompt=prompt,
    stream=True,
    options={"temperature": temperature, "top_p": top_p},
    keep_alive="0s",
)

REST API (curl)

# Generate with immediate unload
curl -s http://127.0.0.1:11434/api/generate -d '{
    "model": "qwen3:8b",
    "prompt": "Refine this prompt...",
    "stream": false,
    "keep_alive": "0s",
    "options": {"temperature": 0.5, "top_p": 0.9}
}'

# Explicit unload (after all generation complete)
curl -s http://127.0.0.1:11434/api/generate -d '{
    "model": "qwen3:8b",
    "prompt": "",
    "keep_alive": 0,
    "options": {"num_predict": 1}
}'

Asynchronous Cleanup Pattern

# Non-blocking: returns immediately, cleanup runs in background daemon thread
cleanup_thread = OllamaClient.cleanup_async(
    model="qwen3:8b",
    logger_prefix="PromptRefiner",
    unload=False,       # keep_alive="0s" already evicted the model
    release_cuda=True,  # torch.cuda.empty_cache() + gc.collect()
)
# ComfyUI queue is NOT blocked — next node can start immediately

---

Fallback & Error Handling Strategy

Error Classification in OllamaClient

StreamResult.kind	Cause	Node Response
ok	Success	Return refined text
timeout	Per-chunk or total timeout	Try subprocess fallback
transient	Network/connection issue	Try subprocess fallback
model_crash	llama runner died (HTTP 500)	Surface actionable error, skip subprocess
server_error	Other non-2xx response	Surface error, skip subprocess
unavailable	ollama package not installed	Surface installation instructions

Cleanup Guarantee

All nodes use try/finally to ensure VRAM cleanup runs even on exceptions:

try:
    result = client.generate_streaming(...)
    # ... process result ...
finally:
    # Always runs — even on KeyboardInterrupt, SystemExit, or unhandled exceptions
    OllamaClient.cleanup_async(model=model, unload=False, release_cuda=True)

Ollama Unload Failure Handling

unload_model() catches all exceptions and logs them without re-raising:

def unload_model(self, model: str) -> bool:
    try:
        ollama.generate(model=model, prompt="", keep_alive=0, options={"num_predict": 1})
        return True
    except ConnectionError:
        self._log("Could not connect to unload model")  # Non-fatal
        return False
    except Exception:
        self._log("Failed to unload model")              # Non-fatal
        return False

CUDA Cache Release Safety

release_vram() gracefully handles environments without CUDA:

@staticmethod
def release_vram() -> None:
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    except ImportError:
        pass  # No torch — no-op

---

Validation & Success Metrics

Immediate Verification

1. Monitor Ollama VRAM with ollama ps

# Before pipeline execution — should show no models loaded
ollama ps

# After PromptRefiner completes — should show NO models (keep_alive="0s")
ollama ps
# Expected: empty output (model evicted)

2. Monitor CUDA VRAM with nvidia-smi

# Watch VRAM in real-time during pipeline execution
watch -n 0.5 nvidia-smi --query-gpu=memory.used,memory.free --format=csv

# Expected behavior:
# Phase 1 (Ollama calls): VRAM spikes to ~5-7 GB
# Phase 2 (post-cleanup):  VRAM drops back to ~0.5 GB (baseline)
# Phase 3 (diffusion):     VRAM rises to ~18-20 GB (no OOM)

3. PyTorch CUDA Memory Stats

import torch
print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Reserved:  {torch.cuda.memory_reserved() / 1e9:.2f} GB")
print(f"Max alloc: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB")

4. Automated Test Suite

# Run all unit tests (79 tests covering OllamaClient, all nodes)
python3 -m pytest tests/ -v

# Expected: 79 passed, 0 failed

Success Criteria Checklist

Criterion	Verification Method	Status
Ollama VRAM drops to 0 after node execution	ollama ps returns empty	PASS
torch.cuda.empty_cache() called post-execution	Background thread in finally block	PASS
keep_alive="0s" on every ollama.generate() call	Code review of all 4 nodes	PASS
Non-blocking cleanup (doesn't stall ComfyUI queue)	cleanup_async() uses daemon thread	PASS
Error handling prevents pipeline crash on Ollama failure	try/finally + exception swallowing in unload_model()	PASS
Configurable model names (no hardcoded strings)	model param in INPUT_TYPES, no string literals	PASS
unload_model toggle in PromptRefiner UI	BOOLEAN input with label "Unload After Use"	PASS
All existing tests pass	pytest tests/ — 79/79	PASS
Diffusion model loads without OOM after Ollama calls	Full pipeline test on RTX 4090	PASS (verified in logs)

---

File Change Summary

File	Lines	Changes
nodes/adapters/ollama_client.py	488 (+102)	Added keep_alive kwarg, unload_model(), release_vram(), cleanup(), cleanup_async()
nodes/prompt_refiner_node.py	228 (+30)	Added keep_alive="0s", unload_model input, try/finally cleanup
nodes/prompt_generator_node.py	552 (+11)	Added keep_alive="0s", try/finally cleanup
nodes/negative_prompt_node.py	184 (+10)	Added keep_alive="0s", try/finally cleanup
nodes/prompt_dual_stream_refiner_node.py	220 (+11)	Added keep_alive="0s", try/finally cleanup
tests/unit/test_prompt_refiner.py	67 (±2)	Updated mock signatures for keep_alive kwarg