nodes/ethanfel-ComfyUI-SAM-Audio
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Harden SAM-Audio loading and inference

Browse Source
This commit is contained in:
Ethanfel
2026-09-06 22:59:57 +02:00
parent dca2f22766
commit 68239634ea
19 changed files with 1179 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/publish_action.yml
+23
View File
@@ -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
.github/workflows/test.yml
+38
View File
@@ -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
README.md
+14 -2
View File
@@ -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
examples/span_separation_workflow.json
+2 -2
View File
@@ -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,
examples/text_separation_workflow.json
+2 -2
View File
@@ -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,
examples/video_separation_workflow.json
+2 -2
View File
@@ -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,
examples/visual_separation_workflow.json
+2 -2
View File
@@ -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,
install.py
+6
View File
@@ -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)
nodes.py
+198 -56
View File
@@ -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,
)
pyproject.toml
+1 -1
View File
@@ -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"
sam_audio_comfy/attention.py
+1 -1
View File
@@ -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,
sam_audio_comfy/audio.py
+48 -4
View File
@@ -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()
sam_audio_comfy/runtime.py
+414 -28
View File
@@ -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"
)
sam_audio_comfy/torchcodec_compat.py
+27
View File
@@ -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__ = []
tests/test_attention.py
+22
View File
@@ -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)
tests/test_audio.py
+34
View File
@@ -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
tests/test_nodes.py
+122 -3
View File
@@ -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"]
tests/test_package.py
+21 -4
View File
@@ -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})]
tests/test_runtime.py
+202
View File
@@ -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