Integrate model lifecycle and lazy span loading

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
 

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

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"

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),
     )
 
 

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):

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