limbicnation
f0aaf1314a
Pass keep_alive="0s" on every Ollama generate call so models are evicted from GPU VRAM immediately instead of lingering 5 minutes, which caused CUDA OOM when downstream diffusion models loaded. Each node now runs async VRAM cleanup in a try/finally block. The cleanup uses unload=True so it also evicts a model loaded via the subprocess fallback path (which carries the default keep_alive). Add unload_model(), release_vram(), cleanup() and cleanup_async() to OllamaClient, plus a per-node unload toggle on PromptRefiner.
68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
"""Unit tests for PromptRefinerNode seed behavior."""
|
|
|
|
from unittest.mock import patch
|
|
|
|
from nodes.adapters.ollama_client import StreamResult
|
|
from nodes.prompt_refiner_node import PromptRefinerNode
|
|
|
|
|
|
class TestPromptRefinerSeed:
|
|
"""Test suite for per-pass seed derivation."""
|
|
|
|
def test_per_pass_seed_increments(self):
|
|
"""Each pass should receive an incremented seed when seed != -1."""
|
|
node = PromptRefinerNode()
|
|
captured_seeds: list[int | None] = []
|
|
|
|
def _capture_seed(*, model, prompt, temperature, top_p, timeout, pbar, seed, keep_alive=None):
|
|
captured_seeds.append(seed)
|
|
return StreamResult(text=f"refined with seed={seed}", kind="ok")
|
|
|
|
with (
|
|
patch.object(node, "REFINEMENT_PROMPT", "{prompt}"),
|
|
patch(
|
|
"nodes.prompt_refiner_node.OllamaClient.generate_streaming",
|
|
side_effect=_capture_seed,
|
|
),
|
|
):
|
|
result = node.refine(
|
|
prompt="a forest",
|
|
model="qwen3:8b",
|
|
passes=3,
|
|
seed=42,
|
|
temperature=0.5,
|
|
top_p=0.9,
|
|
timeout=30,
|
|
)
|
|
|
|
assert captured_seeds == [42, 43, 44]
|
|
assert "refined with seed=44" in result[0]
|
|
|
|
def test_random_seed_none_for_all_passes(self):
|
|
"""When seed is -1, all passes should receive None (random)."""
|
|
node = PromptRefinerNode()
|
|
captured_seeds: list[int | None] = []
|
|
|
|
def _capture_seed(*, model, prompt, temperature, top_p, timeout, pbar, seed, keep_alive=None):
|
|
captured_seeds.append(seed)
|
|
return StreamResult(text="refined", kind="ok")
|
|
|
|
with (
|
|
patch.object(node, "REFINEMENT_PROMPT", "{prompt}"),
|
|
patch(
|
|
"nodes.prompt_refiner_node.OllamaClient.generate_streaming",
|
|
side_effect=_capture_seed,
|
|
),
|
|
):
|
|
node.refine(
|
|
prompt="a forest",
|
|
model="qwen3:8b",
|
|
passes=2,
|
|
seed=-1,
|
|
temperature=0.5,
|
|
top_p=0.9,
|
|
timeout=30,
|
|
)
|
|
|
|
assert captured_seeds == [None, None]
|