mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-31 03:17:23 +08:00
fc4a135c04
- Reduce SeedVR2 coverage down to production unit tests - Route SeedVR2 7B through Comfy varlength attention - Disable SeedVR2 RoPE cache reuse after the upstream DynamicVRAM change
71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
import importlib
|
|
import inspect
|
|
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import torch
|
|
|
|
from comfy.cli_args import args as cli_args
|
|
|
|
if not torch.cuda.is_available():
|
|
cli_args.cpu = True
|
|
|
|
import comfy_extras.nodes_seedvr as nodes_seedvr # noqa: E402
|
|
|
|
|
|
def test_resize_simple_multiplier_resolves_upscaled_shorter_edge():
|
|
images = torch.zeros(1, 3, 16, 20, 3)
|
|
|
|
output = nodes_seedvr.SeedVR2Resize.execute(images, 4.0)
|
|
|
|
input_pixels, original_image, upscaled_shorter_edge = output.result
|
|
assert tuple(input_pixels.shape) == (1, 5, 64, 80, 3)
|
|
assert input_pixels.min().item() == 0.0
|
|
assert input_pixels.max().item() == 0.0
|
|
assert original_image is images
|
|
assert upscaled_shorter_edge == 64
|
|
|
|
|
|
def test_seedvr_node_signature_matches_schema():
|
|
mock_mm = MagicMock()
|
|
mock_mm.xformers_enabled.return_value = False
|
|
mock_mm.xformers_enabled_vae.return_value = False
|
|
mock_mm.sage_attention_enabled.return_value = False
|
|
mock_mm.flash_attention_enabled.return_value = False
|
|
|
|
sentinel = object()
|
|
prior_cpu = cli_args.cpu
|
|
cli_args.cpu = True
|
|
prior_module = sys.modules.get("comfy_extras.nodes_seedvr", sentinel)
|
|
comfy_pkg = sys.modules.get("comfy")
|
|
prior_mm_attr = getattr(comfy_pkg, "model_management", sentinel) if comfy_pkg else sentinel
|
|
|
|
with patch.dict(sys.modules, {"comfy.model_management": mock_mm}):
|
|
if comfy_pkg is not None:
|
|
setattr(comfy_pkg, "model_management", mock_mm)
|
|
sys.modules.pop("comfy_extras.nodes_seedvr", None)
|
|
try:
|
|
nodes_seedvr = importlib.import_module("comfy_extras.nodes_seedvr")
|
|
for node_cls in (nodes_seedvr.SeedVR2Resize, nodes_seedvr.SeedVR2ResizeAdvanced):
|
|
schema_ids = [i.id for i in node_cls.define_schema().inputs]
|
|
exec_params = [
|
|
p for p in inspect.signature(node_cls.execute).parameters.keys()
|
|
if p != "cls"
|
|
]
|
|
assert schema_ids == exec_params, (
|
|
f"{node_cls.__name__} schema/execute drift: "
|
|
f"schema_ids={schema_ids}, exec_params={exec_params}"
|
|
)
|
|
finally:
|
|
cli_args.cpu = prior_cpu
|
|
if prior_module is sentinel:
|
|
sys.modules.pop("comfy_extras.nodes_seedvr", None)
|
|
else:
|
|
sys.modules["comfy_extras.nodes_seedvr"] = prior_module
|
|
if comfy_pkg is not None:
|
|
if prior_mm_attr is sentinel:
|
|
if hasattr(comfy_pkg, "model_management"):
|
|
delattr(comfy_pkg, "model_management")
|
|
else:
|
|
setattr(comfy_pkg, "model_management", prior_mm_attr)
|