228 lines
7.7 KiB
Python
228 lines
7.7 KiB
Python
import contextlib
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
def _load_node_package(monkeypatch, tmp_path):
|
|
comfy = types.ModuleType("comfy")
|
|
comfy.__path__ = []
|
|
management = types.ModuleType("comfy.model_management")
|
|
management.load_calls = []
|
|
management.load_models_gpu = lambda models, force_full_load=False: (
|
|
management.load_calls.append((models, force_full_load))
|
|
)
|
|
patcher = types.ModuleType("comfy.model_patcher")
|
|
utils = types.ModuleType("comfy.utils")
|
|
utils.ProgressBar = lambda total: types.SimpleNamespace(update=lambda value: None)
|
|
folder_paths = types.ModuleType("folder_paths")
|
|
folder_paths.models_dir = str(tmp_path)
|
|
folder_paths._paths = []
|
|
folder_paths.add_model_folder_path = (
|
|
lambda name, path, is_default=False: folder_paths._paths.append(path)
|
|
if path not in folder_paths._paths
|
|
else None
|
|
)
|
|
folder_paths.get_folder_paths = lambda name: folder_paths._paths
|
|
|
|
monkeypatch.setitem(sys.modules, "comfy", comfy)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", management)
|
|
monkeypatch.setitem(sys.modules, "comfy.model_patcher", patcher)
|
|
monkeypatch.setitem(sys.modules, "comfy.utils", utils)
|
|
monkeypatch.setitem(sys.modules, "folder_paths", folder_paths)
|
|
|
|
root = Path(__file__).parents[1]
|
|
package_name = "comfyui_sam_audio_node_test"
|
|
spec = importlib.util.spec_from_file_location(
|
|
package_name,
|
|
root / "__init__.py",
|
|
submodule_search_locations=[str(root)],
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
monkeypatch.setitem(sys.modules, package_name, module)
|
|
spec.loader.exec_module(module)
|
|
return module, management
|
|
|
|
|
|
class FakeBatch:
|
|
def __init__(self, audios, descriptions, anchors, masked_videos):
|
|
self.audios = audios
|
|
self.descriptions = descriptions
|
|
self.anchors = anchors
|
|
self.masked_videos = masked_videos
|
|
|
|
def to(self, device):
|
|
return self
|
|
|
|
|
|
class FakeProcessor:
|
|
audio_sampling_rate = 8
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def __call__(self, audios, descriptions, anchors=None, masked_videos=None):
|
|
self.calls.append(
|
|
{
|
|
"audios": audios,
|
|
"descriptions": descriptions,
|
|
"anchors": anchors,
|
|
"masked_videos": masked_videos,
|
|
}
|
|
)
|
|
return FakeBatch(audios, descriptions, anchors, masked_videos)
|
|
|
|
def mask_videos(self, videos, masks):
|
|
return [
|
|
video * mask.eq(0) for video, mask in zip(videos, masks, strict=False)
|
|
]
|
|
|
|
|
|
class FakeModel(torch.nn.Module):
|
|
def separate(self, batch, **kwargs):
|
|
target = [row.mean(0) for row in batch.audios]
|
|
return types.SimpleNamespace(target=target, residual=[-row for row in target])
|
|
|
|
|
|
def test_text_node_runs_full_audio_batch_with_standard_shapes(monkeypatch, tmp_path):
|
|
package, management = _load_node_package(monkeypatch, tmp_path)
|
|
nodes = sys.modules[f"{package.__name__}.nodes"]
|
|
runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
|
|
pipeline = runtime.SAMAudioPipeline(
|
|
model=FakeModel(),
|
|
processor=FakeProcessor(),
|
|
patcher=object(),
|
|
device=torch.device("cpu"),
|
|
source="fake",
|
|
)
|
|
|
|
target, residual = nodes.SAMAudioTextSeparator().separate(
|
|
pipeline=pipeline,
|
|
audio={"waveform": torch.ones(2, 2, 8), "sample_rate": 8},
|
|
description="test sound",
|
|
predict_spans=False,
|
|
seed=7,
|
|
inference_steps=32,
|
|
)
|
|
|
|
assert target["waveform"].shape == (2, 1, 8)
|
|
assert residual["waveform"].shape == (2, 1, 8)
|
|
assert target["sample_rate"] == 8
|
|
loader_inputs = package.NODE_CLASS_MAPPINGS[
|
|
"SAMAudioPipelineLoader"
|
|
].INPUT_TYPES()["required"]
|
|
assert loader_inputs["attention_backend"][0] == ["pytorch", "comfy_kitchen"]
|
|
sampler_inputs = package.NODE_CLASS_MAPPINGS[
|
|
"SAMAudioTextSeparator"
|
|
].INPUT_TYPES()["required"]
|
|
assert sampler_inputs["seed"][1]["control_after_generate"] is True
|
|
assert sampler_inputs["chunk_duration"][1]["default"] == 10.0
|
|
assert sampler_inputs["chunk_overlap"][1]["default"] == 1.0
|
|
assert set(package.NODE_CLASS_MAPPINGS) == {
|
|
"SAMAudioPipelineLoader",
|
|
"SAMAudioSpanPrompt",
|
|
"SAMAudioTextSeparator",
|
|
"SAMAudioSpanSeparator",
|
|
"SAMAudioVisualSeparator",
|
|
}
|
|
assert management.load_calls == [([pipeline.patcher], True)]
|
|
|
|
|
|
def _fake_pipeline(runtime, processor=None):
|
|
return runtime.SAMAudioPipeline(
|
|
model=FakeModel(),
|
|
processor=processor or FakeProcessor(),
|
|
patcher=object(),
|
|
device=torch.device("cpu"),
|
|
source="fake",
|
|
)
|
|
|
|
|
|
def test_text_node_chunks_crossfades_and_advances_seed(monkeypatch, tmp_path):
|
|
package, _ = _load_node_package(monkeypatch, tmp_path)
|
|
nodes = sys.modules[f"{package.__name__}.nodes"]
|
|
runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
|
|
processor = FakeProcessor()
|
|
pipeline = _fake_pipeline(runtime, processor)
|
|
seeds = []
|
|
|
|
@contextlib.contextmanager
|
|
def record_seed(seed, device):
|
|
seeds.append(seed)
|
|
yield
|
|
|
|
monkeypatch.setattr(nodes, "seeded_inference", record_seed)
|
|
target, residual = nodes.SAMAudioTextSeparator().separate(
|
|
pipeline=pipeline,
|
|
audio={"waveform": torch.ones(2, 2, 18), "sample_rate": 8},
|
|
description="test sound",
|
|
predict_spans=False,
|
|
seed=7,
|
|
inference_steps=32,
|
|
chunk_duration=0.75,
|
|
chunk_overlap=0.25,
|
|
)
|
|
|
|
assert [call["audios"][0].shape[-1] for call in processor.calls] == [6, 6, 6, 6]
|
|
assert seeds == [7, 8, 9, 10]
|
|
assert torch.equal(target["waveform"], torch.ones(2, 1, 18))
|
|
assert torch.equal(residual["waveform"], -torch.ones(2, 1, 18))
|
|
|
|
|
|
def test_span_node_localizes_global_spans_for_each_chunk(monkeypatch, tmp_path):
|
|
package, _ = _load_node_package(monkeypatch, tmp_path)
|
|
nodes = sys.modules[f"{package.__name__}.nodes"]
|
|
runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
|
|
processor = FakeProcessor()
|
|
pipeline = _fake_pipeline(runtime, processor)
|
|
spans = (
|
|
runtime.SpanPrompt("+", 0.5, 1.25),
|
|
runtime.SpanPrompt("-", 1.5, 1.75),
|
|
)
|
|
|
|
nodes.SAMAudioSpanSeparator().separate(
|
|
pipeline=pipeline,
|
|
audio={"waveform": torch.ones(1, 1, 16), "sample_rate": 8},
|
|
description="test sound",
|
|
spans=spans,
|
|
seed=0,
|
|
inference_steps=32,
|
|
chunk_duration=1.0,
|
|
chunk_overlap=0.25,
|
|
)
|
|
|
|
anchors = [call["anchors"][0] for call in processor.calls]
|
|
assert anchors[0] == [runtime.SpanPrompt("+", 0.5, 1.0)]
|
|
assert anchors[1] == [
|
|
runtime.SpanPrompt("+", 0.0, 0.5),
|
|
runtime.SpanPrompt("-", 0.75, 1.0),
|
|
]
|
|
assert anchors[2] == [runtime.SpanPrompt("-", 0.0, 0.25)]
|
|
|
|
|
|
def test_visual_node_slices_frames_to_each_audio_chunk(monkeypatch, tmp_path):
|
|
package, _ = _load_node_package(monkeypatch, tmp_path)
|
|
nodes = sys.modules[f"{package.__name__}.nodes"]
|
|
runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
|
|
processor = FakeProcessor()
|
|
pipeline = _fake_pipeline(runtime, processor)
|
|
images = torch.linspace(0, 1, 8).view(8, 1, 1, 1).expand(-1, 2, 2, 3)
|
|
|
|
nodes.SAMAudioVisualSeparator().separate(
|
|
pipeline=pipeline,
|
|
audio={"waveform": torch.ones(1, 1, 16), "sample_rate": 8},
|
|
images=images,
|
|
mask=torch.zeros(2, 2),
|
|
description="",
|
|
seed=0,
|
|
inference_steps=32,
|
|
chunk_duration=1.0,
|
|
chunk_overlap=0.25,
|
|
)
|
|
|
|
frame_counts = [call["masked_videos"][0].shape[0] for call in processor.calls]
|
|
assert frame_counts == [4, 4, 2]
|