diff --git a/README.md b/README.md index 9f47d74..607a528 100644 --- a/README.md +++ b/README.md @@ -39,9 +39,9 @@ You can authenticate in the ComfyUI environment if you prefer the official repos hf auth login ``` -Models download on first use to `ComfyUI/models/sam_audio/`. Both official and mirrored checkpoints remain subject to Meta's SAM License. +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 and PE span-predictor assets are pinned in the Hugging Face cache and reused locally after their first download. +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. ## Nodes @@ -71,6 +71,7 @@ Example workflows are included in the [`examples`](examples/) directory. - White pixels in a visual mask identify the target object. - Local model folders must contain `config.json` and `checkpoint.pt`. - The model loader offers PyTorch SDPA (default) or Comfy Kitchen INT8 attention. Comfy Kitchen requires a supported current ComfyUI build and GPU and may slightly change results because its attention is quantized. +- 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. ## License diff --git a/nodes.py b/nodes.py index 9183d80..37a0af4 100644 --- a/nodes.py +++ b/nodes.py @@ -148,6 +148,9 @@ def _run_separation( if spans is not None: validate_spans(spans, duration) + if predict_spans: + pipeline.ensure_span_predictor(model_management) + masked_video = None if visual_prompt is not None: frames, mask = visual_prompt diff --git a/pyproject.toml b/pyproject.toml index 760fc51..0e326b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "comfyui-sam-audio" -version = "0.2.3" +version = "0.2.4" description = "Native ComfyUI nodes for local Meta SAM-Audio sound separation" readme = "README.md" requires-python = ">=3.11" diff --git a/sam_audio_comfy/runtime.py b/sam_audio_comfy/runtime.py index 4f0dbba..e5a45c4 100644 --- a/sam_audio_comfy/runtime.py +++ b/sam_audio_comfy/runtime.py @@ -6,6 +6,7 @@ import importlib import json import logging import os +import time from dataclasses import dataclass from pathlib import Path from typing import Any, Iterator, Literal, NamedTuple @@ -118,14 +119,50 @@ class SAMAudioPipeline: device: torch.device source: str attention_backend: str = "pytorch" + span_predictor_source: str | None = None def get_models(self) -> list[Any]: """Let ComfyUI track this model for the lifetime of the prompt.""" return [self.patcher] def load(self, model_management: Any) -> None: + started = time.perf_counter() model_management.load_models_gpu([self.patcher], force_full_load=True) self.model.eval() + elapsed = time.perf_counter() - started + if elapsed >= 1.0: + LOGGER.info("Loaded SAM-Audio weights onto %s in %.1f seconds", self.device, elapsed) + + def ensure_span_predictor(self, model_management: Any) -> None: + """Attach the optional PE span predictor only when it is requested.""" + if hasattr(self.model, "span_predictor"): + return + if self.span_predictor_source is None: + raise RuntimeError( + "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) + + 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_transform = transform + if hasattr(self.patcher, "size"): + self.patcher.size = 0 + LOGGER.info( + "Loaded optional SAM-Audio span predictor in %.1f seconds", + time.perf_counter() - started, + ) def register_model_folder(folder_paths: Any) -> Path: @@ -391,13 +428,38 @@ def _sam_audio_model_kwargs(model_path: Path) -> dict[str, Any]: text_encoder["name"] = _cached_support_snapshot(SUPPORT_MODELS["t5-base"]) model_kwargs["text_encoder"] = text_encoder - if config.get("span_predictor") == "pe-a-frame-large": - model_kwargs["span_predictor"] = _cached_support_snapshot( - SUPPORT_MODELS["pe-a-frame-large"] - ) + # The PE span predictor is a separate 5.8 GiB model and is not required for + # ordinary text, explicit-span, or visual separation. It is attached lazily + # by SAMAudioPipeline.ensure_span_predictor when predict_spans is enabled. + if config.get("span_predictor") is not None: + model_kwargs["span_predictor"] = None return model_kwargs +def _configured_span_predictor(model_path: Path) -> str | None: + config_path = model_path / "config.json" + if not config_path.is_file(): + return None + value = json.loads(config_path.read_text(encoding="utf-8")).get( + "span_predictor" + ) + return value if isinstance(value, str) and value else None + + +def _load_span_predictor_assets(source: str) -> tuple[torch.nn.Module, Any]: + if source == "pe-a-frame-large": + source = _cached_support_snapshot(SUPPORT_MODELS["pe-a-frame-large"]) + try: + module = importlib.import_module("core.audio_visual_encoder") + predictor = module.PEAudioFrame.from_config(source, pretrained=True) + transform = module.PEAudioFrameTransform.from_config(source) + except Exception as error: + raise RuntimeError( + f"Failed to load SAM-Audio's optional span predictor from {source}" + ) from error + 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.""" model_kwargs = _sam_audio_model_kwargs(model_path) @@ -442,8 +504,13 @@ def load_pipeline( SAMAudio, SAMAudioProcessor = import_sam_audio() LOGGER.info("Loading SAM-Audio model from %s", model_path) + started = time.perf_counter() try: model = _load_sam_audio_checkpoint(SAMAudio, model_path) + LOGGER.info( + "Loaded SAM-Audio checkpoint and required text encoder on CPU in %.1f seconds", + time.perf_counter() - started, + ) processor = SAMAudioProcessor.from_pretrained(str(model_path)) except Exception as error: raise RuntimeError( @@ -467,6 +534,7 @@ def load_pipeline( device=load_device, source=model_name, attention_backend=attention_backend, + span_predictor_source=_configured_span_predictor(model_path), ) diff --git a/tests/test_nodes.py b/tests/test_nodes.py index d711445..b3a7fda 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -11,7 +11,10 @@ def _load_node_package(monkeypatch, tmp_path): comfy = types.ModuleType("comfy") comfy.__path__ = [] management = types.ModuleType("comfy.model_management") - management.load_models_gpu = lambda models, force_full_load=False: None + management.load_calls = [] + management.load_models_gpu = lambda models, force_full_load=False: ( + management.load_calls.append((models, force_full_load)) + ) patcher = types.ModuleType("comfy.model_patcher") utils = types.ModuleType("comfy.utils") utils.ProgressBar = lambda total: types.SimpleNamespace(update=lambda value: None) @@ -125,7 +128,7 @@ def test_text_node_runs_full_audio_batch_with_standard_shapes(monkeypatch, tmp_p "SAMAudioSpanSeparator", "SAMAudioVisualSeparator", } - assert management.load_models_gpu is not None + assert management.load_calls == [([pipeline.patcher], True)] def _fake_pipeline(runtime, processor=None): diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 3b14d20..209eaf5 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -1,4 +1,5 @@ import hashlib +import types from pathlib import Path import pytest @@ -233,7 +234,7 @@ def test_support_snapshot_downloads_only_when_cache_is_absent(monkeypatch): assert [call.get("local_files_only", False) for call in calls] == [True, False] -def test_model_kwargs_replace_support_ids_with_local_snapshots(tmp_path, monkeypatch): +def test_model_kwargs_resolve_t5_but_defer_span_predictor(tmp_path, monkeypatch): (tmp_path / "config.json").write_text( '{"text_encoder":{"name":"t5-base","dim":768},' '"span_predictor":"pe-a-frame-large"}' @@ -250,12 +251,49 @@ def test_model_kwargs_replace_support_ids_with_local_snapshots(tmp_path, monkeyp "name": "/cache/t5-base", "dim": 768, } - assert kwargs["span_predictor"] == "/cache/facebook--pe-a-frame-large" + assert kwargs["span_predictor"] is None + assert runtime._configured_span_predictor(tmp_path) == "pe-a-frame-large" assert kwargs["text_ranker"] is None assert kwargs["visual_ranker"] is None assert kwargs["strict"] is True +def test_pipeline_attaches_span_predictor_lazily(monkeypatch): + predictor = torch.nn.Linear(2, 2) + transform = object() + loads = [] + unloads = [] + monkeypatch.setattr( + runtime, + "_load_span_predictor_assets", + lambda source: loads.append(source) or (predictor, transform), + ) + 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) + pipeline = runtime.SAMAudioPipeline( + model=model, + processor=object(), + patcher=patcher, + device=torch.device("cpu"), + source="fake", + span_predictor_source="pe-a-frame-large", + ) + + pipeline.ensure_span_predictor(management) + pipeline.ensure_span_predictor(management) + + assert loads == ["pe-a-frame-large"] + assert unloads == [(patcher, True)] + assert model.span_predictor is predictor + assert model.span_predictor_transform is transform + assert patcher.size == 0 + + def test_local_checkpoint_does_not_hide_unrelated_type_errors(tmp_path): class FakeSAMAudio: @classmethod