diff --git a/README.md b/README.md index 1802240..3843d8e 100644 --- a/README.md +++ b/README.md @@ -58,11 +58,14 @@ All nodes are in `audio/SAM-Audio`. 3. Connect both to a separator and enter a short prompt such as `man speaking` or `dog barking`. 4. Preview or save the `target` and `residual` outputs. +Long audio is processed in overlapping chunks by default. Set `chunk_duration` to `0` for a single full-clip pass. + Example workflows are included in the [`examples`](examples/) directory. ## Notes - Output is mono at 48 kHz. +- Adjacent chunks use a normalized crossfade; the defaults are 10-second chunks with 1-second overlap. - White pixels in a visual mask identify the target object. - Local model folders must contain `config.json` and `checkpoint.pt`. - The model loader offers PyTorch SDPA (default) or Comfy Kitchen INT8 attention. Comfy Kitchen requires a supported current ComfyUI build and GPU and may slightly change results because its attention is quantized. diff --git a/examples/span_separation_workflow.json b/examples/span_separation_workflow.json index 88347cf..65920f3 100644 --- a/examples/span_separation_workflow.json +++ b/examples/span_separation_workflow.json @@ -65,7 +65,7 @@ {"name": "residual", "type": "AUDIO", "links": [5], "slot_index": 1} ], "properties": {"Node name for S&R": "SAMAudioSpanSeparator"}, - "widgets_values": ["car honking", 0, "fixed", 32] + "widgets_values": ["car honking", 0, "fixed", 32, 10.0, 1.0] }, { "id": 5, diff --git a/examples/text_separation_workflow.json b/examples/text_separation_workflow.json index a980428..0998557 100644 --- a/examples/text_separation_workflow.json +++ b/examples/text_separation_workflow.json @@ -49,7 +49,7 @@ {"name": "residual", "type": "AUDIO", "links": [4], "slot_index": 1} ], "properties": {"Node name for S&R": "SAMAudioTextSeparator"}, - "widgets_values": ["man speaking", false, 0, "fixed", 32] + "widgets_values": ["man speaking", false, 0, "fixed", 32, 10.0, 1.0] }, { "id": 4, diff --git a/examples/visual_separation_workflow.json b/examples/visual_separation_workflow.json index a2ab336..aa2baba 100644 --- a/examples/visual_separation_workflow.json +++ b/examples/visual_separation_workflow.json @@ -67,7 +67,7 @@ {"name": "residual", "type": "AUDIO", "links": [6], "slot_index": 1} ], "properties": {"Node name for S&R": "SAMAudioVisualSeparator"}, - "widgets_values": ["", 0, "fixed", 32] + "widgets_values": ["", 0, "fixed", 32, 10.0, 1.0] }, { "id": 5, diff --git a/nodes.py b/nodes.py index 72a67ef..9183d80 100644 --- a/nodes.py +++ b/nodes.py @@ -5,14 +5,19 @@ from typing import Any import comfy.model_management as model_management import comfy.model_patcher as model_patcher +import comfy.utils as comfy_utils import folder_paths import torch from .sam_audio_comfy.attention import ATTENTION_BACKENDS, attention_backend_context from .sam_audio_comfy.audio import ( + AudioChunk, + crossfade_audio_chunks, + plan_audio_chunks, prepare_audio, prepare_visual_prompt, result_to_audio, + slice_visual_frames, ) from .sam_audio_comfy.runtime import ( SAMAudioPipeline, @@ -58,6 +63,31 @@ def _sampling_inputs() -> dict[str, tuple[Any, ...]]: } +def _chunking_inputs() -> dict[str, tuple[Any, ...]]: + return { + "chunk_duration": ( + "FLOAT", + { + "default": 10.0, + "min": 0.0, + "max": 3600.0, + "step": 0.5, + "tooltip": "Seconds processed per pass. Use 0 to process the entire clip at once.", + }, + ), + "chunk_overlap": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 60.0, + "step": 0.1, + "tooltip": "Seconds shared by adjacent chunks for a smooth crossfade.", + }, + ), + } + + def _as_result_list(value: Any) -> list[torch.Tensor]: if isinstance(value, torch.Tensor): return [value] if value.ndim == 1 else list(value) @@ -66,12 +96,30 @@ def _as_result_list(value: Any) -> list[torch.Tensor]: raise RuntimeError(f"SAM-Audio returned an unexpected result type: {type(value).__name__}") +def _spans_for_chunk( + spans: tuple[SpanPrompt, ...], chunk: AudioChunk, sample_rate: int +) -> tuple[SpanPrompt, ...]: + chunk_start = chunk.start / sample_rate + chunk_end = chunk.end / sample_rate + localized = [] + for span in spans: + start = max(span.start, chunk_start) + end = min(span.end, chunk_end) + if end > start: + localized.append( + SpanPrompt(span.token, start - chunk_start, end - chunk_start) + ) + return tuple(localized) + + def _run_separation( pipeline: SAMAudioPipeline, audio: dict[str, Any], description: str, seed: int, inference_steps: int, + chunk_duration: float, + chunk_overlap: float, *, predict_spans: bool = False, spans: tuple[SpanPrompt, ...] | None = None, @@ -87,52 +135,84 @@ def _run_separation( waveforms, target_length = prepare_audio( audio, pipeline.processor.audio_sampling_rate ) + sample_rate = pipeline.processor.audio_sampling_rate batch_size = len(waveforms) - duration = target_length / pipeline.processor.audio_sampling_rate + duration = target_length / sample_rate + chunks = plan_audio_chunks( + target_length, + sample_rate, + chunk_duration, + chunk_overlap, + ) - anchors = None if spans is not None: validate_spans(spans, duration) - anchors = [list(spans) for _ in range(batch_size)] - masked_videos = None + masked_video = None if visual_prompt is not None: frames, mask = visual_prompt - masked = pipeline.processor.mask_videos([frames], [mask])[0] - masked_videos = [masked for _ in range(batch_size)] - - batch = pipeline.processor( - audios=waveforms, - descriptions=[description for _ in range(batch_size)], - anchors=anchors, - masked_videos=masked_videos, - ) + masked_video = pipeline.processor.mask_videos([frames], [mask])[0] pipeline.load(model_management) - batch = batch.to(pipeline.device) + target_chunks = [] + residual_chunks = [] + progress = comfy_utils.ProgressBar(len(chunks)) + for chunk_index, chunk in enumerate(chunks): + chunk_waveforms = [ + waveform[..., chunk.start : chunk.end] for waveform in waveforms + ] - with ( - seeded_inference(seed, pipeline.device), - attention_backend_context(pipeline.attention_backend), - torch.inference_mode(), - ): - result = pipeline.model.separate( - batch, - ode_opt=sampling_options, - reranking_candidates=1, - predict_spans=predict_spans, + anchors = None + if spans is not None: + localized_spans = _spans_for_chunk(spans, chunk, sample_rate) + anchors = [list(localized_spans) for _ in range(batch_size)] + + masked_videos = None + if masked_video is not None: + chunk_video = slice_visual_frames(masked_video, chunk, target_length) + masked_videos = [chunk_video for _ in range(batch_size)] + + batch = pipeline.processor( + audios=chunk_waveforms, + descriptions=[description for _ in range(batch_size)], + anchors=anchors, + masked_videos=masked_videos, ) + batch = batch.to(pipeline.device) - target = result_to_audio( - _as_result_list(result.target), - pipeline.processor.audio_sampling_rate, - target_length, - ) - residual = result_to_audio( - _as_result_list(result.residual), - pipeline.processor.audio_sampling_rate, - target_length, - ) + chunk_seed = (seed + chunk_index) % (2**64) + with ( + seeded_inference(chunk_seed, pipeline.device), + attention_backend_context(pipeline.attention_backend), + torch.inference_mode(), + ): + result = pipeline.model.separate( + batch, + ode_opt=sampling_options, + reranking_candidates=1, + predict_spans=predict_spans, + ) + + target = result_to_audio( + _as_result_list(result.target), sample_rate, chunk.length + ) + residual = result_to_audio( + _as_result_list(result.residual), sample_rate, chunk.length + ) + target_chunks.append((chunk, target["waveform"])) + residual_chunks.append((chunk, residual["waveform"])) + progress.update(1) + + target = { + "waveform": crossfade_audio_chunks(target_chunks, target_length).contiguous(), + "sample_rate": sample_rate, + } + residual = { + "waveform": crossfade_audio_chunks( + residual_chunks, target_length + ).contiguous(), + "sample_rate": sample_rate, + } return target, residual @@ -262,6 +342,7 @@ class SAMAudioTextSeparator: ), } required.update(_sampling_inputs()) + required.update(_chunking_inputs()) return {"required": required} RETURN_TYPES = ("AUDIO", "AUDIO") @@ -277,6 +358,8 @@ class SAMAudioTextSeparator: predict_spans: bool, seed: int, inference_steps: int, + chunk_duration: float = 10.0, + chunk_overlap: float = 1.0, ): return _run_separation( pipeline, @@ -284,6 +367,8 @@ class SAMAudioTextSeparator: description, seed, inference_steps, + chunk_duration, + chunk_overlap, predict_spans=predict_spans, ) @@ -306,6 +391,7 @@ class SAMAudioSpanSeparator: "spans": (SPANS_TYPE,), } required.update(_sampling_inputs()) + required.update(_chunking_inputs()) return {"required": required} RETURN_TYPES = ("AUDIO", "AUDIO") @@ -321,6 +407,8 @@ class SAMAudioSpanSeparator: spans: tuple[SpanPrompt, ...], seed: int, inference_steps: int, + chunk_duration: float = 10.0, + chunk_overlap: float = 1.0, ): return _run_separation( pipeline, @@ -328,6 +416,8 @@ class SAMAudioSpanSeparator: description, seed, inference_steps, + chunk_duration, + chunk_overlap, spans=tuple(spans), ) @@ -361,6 +451,7 @@ class SAMAudioVisualSeparator: ), } required.update(_sampling_inputs()) + required.update(_chunking_inputs()) return {"required": required} RETURN_TYPES = ("AUDIO", "AUDIO") @@ -377,6 +468,8 @@ class SAMAudioVisualSeparator: description: str, seed: int, inference_steps: int, + chunk_duration: float = 10.0, + chunk_overlap: float = 1.0, ): visual_prompt = prepare_visual_prompt(images, mask) return _run_separation( @@ -385,6 +478,8 @@ class SAMAudioVisualSeparator: description, seed, inference_steps, + chunk_duration, + chunk_overlap, visual_prompt=visual_prompt, ) diff --git a/pyproject.toml b/pyproject.toml index 9ab841a..7149e85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "comfyui-sam-audio" -version = "0.1.0" +version = "0.2.0" description = "Native ComfyUI nodes for local Meta SAM-Audio sound separation" readme = "README.md" requires-python = ">=3.11" diff --git a/sam_audio_comfy/audio.py b/sam_audio_comfy/audio.py index 525a2f9..8e7fd25 100644 --- a/sam_audio_comfy/audio.py +++ b/sam_audio_comfy/audio.py @@ -1,12 +1,122 @@ from __future__ import annotations +import math from collections.abc import Mapping, Sequence +from dataclasses import dataclass import torch import torch.nn.functional as F import torchaudio +@dataclass(frozen=True) +class AudioChunk: + start: int + end: int + + @property + def length(self) -> int: + return self.end - self.start + + +def plan_audio_chunks( + total_samples: int, + sample_rate: int, + chunk_duration: float, + chunk_overlap: float, +) -> list[AudioChunk]: + """Plan overlapping sample ranges; zero duration keeps the original full pass.""" + if total_samples < 1: + raise ValueError("total_samples must be positive") + if sample_rate < 1: + raise ValueError("sample_rate must be positive") + if chunk_duration < 0: + raise ValueError("chunk_duration cannot be negative") + if chunk_overlap < 0: + raise ValueError("chunk_overlap cannot be negative") + if chunk_duration == 0: + return [AudioChunk(0, total_samples)] + + chunk_samples = round(chunk_duration * sample_rate) + overlap_samples = round(chunk_overlap * sample_rate) + if chunk_samples < 1: + raise ValueError("chunk_duration is shorter than one audio sample") + if overlap_samples >= chunk_samples: + raise ValueError("chunk_overlap must be shorter than chunk_duration") + if total_samples <= chunk_samples: + return [AudioChunk(0, total_samples)] + + hop_samples = chunk_samples - overlap_samples + chunks = [] + start = 0 + while True: + end = min(start + chunk_samples, total_samples) + chunks.append(AudioChunk(start, end)) + if end == total_samples: + return chunks + start += hop_samples + + +def crossfade_audio_chunks( + chunks: Sequence[tuple[AudioChunk, torch.Tensor]], total_samples: int +) -> torch.Tensor: + """Reassemble [batch, channels, samples] chunks with normalized linear fades.""" + if not chunks: + raise ValueError("chunks cannot be empty") + + first = chunks[0][1] + if first.ndim != 3: + raise ValueError("chunk waveforms must have [batch, channels, samples] shape") + output = first.new_zeros((*first.shape[:-1], total_samples)) + weights = first.new_zeros(total_samples) + + for index, (chunk, waveform) in enumerate(chunks): + if waveform.ndim != 3 or waveform.shape[:-1] != first.shape[:-1]: + raise ValueError("all chunk waveforms must have matching batch and channels") + if waveform.shape[-1] != chunk.length: + raise ValueError("chunk waveform length does not match its sample range") + + window = waveform.new_ones(chunk.length) + if index > 0: + fade_in = max(0, chunks[index - 1][0].end - chunk.start) + if fade_in: + window[:fade_in] *= torch.linspace( + 0, 1, fade_in + 2, dtype=window.dtype, device=window.device + )[1:-1] + if index + 1 < len(chunks): + fade_out = max(0, chunk.end - chunks[index + 1][0].start) + if fade_out: + window[-fade_out:] *= torch.linspace( + 1, 0, fade_out + 2, dtype=window.dtype, device=window.device + )[1:-1] + + output[..., chunk.start : chunk.end] += waveform * window + weights[chunk.start : chunk.end] += window + + if torch.any(weights == 0): + raise RuntimeError("chunk plan left uncovered audio samples") + return output / weights + + +def slice_visual_frames( + frames: torch.Tensor, chunk: AudioChunk, total_samples: int +) -> torch.Tensor: + """Select frames whose uniformly spaced time range intersects an audio chunk.""" + if frames.ndim < 1 or frames.shape[0] < 1: + raise ValueError("frames must contain at least one frame") + if total_samples < 1: + raise ValueError("total_samples must be positive") + if frames.shape[0] == 1 or chunk == AudioChunk(0, total_samples): + return frames + + frame_count = frames.shape[0] + start = math.floor(chunk.start * frame_count / total_samples) + end = math.ceil(chunk.end * frame_count / total_samples) + start = min(start, frame_count - 1) + end = min(max(start + 1, end), frame_count) + return frames[start:end] + + def _float_waveform(waveform: torch.Tensor) -> torch.Tensor: if waveform.dtype.is_floating_point: return waveform.float() diff --git a/tests/test_audio.py b/tests/test_audio.py index ec894f9..d0a51c0 100644 --- a/tests/test_audio.py +++ b/tests/test_audio.py @@ -1,6 +1,51 @@ +import pytest import torch -from sam_audio_comfy.audio import prepare_audio, prepare_visual_prompt, result_to_audio +from sam_audio_comfy.audio import ( + AudioChunk, + crossfade_audio_chunks, + plan_audio_chunks, + prepare_audio, + prepare_visual_prompt, + result_to_audio, + slice_visual_frames, +) + + +def test_chunk_plan_covers_audio_with_requested_overlap(): + assert plan_audio_chunks(21, 1, 10.0, 2.0) == [ + AudioChunk(0, 10), + AudioChunk(8, 18), + AudioChunk(16, 21), + ] + assert plan_audio_chunks(21, 1, 0.0, 2.0) == [AudioChunk(0, 21)] + assert plan_audio_chunks(5, 1, 10.0, 2.0) == [AudioChunk(0, 5)] + + with pytest.raises(ValueError, match="shorter than chunk_duration"): + plan_audio_chunks(21, 1, 10.0, 10.0) + + +def test_crossfade_audio_chunks_normalizes_overlap(): + first = torch.ones(1, 1, 6) + second = torch.full((1, 1, 6), 3.0) + output = crossfade_audio_chunks( + [(AudioChunk(0, 6), first), (AudioChunk(4, 10), second)], 10 + ) + + assert torch.equal(output[..., :4], torch.ones(1, 1, 4)) + assert torch.allclose(output[0, 0, 4:6], torch.tensor([5 / 3, 7 / 3])) + assert torch.equal(output[..., 6:], torch.full((1, 1, 4), 3.0)) + + +def test_visual_frames_are_sliced_to_chunk_time_range(): + frames = torch.arange(10).view(10, 1) + + assert torch.equal( + slice_visual_frames(frames, AudioChunk(2, 6), 10), frames[2:6] + ) + assert torch.equal( + slice_visual_frames(frames, AudioChunk(0, 10), 10), frames + ) def test_prepare_audio_preserves_batch_and_resamples(): diff --git a/tests/test_nodes.py b/tests/test_nodes.py index cb208cc..d711445 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -1,3 +1,4 @@ +import contextlib import importlib.util import sys import types @@ -12,6 +13,8 @@ def _load_node_package(monkeypatch, tmp_path): management = types.ModuleType("comfy.model_management") management.load_models_gpu = lambda models, force_full_load=False: None 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 = [] @@ -25,6 +28,7 @@ def _load_node_package(monkeypatch, tmp_path): 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] @@ -54,7 +58,18 @@ class FakeBatch: 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): @@ -101,6 +116,8 @@ def test_text_node_runs_full_audio_batch_with_standard_shapes(monkeypatch, tmp_p "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", @@ -109,3 +126,99 @@ def test_text_node_runs_full_audio_batch_with_standard_shapes(monkeypatch, tmp_p "SAMAudioVisualSeparator", } assert management.load_models_gpu is not None + + +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] diff --git a/tests/test_package.py b/tests/test_package.py index ffd0d08..5cb5b12 100644 --- a/tests/test_package.py +++ b/tests/test_package.py @@ -44,15 +44,15 @@ def test_example_sampler_widgets_include_seed_control(): expected = { "text_separation_workflow.json": ( "SAMAudioTextSeparator", - ["man speaking", False, 0, "fixed", 32], + ["man speaking", False, 0, "fixed", 32, 10.0, 1.0], ), "span_separation_workflow.json": ( "SAMAudioSpanSeparator", - ["car honking", 0, "fixed", 32], + ["car honking", 0, "fixed", 32, 10.0, 1.0], ), "visual_separation_workflow.json": ( "SAMAudioVisualSeparator", - ["", 0, "fixed", 32], + ["", 0, "fixed", 32, 10.0, 1.0], ), } for filename, (node_type, widget_values) in expected.items():