From 68239634ea44a9def6caa7ea712f9f17acce3baa Mon Sep 17 00:00:00 2001 From: Ethanfel <28683423+ethanfel@users.noreply.github.com> Date: Sun, 6 Sep 2026 22:59:57 +0200 Subject: [PATCH] Harden SAM-Audio loading and inference --- .github/workflows/publish_action.yml | 23 ++ .github/workflows/test.yml | 38 ++ README.md | 16 +- examples/span_separation_workflow.json | 4 +- examples/text_separation_workflow.json | 4 +- examples/video_separation_workflow.json | 4 +- examples/visual_separation_workflow.json | 4 +- install.py | 6 + nodes.py | 254 ++++++++++--- pyproject.toml | 2 +- sam_audio_comfy/attention.py | 2 +- sam_audio_comfy/audio.py | 52 ++- sam_audio_comfy/runtime.py | 442 +++++++++++++++++++++-- sam_audio_comfy/torchcodec_compat.py | 27 ++ tests/test_attention.py | 22 ++ tests/test_audio.py | 34 ++ tests/test_nodes.py | 125 ++++++- tests/test_package.py | 25 +- tests/test_runtime.py | 202 +++++++++++ 19 files changed, 1179 insertions(+), 107 deletions(-) create mode 100644 .github/workflows/test.yml diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml index a57c8f7..b3cf2a9 100644 --- a/.github/workflows/publish_action.yml +++ b/.github/workflows/publish_action.yml @@ -12,8 +12,31 @@ permissions: contents: read jobs: + test: + name: Test release + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install test dependencies + run: | + python -m pip install --index-url https://download.pytorch.org/whl/cpu torch torchaudio + python -m pip install pytest ruff + + - name: Lint and test + run: | + ruff check . + pytest -q + publish-node: name: Publish custom node + needs: test runs-on: ubuntu-latest steps: - name: Check out code diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..fd0d112 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,38 @@ +name: Tests + +on: + pull_request: + push: + branches: + - main + +permissions: + contents: read + +jobs: + test: + name: Python ${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.13"] + steps: + - name: Check out code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install test dependencies + run: | + python -m pip install --index-url https://download.pytorch.org/whl/cpu torch torchaudio + python -m pip install pytest ruff + + - name: Lint + run: ruff check . + + - name: Test + run: pytest -q diff --git a/README.md b/README.md index b0ba5d4..96348f9 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ Supports text prompts, positive and negative time spans, and visual masks. Each ## Requirements - Python 3.11 or newer +- Git - A CUDA GPU is strongly recommended - An internet connection for the first model download @@ -29,6 +30,11 @@ Then restart ComfyUI. If installed through ComfyUI Manager, the installer should The node uses ComfyUI's decoded media tensors and does not install or load TorchCodec. +The installer does not install or replace PyTorch, TorchCodec, or xFormers. It installs +the listed inference dependencies and three pinned Meta source packages into the Python +environment running ComfyUI. Existing packages are kept when they already satisfy the +requirements; an older `transformers` may be upgraded to the required version. + ## Model download Select a model in **SAM-Audio Model Loader** and queue the workflow. The loader first tries Meta's official repository. If Hugging Face denies gated access, it automatically downloads the matching pinned and checksum-verified checkpoint from the [public mirror collection](https://huggingface.co/collections/mrfakename/sam-audio). @@ -41,7 +47,7 @@ hf auth login Models download on first use to the active ComfyUI models directory under `models/sam_audio/`. For example, when ComfyUI is started with `--models-directory /path/to/models`, SAM-Audio checkpoints are stored in `/path/to/models/sam_audio/`. Both official and mirrored checkpoints remain subject to Meta's SAM License. -The required T5 asset is pinned in the Hugging Face cache and reused locally after its first download. The separate 5.8 GiB PE span predictor is loaded only when `predict_spans` is enabled. +The required T5 asset is pinned in the Hugging Face cache and reused locally after its first download. The separate 5.8 GiB PE span predictor is loaded only when `predict_spans` is enabled. `reranking_candidates` above 1 similarly downloads Meta's optional 5.7 GiB Judge checkpoint once and requires a text description. ## Nodes @@ -63,17 +69,23 @@ All nodes are in `audio/SAM-Audio`. Long audio is processed in overlapping chunks by default. Set `chunk_duration` to `0` for a single full-clip pass. +`precision=auto` uses BF16 on supported GPUs and FP32 otherwise. `initial_device=cpu` keeps the normal ComfyUI offload path; choose `gpu` to materialize checkpoint weights directly on the GPU and avoid the CPU-first checkpoint stage. + 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. +- Progress and cancellation are checked during every ODE inference step. - White pixels in a visual mask identify the target object. -- Native `VIDEO` inputs must contain an audio track. Use ComfyUI's **Create Video** node to attach separate audio when needed. +- Native `VIDEO` inputs use the video's real frame rate. Frames and embedded audio are trimmed to their common timestamped duration. Use ComfyUI's **Create Video** node to attach separate audio when needed. +- Span boundaries are quantized to the model's 40 ms audio-feature grid; zero-width and contradictory positive/negative anchors are rejected. - Local model folders must contain `config.json` and `checkpoint.pt`. +- All configured ComfyUI `sam_audio` model roots are searched, including roots from `extra_model_paths.yaml`. - 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. - ComfyUI's **Unload Models** action moves SAM-Audio out of VRAM but keeps the cached loader output in system RAM. Use **Unload Models and Execution Cache** when you also want ComfyUI to discard that cached CPU model. +- Because Meta's source packages are installed without their training-only dependency trees, `pip check` may report their optional packages as missing. This is expected for this inference-only integration. ## License diff --git a/examples/span_separation_workflow.json b/examples/span_separation_workflow.json index 65920f3..9ae3e42 100644 --- a/examples/span_separation_workflow.json +++ b/examples/span_separation_workflow.json @@ -15,7 +15,7 @@ {"name": "pipeline", "type": "SAM_AUDIO_PIPELINE", "links": [1], "slot_index": 0} ], "properties": {"Node name for S&R": "SAMAudioPipelineLoader"}, - "widgets_values": ["facebook/sam-audio-large", "pytorch"] + "widgets_values": ["facebook/sam-audio-large", "pytorch", "auto", "cpu"] }, { "id": 2, @@ -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, 10.0, 1.0] + "widgets_values": ["car honking", 0, "fixed", 32, 10.0, 1.0, 1] }, { "id": 5, diff --git a/examples/text_separation_workflow.json b/examples/text_separation_workflow.json index 0998557..1e1b921 100644 --- a/examples/text_separation_workflow.json +++ b/examples/text_separation_workflow.json @@ -15,7 +15,7 @@ {"name": "pipeline", "type": "SAM_AUDIO_PIPELINE", "links": [1], "slot_index": 0} ], "properties": {"Node name for S&R": "SAMAudioPipelineLoader"}, - "widgets_values": ["facebook/sam-audio-large", "pytorch"] + "widgets_values": ["facebook/sam-audio-large", "pytorch", "auto", "cpu"] }, { "id": 2, @@ -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, 10.0, 1.0] + "widgets_values": ["man speaking", false, 0, "fixed", 32, 10.0, 1.0, 1] }, { "id": 4, diff --git a/examples/video_separation_workflow.json b/examples/video_separation_workflow.json index c05364e..21af861 100644 --- a/examples/video_separation_workflow.json +++ b/examples/video_separation_workflow.json @@ -15,7 +15,7 @@ {"name": "pipeline", "type": "SAM_AUDIO_PIPELINE", "links": [1], "slot_index": 0} ], "properties": {"Node name for S&R": "SAMAudioPipelineLoader"}, - "widgets_values": ["facebook/sam-audio-base-tv", "pytorch"] + "widgets_values": ["facebook/sam-audio-base-tv", "pytorch", "auto", "cpu"] }, { "id": 2, @@ -66,7 +66,7 @@ {"name": "residual", "type": "AUDIO", "links": [5], "slot_index": 1} ], "properties": {"Node name for S&R": "SAMAudioVideoSeparator"}, - "widgets_values": ["", 0, "fixed", 32, 10.0, 1.0] + "widgets_values": ["", 0, "fixed", 32, 10.0, 1.0, 1] }, { "id": 5, diff --git a/examples/visual_separation_workflow.json b/examples/visual_separation_workflow.json index aa2baba..5c5d907 100644 --- a/examples/visual_separation_workflow.json +++ b/examples/visual_separation_workflow.json @@ -15,7 +15,7 @@ {"name": "pipeline", "type": "SAM_AUDIO_PIPELINE", "links": [1], "slot_index": 0} ], "properties": {"Node name for S&R": "SAMAudioPipelineLoader"}, - "widgets_values": ["facebook/sam-audio-large-tv", "pytorch"] + "widgets_values": ["facebook/sam-audio-large-tv", "pytorch", "auto", "cpu"] }, { "id": 2, @@ -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, 10.0, 1.0] + "widgets_values": ["", 0, "fixed", 32, 10.0, 1.0, 1] }, { "id": 5, diff --git a/install.py b/install.py index 61e8587..311f363 100644 --- a/install.py +++ b/install.py @@ -2,6 +2,7 @@ from __future__ import annotations +import shutil import subprocess import sys from pathlib import Path @@ -25,6 +26,11 @@ def run_pip(*arguments: str) -> None: def main() -> None: + if shutil.which("git") is None: + raise SystemExit( + "Git is required to install SAM-Audio's pinned source packages. " + "Install Git, then run install.py again with ComfyUI's Python." + ) run_pip("-r", str(ROOT / "requirements.txt")) for package in UPSTREAM_PACKAGES: run_pip("--no-deps", package) diff --git a/nodes.py b/nodes.py index 332c64c..a785c2c 100644 --- a/nodes.py +++ b/nodes.py @@ -1,5 +1,7 @@ from __future__ import annotations +import logging +import math from collections.abc import Sequence from typing import Any @@ -18,8 +20,11 @@ from .sam_audio_comfy.audio import ( prepare_visual_prompt, result_to_audio, slice_visual_frames, + synchronize_video_audio, ) from .sam_audio_comfy.runtime import ( + INITIAL_DEVICES, + PRECISIONS, SAMAudioPipeline, SpanPrompt, load_pipeline, @@ -34,6 +39,7 @@ CATEGORY = "audio/SAM-Audio" PIPELINE_TYPE = "SAM_AUDIO_PIPELINE" SPANS_TYPE = "SAM_AUDIO_SPANS" MAX_SEED = 2**64 - 1 +LOGGER = logging.getLogger(__name__) register_model_folder(folder_paths) @@ -88,6 +94,21 @@ def _chunking_inputs() -> dict[str, tuple[Any, ...]]: } +def _quality_inputs() -> dict[str, tuple[Any, ...]]: + return { + "reranking_candidates": ( + "INT", + { + "default": 1, + "min": 1, + "max": 8, + "step": 1, + "tooltip": "Generate multiple candidates and select one with Meta's Judge. Values above 1 lazily download and load the optional Judge model.", + }, + ) + } + + def _as_result_list(value: Any) -> list[torch.Tensor]: if isinstance(value, torch.Tensor): return [value] if value.ndim == 1 else list(value) @@ -97,7 +118,10 @@ def _as_result_list(value: Any) -> list[torch.Tensor]: def _spans_for_chunk( - spans: tuple[SpanPrompt, ...], chunk: AudioChunk, sample_rate: int + spans: tuple[SpanPrompt, ...], + chunk: AudioChunk, + sample_rate: int, + hop_length: int, ) -> tuple[SpanPrompt, ...]: chunk_start = chunk.start / sample_rate chunk_end = chunk.end / sample_rate @@ -105,9 +129,13 @@ def _spans_for_chunk( for span in spans: start = max(span.start, chunk_start) end = min(span.end, chunk_end) - if end > start: + local_start = start - chunk_start + local_end = end - chunk_start + if end > start and math.ceil(local_start * sample_rate / hop_length) < math.ceil( + local_end * sample_rate / hop_length + ): localized.append( - SpanPrompt(span.token, start - chunk_start, end - chunk_start) + SpanPrompt(span.token, local_start, local_end) ) return tuple(localized) @@ -124,6 +152,8 @@ def _run_separation( predict_spans: bool = False, spans: tuple[SpanPrompt, ...] | None = None, visual_prompt: tuple[torch.Tensor, torch.Tensor] | None = None, + visual_frame_rate: float | None = None, + reranking_candidates: int = 1, ) -> tuple[dict[str, Any], dict[str, Any]]: if not isinstance(pipeline, SAMAudioPipeline): raise TypeError("pipeline must come from a SAM-Audio Model Loader node") @@ -137,6 +167,34 @@ def _run_separation( ) sample_rate = pipeline.processor.audio_sampling_rate batch_size = len(waveforms) + + interrupt = getattr( + model_management, "throw_exception_if_processing_interrupted", lambda: None + ) + interrupt() + + masked_video = None + if visual_prompt is not None: + frames, mask = visual_prompt + masked_video = pipeline.processor.mask_videos([frames], [mask])[0] + if visual_frame_rate is not None: + original_audio_length = target_length + original_frame_count = masked_video.shape[0] + waveforms, target_length, masked_video = synchronize_video_audio( + waveforms, + target_length, + masked_video, + sample_rate, + visual_frame_rate, + ) + if ( + target_length != original_audio_length + or masked_video.shape[0] != original_frame_count + ): + LOGGER.info( + "Trimmed native VIDEO frames/audio to their common timestamped duration" + ) + duration = target_length / sample_rate chunks = plan_audio_chunks( target_length, @@ -146,65 +204,109 @@ def _run_separation( ) if spans is not None: - validate_spans(spans, duration) + validate_spans( + spans, + duration, + hop_length=pipeline.processor.audio_hop_length, + sample_rate=sample_rate, + ) if predict_spans: pipeline.ensure_span_predictor(model_management) - - masked_video = None - if visual_prompt is not None: - frames, mask = visual_prompt - masked_video = pipeline.processor.mask_videos([frames], [mask])[0] + if reranking_candidates < 1 or reranking_candidates > 8: + raise ValueError("reranking_candidates must be between 1 and 8") + if reranking_candidates > 1: + if not description: + raise ValueError( + "Judge reranking requires a text description, including for visual prompts" + ) + pipeline.ensure_judge_ranker(model_management) pipeline.load(model_management) 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 - ] + progress = comfy_utils.ProgressBar(len(chunks) * inference_steps) + progress_ticks = 0 - 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)] + def inference_tick(_module=None, _inputs=None): + nonlocal progress_ticks + interrupt() + if progress_ticks < len(chunks) * inference_steps: + progress.update(1) + progress_ticks += 1 - 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)] + transformer = getattr(pipeline.model, "transformer", None) + register_hook = getattr(transformer, "register_forward_pre_hook", None) + progress_hook = register_hook(inference_tick) if callable(register_hook) else None + try: + for chunk_index, chunk in enumerate(chunks): + interrupt() + chunk_waveforms = [ + waveform[..., chunk.start : chunk.end] for waveform in waveforms + ] - batch = pipeline.processor( - audios=chunk_waveforms, - descriptions=[description for _ in range(batch_size)], - anchors=anchors, - masked_videos=masked_videos, - ) - batch = batch.to(pipeline.device) + anchors = None + if spans is not None: + localized_spans = _spans_for_chunk( + spans, + chunk, + sample_rate, + pipeline.processor.audio_hop_length, + ) + anchors = [list(localized_spans) for _ in range(batch_size)] - 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, + masked_videos = None + if masked_video is not None: + chunk_video = slice_visual_frames( + masked_video, + chunk, + target_length, + sample_rate=sample_rate if visual_frame_rate is not None else None, + frame_rate=visual_frame_rate, + ) + 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), 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) + chunk_seed = (seed + chunk_index) % (2**64) + autocast_enabled = pipeline.dtype != torch.float32 + with ( + seeded_inference(chunk_seed, pipeline.device), + attention_backend_context(pipeline.attention_backend), + torch.autocast( + device_type=pipeline.device.type, + dtype=pipeline.dtype, + enabled=autocast_enabled, + ), + torch.inference_mode(), + ): + result = pipeline.model.separate( + batch, + ode_opt=sampling_options, + reranking_candidates=reranking_candidates, + 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"])) + if progress_hook is None: + for _ in range(inference_steps): + inference_tick() + finally: + if progress_hook is not None: + progress_hook.remove() target = { "waveform": crossfade_audio_chunks(target_chunks, target_length).contiguous(), @@ -245,7 +347,23 @@ class SAMAudioPipelineLoader: "tooltip": "PyTorch SDPA is the accurate default. Comfy Kitchen uses its quantized INT8 attention kernel when supported.", }, ), - } + }, + "optional": { + "precision": ( + list(PRECISIONS), + { + "default": "auto", + "tooltip": "Auto uses BF16 on supported GPUs and FP32 otherwise. BF16 roughly halves model memory and is usually faster.", + }, + ), + "initial_device": ( + list(INITIAL_DEVICES), + { + "default": "cpu", + "tooltip": "GPU assigns checkpoint tensors directly to the GPU during loading. CPU gives ComfyUI its normal offloadable starting state.", + }, + ), + }, } RETURN_TYPES = (PIPELINE_TYPE,) @@ -253,7 +371,13 @@ class SAMAudioPipelineLoader: FUNCTION = "load_model" CATEGORY = CATEGORY - def load_model(self, model_name: str, attention_backend: str = "pytorch"): + def load_model( + self, + model_name: str, + attention_backend: str = "pytorch", + precision: str = "auto", + initial_device: str = "cpu", + ): return ( load_pipeline( model_name, @@ -261,6 +385,8 @@ class SAMAudioPipelineLoader: model_management, model_patcher, attention_backend, + precision, + initial_device, ), ) @@ -346,7 +472,7 @@ class SAMAudioTextSeparator: } required.update(_sampling_inputs()) required.update(_chunking_inputs()) - return {"required": required} + return {"required": required, "optional": _quality_inputs()} RETURN_TYPES = ("AUDIO", "AUDIO") RETURN_NAMES = ("target", "residual") @@ -363,6 +489,7 @@ class SAMAudioTextSeparator: inference_steps: int, chunk_duration: float = 10.0, chunk_overlap: float = 1.0, + reranking_candidates: int = 1, ): return _run_separation( pipeline, @@ -373,6 +500,7 @@ class SAMAudioTextSeparator: chunk_duration, chunk_overlap, predict_spans=predict_spans, + reranking_candidates=reranking_candidates, ) @@ -395,7 +523,7 @@ class SAMAudioSpanSeparator: } required.update(_sampling_inputs()) required.update(_chunking_inputs()) - return {"required": required} + return {"required": required, "optional": _quality_inputs()} RETURN_TYPES = ("AUDIO", "AUDIO") RETURN_NAMES = ("target", "residual") @@ -412,6 +540,7 @@ class SAMAudioSpanSeparator: inference_steps: int, chunk_duration: float = 10.0, chunk_overlap: float = 1.0, + reranking_candidates: int = 1, ): return _run_separation( pipeline, @@ -422,6 +551,7 @@ class SAMAudioSpanSeparator: chunk_duration, chunk_overlap, spans=tuple(spans), + reranking_candidates=reranking_candidates, ) @@ -455,7 +585,7 @@ class SAMAudioVisualSeparator: } required.update(_sampling_inputs()) required.update(_chunking_inputs()) - return {"required": required} + return {"required": required, "optional": _quality_inputs()} RETURN_TYPES = ("AUDIO", "AUDIO") RETURN_NAMES = ("target", "residual") @@ -473,6 +603,7 @@ class SAMAudioVisualSeparator: inference_steps: int, chunk_duration: float = 10.0, chunk_overlap: float = 1.0, + reranking_candidates: int = 1, ): visual_prompt = prepare_visual_prompt(images, mask) return _run_separation( @@ -484,6 +615,7 @@ class SAMAudioVisualSeparator: chunk_duration, chunk_overlap, visual_prompt=visual_prompt, + reranking_candidates=reranking_candidates, ) @@ -520,7 +652,7 @@ class SAMAudioVideoSeparator: } required.update(_sampling_inputs()) required.update(_chunking_inputs()) - return {"required": required} + return {"required": required, "optional": _quality_inputs()} RETURN_TYPES = ("AUDIO", "AUDIO") RETURN_NAMES = ("target", "residual") @@ -537,6 +669,7 @@ class SAMAudioVideoSeparator: inference_steps: int, chunk_duration: float = 10.0, chunk_overlap: float = 1.0, + reranking_candidates: int = 1, ): get_components = getattr(video, "get_components", None) if not callable(get_components): @@ -545,7 +678,8 @@ class SAMAudioVideoSeparator: components = get_components() images = getattr(components, "images", None) audio = getattr(components, "audio", None) - if not isinstance(images, torch.Tensor) or images.shape[0] < 1: + frame_rate = getattr(components, "frame_rate", None) + if not isinstance(images, torch.Tensor) or images.ndim < 1 or images.shape[0] < 1: raise ValueError("VIDEO input does not contain any decoded frames") if audio is None: raise ValueError( @@ -553,6 +687,12 @@ class SAMAudioVideoSeparator: "Create Video node or use SAM-Audio Visual Separate with separate " "IMAGE and AUDIO inputs" ) + try: + frame_rate = float(frame_rate) + except (TypeError, ValueError) as error: + raise ValueError("VIDEO input does not contain a valid frame rate") from error + if not math.isfinite(frame_rate) or frame_rate <= 0: + raise ValueError("VIDEO frame rate must be a positive finite number") visual_prompt = prepare_visual_prompt(images, mask) return _run_separation( @@ -564,6 +704,8 @@ class SAMAudioVideoSeparator: chunk_duration, chunk_overlap, visual_prompt=visual_prompt, + visual_frame_rate=frame_rate, + reranking_candidates=reranking_candidates, ) diff --git a/pyproject.toml b/pyproject.toml index f1373b8..7d1f729 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "comfyui-meta-sam-audio" -version = "0.3.1" +version = "0.4.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/attention.py b/sam_audio_comfy/attention.py index 4d63d8e..091fdd1 100644 --- a/sam_audio_comfy/attention.py +++ b/sam_audio_comfy/attention.py @@ -67,7 +67,7 @@ def _kitchen_sdpa(comfy_kitchen: Any): ) -> torch.Tensor: # Comfy Kitchen is an inference-only, non-causal kernel. Preserve exact # PyTorch behavior for uncommon calls outside its supported surface. - if dropout_p != 0.0 or is_causal: + if dropout_p != 0.0 or is_causal or enable_gqa: return _PYTORCH_SDPA( query, key, diff --git a/sam_audio_comfy/audio.py b/sam_audio_comfy/audio.py index 8e7fd25..467186d 100644 --- a/sam_audio_comfy/audio.py +++ b/sam_audio_comfy/audio.py @@ -99,9 +99,14 @@ def crossfade_audio_chunks( def slice_visual_frames( - frames: torch.Tensor, chunk: AudioChunk, total_samples: int + frames: torch.Tensor, + chunk: AudioChunk, + total_samples: int, + *, + sample_rate: int | None = None, + frame_rate: float | None = None, ) -> torch.Tensor: - """Select frames whose uniformly spaced time range intersects an audio chunk.""" + """Select the frames whose 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: @@ -110,13 +115,52 @@ def slice_visual_frames( 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) + if frame_rate is not None: + if sample_rate is None or sample_rate < 1: + raise ValueError("sample_rate is required with frame_rate") + if not math.isfinite(frame_rate) or frame_rate <= 0: + raise ValueError("frame_rate must be a positive finite number") + start = math.floor((chunk.start / sample_rate) * frame_rate) + end = math.ceil((chunk.end / sample_rate) * frame_rate) + else: + 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 synchronize_video_audio( + waveforms: Sequence[torch.Tensor], + target_length: int, + frames: torch.Tensor, + sample_rate: int, + frame_rate: float, +) -> tuple[list[torch.Tensor], int, torch.Tensor]: + """Trim native VIDEO components to their common timestamped duration.""" + if sample_rate < 1: + raise ValueError("sample_rate must be positive") + if not math.isfinite(frame_rate) or frame_rate <= 0: + raise ValueError("VIDEO frame rate must be a positive finite number") + if frames.ndim < 1 or frames.shape[0] < 1: + raise ValueError("VIDEO must contain at least one frame") + + # ComfyUI defines a tensor-backed video's duration as frame_count / frame_rate. + # Match its encoder by discarding audio beyond that duration, while also dropping + # video frames which occur after a shorter embedded audio track ends. + video_samples = max(1, math.floor((frames.shape[0] / frame_rate) * sample_rate)) + common_length = min(target_length, video_samples) + if common_length < 1: + raise ValueError("VIDEO and its embedded audio have no overlapping duration") + + frame_count = min( + frames.shape[0], + max(1, math.ceil((common_length / sample_rate) * frame_rate)), + ) + trimmed_waveforms = [waveform[..., :common_length].contiguous() for waveform in waveforms] + return trimmed_waveforms, common_length, frames[:frame_count].contiguous() + + def _float_waveform(waveform: torch.Tensor) -> torch.Tensor: if waveform.dtype.is_floating_point: return waveform.float() diff --git a/sam_audio_comfy/runtime.py b/sam_audio_comfy/runtime.py index db68dcf..b5311db 100644 --- a/sam_audio_comfy/runtime.py +++ b/sam_audio_comfy/runtime.py @@ -1,11 +1,16 @@ from __future__ import annotations import contextlib +import functools import hashlib import importlib import json import logging +import math import os +import re +import sys +import threading import time from dataclasses import dataclass from pathlib import Path @@ -13,9 +18,9 @@ from typing import Any, Iterator, Literal, NamedTuple import torch +from . import torchcodec_compat from .attention import validate_attention_backend from .audiotools_compat import audiotools_import_compatibility -from .torchcodec_compat import torchcodec_import_compatibility from .xformers_compat import xformers_import_compatibility LOGGER = logging.getLogger(__name__) @@ -32,6 +37,10 @@ OFFICIAL_MODELS = ( MODEL_FOLDER = "sam_audio" MODEL_FILES = ("config.json", "checkpoint.pt") MIRROR_MARKER = ".verified-mirror-checkpoint" +PRECISIONS = ("auto", "bf16", "fp32") +INITIAL_DEVICES = ("cpu", "gpu") +CPU_DEVICE = torch.device("cpu") +_IMPORT_LOCK = threading.RLock() @dataclass(frozen=True) @@ -102,6 +111,18 @@ SUPPORT_MODELS = { "facebook/pe-a-frame-large", "40187271298f84e2966d4518c88dde698540c9ad", ), + "sam-audio-judge": SupportModelSpec( + "facebook/sam-audio-judge", + "cd52eb977697cc0b25d4ecd309a88d4af46223ee", + ( + "checkpoint.pt", + "config.json", + "preprocessor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json", + ), + ), } @@ -120,6 +141,7 @@ class SAMAudioPipeline: source: str attention_backend: str = "pytorch" span_predictor_source: str | None = None + dtype: torch.dtype = torch.float32 def get_models(self) -> list[Any]: """Let ComfyUI track this model for the lifetime of the prompt.""" @@ -142,28 +164,48 @@ class SAMAudioPipeline: "This SAM-Audio checkpoint does not declare a span predictor" ) - # If the core model was already used, move it off the GPU before adding a - # new submodule. This also makes ComfyUI recalculate the patcher's size. - unload = getattr(model_management, "unload_model_and_clones", None) - if unload is not None: - try: - unload(self.patcher, all_devices=True) - except TypeError: - unload(self.patcher) + self._unload_before_extension(model_management) started = time.perf_counter() predictor, transform = _load_span_predictor_assets( self.span_predictor_source ) - self.model.span_predictor = predictor.eval().to("cpu") + self.model.span_predictor = predictor.eval().to(device="cpu", dtype=self.dtype) self.model.span_predictor_transform = transform - if hasattr(self.patcher, "size"): - self.patcher.size = 0 + self._reset_patcher_size() LOGGER.info( "Loaded optional SAM-Audio span predictor in %.1f seconds", time.perf_counter() - started, ) + def ensure_judge_ranker(self, model_management: Any) -> None: + """Attach Meta's Judge ranker only when multiple candidates are requested.""" + if getattr(self.model, "text_ranker", None) is not None: + return + self._unload_before_extension(model_management) + started = time.perf_counter() + self.model.text_ranker = _load_judge_ranker(self.dtype) + self._reset_patcher_size() + LOGGER.info( + "Loaded optional SAM-Audio Judge reranker in %.1f seconds", + time.perf_counter() - started, + ) + + def _unload_before_extension(self, model_management: Any) -> None: + # Move an already-used model off the GPU before adding a submodule so + # ComfyUI can recalculate its complete size and residency. + unload = getattr(model_management, "unload_model_and_clones", None) + if unload is None: + return + try: + unload(self.patcher, all_devices=True) + except TypeError: + unload(self.patcher) + + def _reset_patcher_size(self) -> None: + if hasattr(self.patcher, "size"): + self.patcher.size = 0 + def register_model_folder(folder_paths: Any) -> Path: model_dir = Path(folder_paths.models_dir) / MODEL_FOLDER @@ -357,6 +399,19 @@ def resolve_model(model_name: str, folder_paths: Any) -> Path: f"Unknown SAM-Audio model {model_name!r}. Refresh the browser after adding local models." ) model_root = register_model_folder(folder_paths) + official_directory = model_name.rsplit("/", 1)[-1] + for raw_root in folder_paths.get_folder_paths(MODEL_FOLDER): + candidate = Path(raw_root) / official_directory + if not _is_model_directory(candidate): + continue + marker = candidate / MIRROR_MARKER + if marker.exists(): + if not (candidate / "LICENSE").is_file(): + # Never trust a partial mirror download. The primary destination's + # downloader can resume or replace it after this search finishes. + continue + _verify_mirror_checkpoint(candidate, MIRROR_MODELS[model_name]) + return candidate return _download_official_model(model_name, model_root) @@ -365,12 +420,16 @@ def import_sam_audio() -> tuple[Any, Any]: LOGGER.info( "Loading SAM-Audio through isolated inference-only compatibility layers" ) - with ( - xformers_import_compatibility(), - audiotools_import_compatibility(), - torchcodec_import_compatibility(), - ): - module = importlib.import_module("sam_audio") + # Serializing our temporary sys.modules entries prevents two simultaneous + # loader executions from observing a half-installed compatibility surface. + with _IMPORT_LOCK: + with ( + xformers_import_compatibility(), + audiotools_import_compatibility(), + torchcodec_compat.torchcodec_import_compatibility(), + ): + module = importlib.import_module("sam_audio") + _install_lazy_torchcodec_decoders() return module.SAMAudio, module.SAMAudioProcessor except Exception as error: raise RuntimeError( @@ -379,6 +438,26 @@ def import_sam_audio() -> tuple[Any, Any]: ) from error +def _install_lazy_torchcodec_decoders() -> None: + """Keep upstream file I/O usable without importing TorchCodec during node load.""" + for name, module in tuple(sys.modules.items()): + if not ( + name == "sam_audio" + or name.startswith("sam_audio.") + or name == "core" + or name.startswith("core.") + ): + continue + if not hasattr(module, "__dict__"): + continue + for attribute, value in tuple(vars(module).items()): + if value is torchcodec_compat.AudioDecoder: + setattr(module, attribute, torchcodec_compat.LazyAudioDecoder) + elif value is torchcodec_compat.VideoDecoder: + setattr(module, attribute, torchcodec_compat.LazyVideoDecoder) + + +@functools.lru_cache(maxsize=None) def _cached_support_snapshot(spec: SupportModelSpec) -> str: try: huggingface_hub = importlib.import_module("huggingface_hub") @@ -460,9 +539,178 @@ def _load_span_predictor_assets(source: str) -> tuple[torch.nn.Module, Any]: return predictor, transform -def _load_sam_audio_checkpoint(SAMAudio: Any, model_path: Path) -> Any: - """Load a local checkpoint across old and new huggingface_hub call signatures.""" +def _load_local_checkpoint_efficient( + model_class: Any, + model_path: Path, + *, + config_overrides: dict[str, Any], + target_device: torch.device, + dtype: torch.dtype, + ignored_missing_pattern: str | None = None, + defer_text_encoder: bool = False, + use_meta_device: bool = True, +) -> torch.nn.Module: + """Materialize checkpoint tensors once, directly in their inference format.""" + config_path = model_path / "config.json" + checkpoint_path = model_path / "checkpoint.pt" + config = json.loads(config_path.read_text(encoding="utf-8")) + for key, value in config_overrides.items(): + if key in config: + config[key] = value + state_dict = torch.load( + checkpoint_path, + weights_only=True, + mmap=True, + map_location="cpu", + ) + + model_config = model_class.config_cls(**config) + if defer_text_encoder: + # SAMAudio's constructor downloads and materializes T5 while also allocating + # every much-larger checkpoint-backed module. Replace only T5 during the + # meta-device construction, then restore and load the real pretrained encoder. + constructor_globals = model_class.__init__.__globals__ + text_encoder_class = constructor_globals.get("T5TextEncoder") + if text_encoder_class is None: + raise RuntimeError("SAM-Audio's T5TextEncoder constructor was not found") + with _IMPORT_LOCK: + constructor_globals["T5TextEncoder"] = lambda _config: torch.nn.Identity() + try: + with torch.device("meta"): + model = model_class(model_config) + finally: + constructor_globals["T5TextEncoder"] = text_encoder_class + model.text_encoder = text_encoder_class(model_config.text_encoder) + elif not use_meta_device: + with _IMPORT_LOCK: + original_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + with torch.device(target_device): + model = model_class(model_config) + finally: + torch.set_default_dtype(original_dtype) + # Preserve the small derived buffers which checkpoint.pt intentionally + # omits, then release all random constructor parameters before materializing + # the real state dict. This avoids holding two Judge-sized models in RAM. + derived_buffers = { + name: value.detach().clone() + for name, value in model.named_buffers() + if name not in state_dict + } + model.to_empty(device=torch.device("meta")) + for name, value in derived_buffers.items(): + module_name, _, buffer_name = name.rpartition(".") + module = model.get_submodule(module_name) if module_name else model + module._buffers[buffer_name] = value + else: + with torch.device("meta"): + model = model_class(model_config) + for key in tuple(state_dict): + value = state_dict[key] + if not isinstance(value, torch.Tensor): + continue + destination_dtype = dtype if value.dtype.is_floating_point else value.dtype + if value.device != target_device or value.dtype != destination_dtype: + state_dict[key] = value.to( + device=target_device, + dtype=destination_dtype, + non_blocking=False, + ) + + result = torch.nn.Module.load_state_dict( + model, + state_dict, + strict=False, + assign=True, + ) + missing = list(result.missing_keys) + if ignored_missing_pattern is not None: + missing = [key for key in missing if not re.search(ignored_missing_pattern, key)] + if missing or result.unexpected_keys: + raise RuntimeError( + f"Checkpoint keys do not match {model_class.__name__}: " + f"missing={missing}, unexpected={list(result.unexpected_keys)}" + ) + + if defer_text_encoder: + _materialize_sam_audio_buffers(model, target_device) + + meta_tensors = [ + name + for name, tensor in (*model.named_parameters(), *model.named_buffers()) + if tensor.is_meta + ] + if meta_tensors: + raise RuntimeError( + "Checkpoint left meta tensors unmaterialized: " + ", ".join(meta_tensors) + ) + + # T5 and other assets loaded by the constructor are not part of checkpoint.pt. + # Move/cast those remaining parameters after assigning the mmap-backed weights. + return model.to(device=target_device, dtype=dtype) + + +def _materialize_sam_audio_buffers( + model: torch.nn.Module, target_device: torch.device +) -> None: + """Rebuild SAM-Audio's four derived, non-checkpoint buffers after meta init.""" + with torch.device(target_device): + vision_model = model.vision_encoder.model + if vision_model.attn_mask is not None and vision_model.attn_mask.is_meta: + vision_model.attn_mask = vision_model.build_causal_mask() + + rope = model.transformer.rope_embeddings + if rope is not None and rope.freqs_cis.is_meta: + rope.reset_parameters() + + time_embedder = model.transformer.t_embedder + if time_embedder.freqs.is_meta: + half = time_embedder.frequency_embedding_size // 2 + time_embedder.freqs = torch.exp( + -math.log(10_000) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ) + + positional_embedding = model.timestep_emb + if positional_embedding.inv_freq.is_meta: + half = positional_embedding.inv_freq.numel() + positional_embedding.inv_freq = torch.exp( + -math.log(10_000) * torch.arange(half).float() / half + ) + + +def _load_sam_audio_checkpoint( + SAMAudio: Any, + model_path: Path, + target_device: torch.device = CPU_DEVICE, + dtype: torch.dtype = torch.float32, +) -> Any: + """Load local weights efficiently, with a compatibility path for test doubles.""" model_kwargs = _sam_audio_model_kwargs(model_path) + if ( + (model_path / "config.json").is_file() + and (model_path / "checkpoint.pt").is_file() + and hasattr(SAMAudio, "config_cls") + ): + config_overrides = { + key: value + for key, value in model_kwargs.items() + if key not in {"map_location", "strict"} + } + return _load_local_checkpoint_efficient( + SAMAudio, + model_path, + config_overrides=config_overrides, + target_device=target_device, + dtype=dtype, + ignored_missing_pattern=( + r"(^text_encoder|^visual_ranker|^text_ranker|^span_predictor)" + ), + defer_text_encoder=True, + ) + try: return SAMAudio.from_pretrained(str(model_path), **model_kwargs) except TypeError as error: @@ -491,43 +739,149 @@ def _load_sam_audio_checkpoint(SAMAudio: Any, model_path: Path) -> Any: ) +class _LocalJudgeRanker(torch.nn.Module): + """Tensor-native wrapper around Meta's official SAM-Audio Judge.""" + + def __init__(self, model: torch.nn.Module, processor: Any): + super().__init__() + self.model = model + self.processor = processor + + @torch.inference_mode() + def forward( + self, + input_audio: list[torch.Tensor], + extracted_audio: list[torch.Tensor], + descriptions: list[str], + sample_rate: int = 48_000, + **_kwargs: Any, + ) -> torch.Tensor: + batch_size, candidates = len(input_audio), len(input_audio[0]) + inputs = [sample[None] for group in input_audio for sample in group] + outputs = [sample[None] for group in extracted_audio for sample in group] + repeated_text = [text for text in descriptions for _ in range(candidates)] + processed = self.processor( + text=repeated_text, + input_audio=inputs, + separated_audio=outputs, + return_tensors="pt", + padding=True, + sampling_rate=sample_rate, + ) + device = next(self.model.parameters()).device + result = self.model(**processed.to(device)) + return result.overall.view(batch_size, candidates) + + +def _load_judge_ranker(dtype: torch.dtype) -> torch.nn.Module: + source = Path(_cached_support_snapshot(SUPPORT_MODELS["sam-audio-judge"])) + try: + judge_module = importlib.import_module("sam_audio.model.judge") + processor_module = importlib.import_module("sam_audio.processor") + judge = _load_local_checkpoint_efficient( + judge_module.SAMAudioJudgeModel, + source, + config_overrides={}, + target_device=torch.device("cpu"), + dtype=dtype, + use_meta_device=False, + ) + processor = processor_module.SAMAudioJudgeProcessor.from_pretrained(str(source)) + except Exception as error: + raise RuntimeError( + f"Failed to load Meta's optional SAM-Audio Judge from {source}: " + f"{type(error).__name__}: {error}" + ) from error + return _LocalJudgeRanker(judge.eval(), processor).eval() + + +def resolve_precision( + precision: str, load_device: torch.device, model_management: Any +) -> torch.dtype: + if precision not in PRECISIONS: + raise ValueError(f"Unknown precision {precision!r}; choose one of {PRECISIONS}") + if precision == "fp32": + return torch.float32 + if precision == "bf16": + if load_device.type == "cpu": + raise RuntimeError("BF16 SAM-Audio inference requires a supported GPU") + supports_dtype = getattr(model_management, "supports_dtype", None) + if callable(supports_dtype) and not supports_dtype(load_device, torch.bfloat16): + raise RuntimeError(f"BF16 is not supported on {load_device}") + return torch.bfloat16 + + should_use_bf16 = getattr(model_management, "should_use_bf16", None) + if load_device.type != "cpu" and callable(should_use_bf16): + try: + if should_use_bf16(load_device): + return torch.bfloat16 + except Exception: + pass + return torch.float32 + + def load_pipeline( model_name: str, folder_paths: Any, model_management: Any, model_patcher: Any, attention_backend: str = "pytorch", + precision: str = "auto", + initial_device: str = "cpu", ) -> SAMAudioPipeline: load_device = model_management.get_torch_device() validate_attention_backend(attention_backend, load_device) + if initial_device not in INITIAL_DEVICES: + raise ValueError( + f"Unknown initial device {initial_device!r}; choose one of {INITIAL_DEVICES}" + ) + if initial_device == "gpu" and load_device.type == "cpu": + raise RuntimeError("Initial device 'gpu' was selected, but ComfyUI has no GPU device") + dtype = resolve_precision(precision, load_device, model_management) + checkpoint_device = ( + load_device + if initial_device == "gpu" and load_device.type != "cpu" + else torch.device("cpu") + ) model_path = resolve_model(model_name, folder_paths) SAMAudio, SAMAudioProcessor = import_sam_audio() - LOGGER.info("Loading SAM-Audio model from %s", model_path) + LOGGER.info( + "Loading SAM-Audio model from %s as %s on %s", + model_path, + str(dtype).removeprefix("torch."), + checkpoint_device, + ) started = time.perf_counter() try: - model = _load_sam_audio_checkpoint(SAMAudio, model_path) + model = _load_sam_audio_checkpoint( + SAMAudio, + model_path, + target_device=checkpoint_device, + dtype=dtype, + ) LOGGER.info( - "Loaded SAM-Audio checkpoint and required text encoder on CPU in %.1f seconds", + "Loaded SAM-Audio checkpoint and required text encoder on %s in %.1f seconds", + checkpoint_device, time.perf_counter() - started, ) processor = SAMAudioProcessor.from_pretrained(str(model_path)) except Exception as error: raise RuntimeError( - f"Failed to load SAM-Audio from {model_path}. Upstream may also need its " - "T5, Perception Encoder, and ranker assets available through Hugging Face." + f"Failed to load SAM-Audio from {model_path}: " + f"{type(error).__name__}: {error}" ) from error model = model.eval() # BaseModel exposes device() as a method, while ComfyUI's ModelPatcher tracks # residency through a mutable model.device attribute. Upstream does not call # that helper internally, so adopting ComfyUI's convention is safe here. - model.device = torch.device("cpu") + model.device = checkpoint_device offload_device = model_management.unet_offload_device() patcher = model_patcher.ModelPatcher( model, load_device=load_device, offload_device=offload_device ) - return SAMAudioPipeline( + pipeline = SAMAudioPipeline( model=model, processor=processor, patcher=patcher, @@ -535,7 +889,13 @@ def load_pipeline( source=model_name, attention_backend=attention_backend, span_predictor_source=_configured_span_predictor(model_path), + dtype=dtype, ) + if checkpoint_device == load_device and load_device.type != "cpu": + # Register direct-to-GPU loads immediately so ComfyUI's unload controls own + # the allocation even if a downstream node fails before inference starts. + pipeline.load(model_management) + return pipeline def ode_options(inference_steps: int) -> dict[str, Any]: @@ -556,7 +916,13 @@ def seeded_inference(seed: int, device: torch.device) -> Iterator[None]: yield -def validate_spans(spans: tuple[SpanPrompt, ...], duration: float) -> None: +def validate_spans( + spans: tuple[SpanPrompt, ...], + duration: float, + *, + hop_length: int | None = None, + sample_rate: int | None = None, +) -> None: if not spans: raise ValueError("At least one span prompt is required") if not any(span.token == "+" for span in spans): @@ -570,3 +936,23 @@ def validate_spans(spans: tuple[SpanPrompt, ...], duration: float) -> None: raise ValueError( f"Span ends at {span.end:g}s but the input audio is only {duration:.3f}s" ) + if hop_length is not None and sample_rate is not None: + start_step = math.ceil(span.start * sample_rate / hop_length) + end_step = math.ceil(span.end * sample_rate / hop_length) + if start_step == end_step: + resolution = hop_length / sample_rate + raise ValueError( + f"Span {span.start:g}-{span.end:g}s is shorter than SAM-Audio's " + f"effective {resolution:.3f}s anchor resolution at that position" + ) + + for index, first in enumerate(spans): + for second in spans[index + 1 :]: + if first.token == second.token: + continue + if max(first.start, second.start) < min(first.end, second.end): + raise ValueError( + "Positive and negative span prompts cannot overlap: " + f"{first.start:g}-{first.end:g}s and " + f"{second.start:g}-{second.end:g}s" + ) diff --git a/sam_audio_comfy/torchcodec_compat.py b/sam_audio_comfy/torchcodec_compat.py index 71fe117..375fcfd 100644 --- a/sam_audio_comfy/torchcodec_compat.py +++ b/sam_audio_comfy/torchcodec_compat.py @@ -3,6 +3,7 @@ from __future__ import annotations import contextlib +import importlib import sys import types from collections.abc import Iterator @@ -28,6 +29,32 @@ class AudioEncoder(_TensorInputsOnly): """Placeholder for an optional upstream ranker disabled by this integration.""" +class LazyAudioDecoder: + """Resolve the real TorchCodec decoder only if another caller needs file I/O.""" + + def __new__(cls, *args, **kwargs): + try: + decoder = importlib.import_module("torchcodec.decoders").AudioDecoder + except Exception as error: + raise RuntimeError( + "TorchCodec audio decoding is unavailable; pass a decoded ComfyUI AUDIO value" + ) from error + return decoder(*args, **kwargs) + + +class LazyVideoDecoder: + """Resolve the real TorchCodec decoder only if another caller needs file I/O.""" + + def __new__(cls, *args, **kwargs): + try: + decoder = importlib.import_module("torchcodec.decoders").VideoDecoder + except Exception as error: + raise RuntimeError( + "TorchCodec video decoding is unavailable; pass a decoded ComfyUI VIDEO value" + ) from error + return decoder(*args, **kwargs) + + def _compatibility_modules() -> dict[str, types.ModuleType]: torchcodec = types.ModuleType("torchcodec") torchcodec.__path__ = [] diff --git a/tests/test_attention.py b/tests/test_attention.py index cd700bd..e994206 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -76,6 +76,28 @@ def test_kitchen_backend_preserves_pytorch_for_causal_attention(monkeypatch): assert kitchen.calls == [] +def test_kitchen_backend_preserves_pytorch_for_grouped_query_attention(monkeypatch): + kitchen = FakeComfyKitchen() + monkeypatch.setattr(attention.importlib, "import_module", lambda _: kitchen) + target = types.ModuleType("sam_audio.fake_transformer") + target.F = torch_functional + monkeypatch.setitem(sys.modules, target.__name__, target) + + query = torch.randn(1, 4, 3, 4) + key = torch.randn(1, 2, 3, 4) + value = torch.randn(1, 2, 3, 4) + with attention.attention_backend_context("comfy_kitchen"): + actual = target.F.scaled_dot_product_attention( + query, key, value, enable_gqa=True + ) + expected = torch_functional.scaled_dot_product_attention( + query, key, value, enable_gqa=True + ) + + assert torch.allclose(actual, expected) + assert kitchen.calls == [] + + def test_kitchen_backend_contiguates_the_head_dimension(monkeypatch): kitchen = FakeComfyKitchen() monkeypatch.setattr(attention.importlib, "import_module", lambda _: kitchen) diff --git a/tests/test_audio.py b/tests/test_audio.py index d0a51c0..bf1e3ad 100644 --- a/tests/test_audio.py +++ b/tests/test_audio.py @@ -9,6 +9,7 @@ from sam_audio_comfy.audio import ( prepare_visual_prompt, result_to_audio, slice_visual_frames, + synchronize_video_audio, ) @@ -48,6 +49,39 @@ def test_visual_frames_are_sliced_to_chunk_time_range(): ) +def test_visual_frames_use_native_frame_timestamps_when_available(): + frames = torch.arange(30).view(30, 1) + + selected = slice_visual_frames( + frames, + AudioChunk(10, 20), + total_samples=40, + sample_rate=10, + frame_rate=5, + ) + + assert torch.equal(selected, frames[5:10]) + + +def test_native_video_components_are_trimmed_to_common_duration(): + waveforms = [torch.ones(1, 20)] + frames = torch.ones(8, 3, 2, 2) + + audio_trimmed, length, frames_kept = synchronize_video_audio( + waveforms, 20, frames, sample_rate=10, frame_rate=8 + ) + assert length == 10 + assert audio_trimmed[0].shape[-1] == 10 + assert frames_kept.shape[0] == 8 + + audio_kept, length, frames_trimmed = synchronize_video_audio( + waveforms, 20, torch.ones(30, 3, 2, 2), sample_rate=10, frame_rate=10 + ) + assert length == 20 + assert audio_kept[0].shape[-1] == 20 + assert frames_trimmed.shape[0] == 20 + + def test_prepare_audio_preserves_batch_and_resamples(): waveform = torch.stack( [torch.ones(2, 8), torch.full((2, 8), 0.5)], dim=0 diff --git a/tests/test_nodes.py b/tests/test_nodes.py index d87cc0f..aaae981 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -13,12 +13,24 @@ def _load_node_package(monkeypatch, tmp_path): comfy.__path__ = [] management = types.ModuleType("comfy.model_management") management.load_calls = [] + management.interrupt_checks = 0 + management.progress_totals = [] + management.progress_updates = [] management.load_models_gpu = lambda models, force_full_load=False: ( management.load_calls.append((models, force_full_load)) ) + management.throw_exception_if_processing_interrupted = lambda: setattr( + management, "interrupt_checks", management.interrupt_checks + 1 + ) patcher = types.ModuleType("comfy.model_patcher") utils = types.ModuleType("comfy.utils") - utils.ProgressBar = lambda total: types.SimpleNamespace(update=lambda value: None) + def progress_bar(total): + management.progress_totals.append(total) + return types.SimpleNamespace( + update=lambda value: management.progress_updates.append(value) + ) + + utils.ProgressBar = progress_bar folder_paths = types.ModuleType("folder_paths") folder_paths.models_dir = str(tmp_path) folder_paths._paths = [] @@ -36,7 +48,7 @@ def _load_node_package(monkeypatch, tmp_path): monkeypatch.setitem(sys.modules, "folder_paths", folder_paths) root = Path(__file__).parents[1] - package_name = "comfyui_sam_audio_node_test" + package_name = f"comfyui_sam_audio_node_test_{tmp_path.name}" spec = importlib.util.spec_from_file_location( package_name, root / "__init__.py", @@ -61,6 +73,7 @@ class FakeBatch: class FakeProcessor: audio_sampling_rate = 8 + audio_hop_length = 2 def __init__(self): self.calls = [] @@ -83,7 +96,12 @@ class FakeProcessor: class FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = [] + def separate(self, batch, **kwargs): + self.calls.append(kwargs) target = [row.mean(0) for row in batch.audios] return types.SimpleNamespace(target=target, residual=[-row for row in target]) @@ -119,9 +137,18 @@ def test_text_node_runs_full_audio_batch_with_standard_shapes(monkeypatch, tmp_p sampler_inputs = package.NODE_CLASS_MAPPINGS[ "SAMAudioTextSeparator" ].INPUT_TYPES()["required"] + sampler_optional = package.NODE_CLASS_MAPPINGS[ + "SAMAudioTextSeparator" + ].INPUT_TYPES()["optional"] + loader_optional = package.NODE_CLASS_MAPPINGS[ + "SAMAudioPipelineLoader" + ].INPUT_TYPES()["optional"] 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 sampler_optional["reranking_candidates"][1]["default"] == 1 + assert loader_optional["precision"][0] == ["auto", "bf16", "fp32"] + assert loader_optional["initial_device"][0] == ["cpu", "gpu"] assert set(package.NODE_CLASS_MAPPINGS) == { "SAMAudioPipelineLoader", "SAMAudioSpanPrompt", @@ -131,6 +158,9 @@ def test_text_node_runs_full_audio_batch_with_standard_shapes(monkeypatch, tmp_p "SAMAudioVideoSeparator", } assert management.load_calls == [([pipeline.patcher], True)] + assert management.progress_totals == [32] + assert sum(management.progress_updates) == 32 + assert management.interrupt_checks == 34 def _fake_pipeline(runtime, processor=None): @@ -241,7 +271,7 @@ def test_video_node_uses_native_comfy_video_frames_and_audio(monkeypatch, tmp_pa get_components=lambda: types.SimpleNamespace( images=images, audio=audio, - frame_rate=8, + frame_rate=4, ) ) @@ -266,6 +296,95 @@ def test_video_node_uses_native_comfy_video_frames_and_audio(monkeypatch, tmp_pa assert video_inputs["video"][0] == "VIDEO" +def test_video_node_trims_audio_to_native_frame_duration(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) + video = types.SimpleNamespace( + get_components=lambda: types.SimpleNamespace( + images=torch.ones(8, 2, 2, 3), + audio={"waveform": torch.ones(1, 1, 16), "sample_rate": 8}, + frame_rate=8, + ) + ) + + target, _ = nodes.SAMAudioVideoSeparator().separate( + pipeline=pipeline, + video=video, + mask=torch.zeros(2, 2), + description="", + seed=0, + inference_steps=2, + chunk_duration=0, + chunk_overlap=0, + ) + + assert processor.calls[0]["audios"][0].shape[-1] == 8 + assert target["waveform"].shape[-1] == 8 + + +def test_reranking_is_lazy_and_passed_to_upstream(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"] + pipeline = _fake_pipeline(runtime) + loaded = [] + pipeline.ensure_judge_ranker = lambda management: loaded.append(management) + + nodes.SAMAudioTextSeparator().separate( + pipeline=pipeline, + audio={"waveform": torch.ones(1, 1, 8), "sample_rate": 8}, + description="test sound", + predict_spans=False, + seed=0, + inference_steps=2, + reranking_candidates=4, + ) + + assert len(loaded) == 1 + assert pipeline.model.calls[0]["reranking_candidates"] == 4 + + +def test_transformer_steps_report_progress_and_check_cancellation(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"] + + class SteppingModel(FakeModel): + def __init__(self): + super().__init__() + self.transformer = torch.nn.Identity() + + def separate(self, batch, **kwargs): + self.transformer(torch.zeros(1)) + self.transformer(torch.zeros(1)) + return super().separate(batch, **kwargs) + + pipeline = runtime.SAMAudioPipeline( + model=SteppingModel(), + processor=FakeProcessor(), + patcher=object(), + device=torch.device("cpu"), + source="fake", + ) + + nodes.SAMAudioTextSeparator().separate( + pipeline=pipeline, + audio={"waveform": torch.ones(1, 1, 8), "sample_rate": 8}, + description="test sound", + predict_spans=False, + seed=0, + inference_steps=2, + ) + + assert management.progress_totals == [2] + assert sum(management.progress_updates) == 2 + assert management.interrupt_checks == 4 + assert not pipeline.model.transformer._forward_pre_hooks + + def test_video_node_requires_embedded_audio(monkeypatch, tmp_path): package, _ = _load_node_package(monkeypatch, tmp_path) nodes = sys.modules[f"{package.__name__}.nodes"] diff --git a/tests/test_package.py b/tests/test_package.py index dfbed79..17c2b00 100644 --- a/tests/test_package.py +++ b/tests/test_package.py @@ -44,19 +44,19 @@ def test_example_sampler_widgets_include_seed_control(): expected = { "text_separation_workflow.json": ( "SAMAudioTextSeparator", - ["man speaking", False, 0, "fixed", 32, 10.0, 1.0], + ["man speaking", False, 0, "fixed", 32, 10.0, 1.0, 1], ), "span_separation_workflow.json": ( "SAMAudioSpanSeparator", - ["car honking", 0, "fixed", 32, 10.0, 1.0], + ["car honking", 0, "fixed", 32, 10.0, 1.0, 1], ), "visual_separation_workflow.json": ( "SAMAudioVisualSeparator", - ["", 0, "fixed", 32, 10.0, 1.0], + ["", 0, "fixed", 32, 10.0, 1.0, 1], ), "video_separation_workflow.json": ( "SAMAudioVideoSeparator", - ["", 0, "fixed", 32, 10.0, 1.0], + ["", 0, "fixed", 32, 10.0, 1.0, 1], ), } for filename, (node_type, widget_values) in expected.items(): @@ -132,3 +132,20 @@ def test_existing_torchcodec_modules_are_restored(monkeypatch): assert sys.modules["torchcodec"] is original assert sys.modules["torchcodec.decoders"] is original_decoders + + +def test_lazy_torchcodec_decoder_resolves_real_file_io_on_demand(monkeypatch): + calls = [] + real_decoders = types.SimpleNamespace( + AudioDecoder=lambda *args, **kwargs: calls.append((args, kwargs)) or "decoded" + ) + monkeypatch.setattr( + torchcodec_compat.importlib, + "import_module", + lambda name: real_decoders, + ) + + result = torchcodec_compat.LazyAudioDecoder("input.wav", sample_rate=48_000) + + assert result == "decoded" + assert calls == [(("input.wav",), {"sample_rate": 48_000})] diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 968e143..e780442 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -1,5 +1,7 @@ import hashlib +import sys import types +from dataclasses import dataclass from pathlib import Path import pytest @@ -54,6 +56,42 @@ def test_official_download_directories_are_not_duplicated_as_local(tmp_path): assert discover_local_models(folders) == {} +def test_official_model_can_live_in_a_secondary_comfy_model_root(tmp_path, monkeypatch): + primary = tmp_path / "primary" + secondary = tmp_path / "secondary" + folders = FakeFolderPaths(primary) + register_model_folder(folders) + folders.paths.append(str(secondary / "sam_audio")) + expected = secondary / "sam_audio" / "sam-audio-base" + _make_model(expected) + monkeypatch.setattr( + runtime, + "_download_official_model", + lambda *_args: pytest.fail("an existing secondary-root model was downloaded"), + ) + + assert runtime.resolve_model("facebook/sam-audio-base", folders) == expected + + +def test_partial_mirror_in_secondary_root_is_not_loaded(tmp_path, monkeypatch): + primary = tmp_path / "primary" + secondary = tmp_path / "secondary" + folders = FakeFolderPaths(primary) + register_model_folder(folders) + folders.paths.append(str(secondary / "sam_audio")) + partial = secondary / "sam_audio" / "sam-audio-base" + _make_model(partial) + (partial / runtime.MIRROR_MARKER).write_text("pending checksum\n") + downloaded = primary / "sam_audio" / "sam-audio-base" + monkeypatch.setattr( + runtime, + "_download_official_model", + lambda _name, _root: downloaded, + ) + + assert runtime.resolve_model("facebook/sam-audio-base", folders) == downloaded + + def test_gated_download_uses_pinned_verified_mirror(tmp_path, monkeypatch): class GatedRepoError(Exception): pass @@ -193,7 +231,67 @@ def test_checkpoint_loader_uses_upstream_weight_loading_path(tmp_path): assert torch.equal(model.weight, torch.ones(1)) +def test_efficient_checkpoint_loader_assigns_and_casts_local_weights(tmp_path): + @dataclass + class Config: + width: int + + class EfficientModel(torch.nn.Module): + config_cls = Config + initial_device = None + + def __init__(self, config): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(config.width)) + type(self).initial_device = self.weight.device + + (tmp_path / "config.json").write_text('{"width": 2}') + torch.save({"weight": torch.tensor([1.0, 2.0])}, tmp_path / "checkpoint.pt") + + model = runtime._load_local_checkpoint_efficient( + EfficientModel, + tmp_path, + config_overrides={}, + target_device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + assert model.weight.dtype == torch.bfloat16 + assert torch.equal(model.weight.float(), torch.tensor([1.0, 2.0])) + assert EfficientModel.initial_device == torch.device("meta") + + +def test_non_meta_checkpoint_loader_preserves_derived_buffers(tmp_path): + @dataclass + class Config: + width: int + + class BufferedModel(torch.nn.Module): + config_cls = Config + + def __init__(self, config): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(config.width)) + self.register_buffer("derived", torch.arange(config.width), persistent=False) + + (tmp_path / "config.json").write_text('{"width": 2}') + torch.save({"weight": torch.tensor([1.0, 2.0])}, tmp_path / "checkpoint.pt") + + model = runtime._load_local_checkpoint_efficient( + BufferedModel, + tmp_path, + config_overrides={}, + target_device=torch.device("cpu"), + dtype=torch.bfloat16, + use_meta_device=False, + ) + + assert torch.equal(model.weight.float(), torch.tensor([1.0, 2.0])) + assert torch.equal(model.derived, torch.tensor([0, 1])) + + def test_support_snapshot_uses_local_cache_without_network(monkeypatch): + runtime._cached_support_snapshot.cache_clear() calls = [] class FakeHub: @@ -214,9 +312,12 @@ def test_support_snapshot_uses_local_cache_without_network(monkeypatch): "local_files_only": True, } ] + assert runtime._cached_support_snapshot(spec) == "/cached/support-model" + assert len(calls) == 1 def test_support_snapshot_downloads_only_when_cache_is_absent(monkeypatch): + runtime._cached_support_snapshot.cache_clear() calls = [] class FakeHub: @@ -295,6 +396,41 @@ def test_pipeline_attaches_span_predictor_lazily(monkeypatch): assert patcher.size == 0 +def test_pipeline_attaches_judge_ranker_lazily(monkeypatch): + ranker = torch.nn.Linear(2, 2) + loads = [] + unloads = [] + monkeypatch.setattr( + runtime, + "_load_judge_ranker", + lambda dtype: loads.append(dtype) or ranker, + ) + management = types.SimpleNamespace( + unload_model_and_clones=lambda patcher, all_devices=False: unloads.append( + (patcher, all_devices) + ) + ) + patcher = types.SimpleNamespace(size=123) + model = torch.nn.Linear(2, 2) + model.text_ranker = None + pipeline = runtime.SAMAudioPipeline( + model=model, + processor=object(), + patcher=patcher, + device=torch.device("cpu"), + source="fake", + dtype=torch.bfloat16, + ) + + pipeline.ensure_judge_ranker(management) + pipeline.ensure_judge_ranker(management) + + assert loads == [torch.bfloat16] + assert unloads == [(patcher, True)] + assert model.text_ranker is ranker + assert patcher.size == 0 + + def test_local_checkpoint_does_not_hide_unrelated_type_errors(tmp_path): class FakeSAMAudio: @classmethod @@ -305,6 +441,24 @@ def test_local_checkpoint_does_not_hide_unrelated_type_errors(tmp_path): runtime._load_sam_audio_checkpoint(FakeSAMAudio, tmp_path) +def test_import_decoder_placeholders_are_replaced_with_lazy_file_io(monkeypatch): + processor = types.ModuleType("sam_audio.processor") + processor.AudioDecoder = runtime.torchcodec_compat.AudioDecoder + processor.VideoDecoder = runtime.torchcodec_compat.VideoDecoder + transforms = types.ModuleType("core.audio_visual_encoder.transforms") + transforms.AudioDecoder = runtime.torchcodec_compat.AudioDecoder + transforms.VideoDecoder = runtime.torchcodec_compat.VideoDecoder + monkeypatch.setitem(sys.modules, "sam_audio.processor", processor) + monkeypatch.setitem(sys.modules, "core.audio_visual_encoder.transforms", transforms) + + runtime._install_lazy_torchcodec_decoders() + + assert processor.AudioDecoder is runtime.torchcodec_compat.LazyAudioDecoder + assert processor.VideoDecoder is runtime.torchcodec_compat.LazyVideoDecoder + assert transforms.AudioDecoder is runtime.torchcodec_compat.LazyAudioDecoder + assert transforms.VideoDecoder is runtime.torchcodec_compat.LazyVideoDecoder + + def test_ode_options_match_upstream_default(): assert ode_options(32) == { "method": "midpoint", @@ -323,6 +477,54 @@ def test_span_validation_requires_positive_and_checks_duration(): validate_spans((SpanPrompt("+", 0.25, 1.25),), duration=1.0) +def test_span_validation_rejects_unrepresentable_and_conflicting_spans(): + with pytest.raises(ValueError, match="effective 0.040s"): + validate_spans( + (SpanPrompt("+", 0.001, 0.01),), + duration=1.0, + hop_length=1920, + sample_rate=48_000, + ) + with pytest.raises(ValueError, match="cannot overlap"): + validate_spans( + (SpanPrompt("+", 0.1, 0.5), SpanPrompt("-", 0.4, 0.8)), + duration=1.0, + ) + + +def test_precision_selection_uses_bf16_only_on_supported_gpu(): + management = types.SimpleNamespace( + should_use_bf16=lambda device: device.type == "cuda", + supports_dtype=lambda device, dtype: dtype == torch.bfloat16, + ) + + assert ( + runtime.resolve_precision("auto", torch.device("cuda"), management) + == torch.bfloat16 + ) + assert ( + runtime.resolve_precision("auto", torch.device("cpu"), management) + == torch.float32 + ) + with pytest.raises(RuntimeError, match="requires a supported GPU"): + runtime.resolve_precision("bf16", torch.device("cpu"), management) + + +def test_gpu_initial_load_rejects_cpu_only_comfy(monkeypatch, tmp_path): + management = types.SimpleNamespace( + get_torch_device=lambda: torch.device("cpu") + ) + + with pytest.raises(RuntimeError, match="no GPU device"): + runtime.load_pipeline( + "facebook/sam-audio-small", + FakeFolderPaths(tmp_path), + management, + model_patcher=object(), + initial_device="gpu", + ) + + def test_seed_context_restores_cpu_rng(): from sam_audio_comfy.runtime import seeded_inference