diodiogod
f285fa8bd2
Technical details: - Add model_mmap_residency fallback to the shared ComfyUI wrapper - Add pinned memory, dtype, patcher, and device compatibility methods - Return full model memory for non-mmap-backed engines instead of a zero stub - Add regression coverage for wrapper compatibility behavior Fixes #276 Fixes #278
104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
"""
|
|
Unit tests for ComfyUIModelWrapper compatibility with newer ComfyUI contracts.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
# Add custom node root to path BEFORE any project imports
|
|
custom_node_root = Path(__file__).parent.parent.parent
|
|
sys.path.insert(0, str(custom_node_root))
|
|
|
|
# Set up minimal environment to avoid ComfyUI imports
|
|
os.environ.setdefault("COMFYUI_TESTING", "1")
|
|
|
|
from utils.models.comfyui_model_wrapper.base_wrapper import ComfyUIModelWrapper, ModelInfo
|
|
|
|
|
|
class DummyModel:
|
|
"""Minimal wrapped model for ComfyUI wrapper tests."""
|
|
|
|
|
|
class TrackingModel:
|
|
"""Wrapped model that exposes ComfyUI memory hooks."""
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def model_mmap_residency(self, free=False):
|
|
self.calls.append(("mmap", free))
|
|
if free:
|
|
return 12, 34
|
|
return 56, 78
|
|
|
|
def pinned_memory_size(self):
|
|
self.calls.append(("pinned", None))
|
|
return 90
|
|
|
|
def lowvram_patch_counter(self):
|
|
self.calls.append(("lowvram", None))
|
|
return 3
|
|
|
|
|
|
def build_wrapper(model=None, memory_size=4096, device="cpu"):
|
|
wrapped_model = model or DummyModel()
|
|
model_info = ModelInfo(
|
|
model=wrapped_model,
|
|
model_type="tts",
|
|
engine="test_engine",
|
|
device=device,
|
|
memory_size=memory_size,
|
|
load_device=device,
|
|
)
|
|
return ComfyUIModelWrapper(wrapped_model, model_info)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestComfyUIModelWrapperCompatibility:
|
|
def test_model_mmap_residency_falls_back_to_full_model_size(self):
|
|
wrapper = build_wrapper(memory_size=12345)
|
|
assert wrapper.model_mmap_residency() == (0, 12345)
|
|
|
|
def test_model_mmap_residency_forwards_wrapped_model(self):
|
|
wrapped = TrackingModel()
|
|
wrapper = build_wrapper(model=wrapped)
|
|
|
|
assert wrapper.model_mmap_residency() == (56, 78)
|
|
assert wrapper.model_mmap_residency(free=True) == (12, 34)
|
|
assert wrapped.calls[:2] == [("mmap", False), ("mmap", True)]
|
|
|
|
def test_safe_defaults_for_pinned_and_lowvram_metrics(self):
|
|
wrapper = build_wrapper(memory_size=2048)
|
|
|
|
assert wrapper.pinned_memory_size() == 0
|
|
assert wrapper.lowvram_patch_counter() == 0
|
|
assert wrapper.get_ram_usage() == 2048
|
|
|
|
def test_forwarded_metrics_use_wrapped_model_hooks(self):
|
|
wrapped = TrackingModel()
|
|
wrapper = build_wrapper(model=wrapped)
|
|
|
|
assert wrapper.pinned_memory_size() == 90
|
|
assert wrapper.lowvram_patch_counter() == 3
|
|
assert ("pinned", None) in wrapped.calls
|
|
assert ("lowvram", None) in wrapped.calls
|
|
|
|
def test_current_loaded_device_returns_torch_device(self):
|
|
wrapper = build_wrapper(device="cpu")
|
|
assert wrapper.current_loaded_device() == torch.device("cpu")
|
|
|
|
def test_model_patches_to_accepts_device_and_dtype_targets(self):
|
|
wrapper = build_wrapper()
|
|
|
|
wrapper.model_patches_to(torch.device("cpu"))
|
|
wrapper.model_patches_to(torch.float16)
|
|
|
|
assert wrapper.load_device == torch.device("cpu")
|
|
assert wrapper.current_loaded_device() == torch.device("cpu")
|
|
assert wrapper.model_dtype() == torch.float16
|
|
assert wrapper.model_patches_models() == ()
|