From 50ffce7c883b243f122817da9e129d5fee71dfbc Mon Sep 17 00:00:00 2001 From: Ethanfel <28683423+ethanfel@users.noreply.github.com> Date: Sun, 23 Aug 2026 15:50:30 +0200 Subject: [PATCH] Fix current Hub loading and demo sampler values --- examples/span_separation_workflow.json | 2 +- examples/text_separation_workflow.json | 2 +- examples/visual_separation_workflow.json | 2 +- nodes.py | 1 + sam_audio_comfy/runtime.py | 44 ++++++++++++++++++++---- tests/test_nodes.py | 4 +++ tests/test_package.py | 23 +++++++++++++ tests/test_runtime.py | 38 ++++++++++++++++++++ 8 files changed, 106 insertions(+), 10 deletions(-) diff --git a/examples/span_separation_workflow.json b/examples/span_separation_workflow.json index c825e1d..88347cf 100644 --- a/examples/span_separation_workflow.json +++ b/examples/span_separation_workflow.json @@ -65,7 +65,7 @@ {"name": "residual", "type": "AUDIO", "links": [5], "slot_index": 1} ], "properties": {"Node name for S&R": "SAMAudioSpanSeparator"}, - "widgets_values": ["car honking", 0, 32] + "widgets_values": ["car honking", 0, "fixed", 32] }, { "id": 5, diff --git a/examples/text_separation_workflow.json b/examples/text_separation_workflow.json index 3ab9594..a980428 100644 --- a/examples/text_separation_workflow.json +++ b/examples/text_separation_workflow.json @@ -49,7 +49,7 @@ {"name": "residual", "type": "AUDIO", "links": [4], "slot_index": 1} ], "properties": {"Node name for S&R": "SAMAudioTextSeparator"}, - "widgets_values": ["man speaking", false, 0, 32] + "widgets_values": ["man speaking", false, 0, "fixed", 32] }, { "id": 4, diff --git a/examples/visual_separation_workflow.json b/examples/visual_separation_workflow.json index 8708840..a2ab336 100644 --- a/examples/visual_separation_workflow.json +++ b/examples/visual_separation_workflow.json @@ -67,7 +67,7 @@ {"name": "residual", "type": "AUDIO", "links": [6], "slot_index": 1} ], "properties": {"Node name for S&R": "SAMAudioVisualSeparator"}, - "widgets_values": ["", 0, 32] + "widgets_values": ["", 0, "fixed", 32] }, { "id": 5, diff --git a/nodes.py b/nodes.py index 9f60a8f..72a67ef 100644 --- a/nodes.py +++ b/nodes.py @@ -41,6 +41,7 @@ def _sampling_inputs() -> dict[str, tuple[Any, ...]]: "default": 0, "min": 0, "max": MAX_SEED, + "control_after_generate": True, "tooltip": "Controls SAM-Audio's initial noise for reproducible separation.", }, ), diff --git a/sam_audio_comfy/runtime.py b/sam_audio_comfy/runtime.py index 8541263..5f87e5e 100644 --- a/sam_audio_comfy/runtime.py +++ b/sam_audio_comfy/runtime.py @@ -320,6 +320,42 @@ def import_sam_audio() -> tuple[Any, Any]: ) from error +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 = { + "map_location": "cpu", + "strict": False, + "text_ranker": None, + "visual_ranker": None, + } + try: + return SAMAudio.from_pretrained(str(model_path), **model_kwargs) + except TypeError as error: + message = str(error) + if not ( + "_from_pretrained()" in message + and "proxies" in message + and "resume_download" in message + ): + raise + + LOGGER.info( + "Using SAM-Audio's legacy local-checkpoint loader with the current " + "huggingface_hub API" + ) + return SAMAudio._from_pretrained( + model_id=str(model_path), + revision=None, + cache_dir=None, + force_download=False, + proxies=None, + resume_download=False, + local_files_only=True, + token=None, + **model_kwargs, + ) + + def load_pipeline( model_name: str, folder_paths: Any, @@ -334,13 +370,7 @@ def load_pipeline( LOGGER.info("Loading SAM-Audio model from %s", model_path) try: - model = SAMAudio.from_pretrained( - str(model_path), - map_location="cpu", - strict=False, - text_ranker=None, - visual_ranker=None, - ) + model = _load_sam_audio_checkpoint(SAMAudio, model_path) processor = SAMAudioProcessor.from_pretrained(str(model_path)) except Exception as error: raise RuntimeError( diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 262b869..cb208cc 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -97,6 +97,10 @@ def test_text_node_runs_full_audio_batch_with_standard_shapes(monkeypatch, tmp_p "SAMAudioPipelineLoader" ].INPUT_TYPES()["required"] assert loader_inputs["attention_backend"][0] == ["pytorch", "comfy_kitchen"] + sampler_inputs = package.NODE_CLASS_MAPPINGS[ + "SAMAudioTextSeparator" + ].INPUT_TYPES()["required"] + assert sampler_inputs["seed"][1]["control_after_generate"] is True assert set(package.NODE_CLASS_MAPPINGS) == { "SAMAudioPipelineLoader", "SAMAudioSpanPrompt", diff --git a/tests/test_package.py b/tests/test_package.py index a281ded..ffd0d08 100644 --- a/tests/test_package.py +++ b/tests/test_package.py @@ -1,3 +1,4 @@ +import json import sys import types from pathlib import Path @@ -38,6 +39,28 @@ def test_node_does_not_require_torchcodec(): assert not any("torchcodec" in line for line in dependency_lines) +def test_example_sampler_widgets_include_seed_control(): + root = Path(__file__).parents[1] + expected = { + "text_separation_workflow.json": ( + "SAMAudioTextSeparator", + ["man speaking", False, 0, "fixed", 32], + ), + "span_separation_workflow.json": ( + "SAMAudioSpanSeparator", + ["car honking", 0, "fixed", 32], + ), + "visual_separation_workflow.json": ( + "SAMAudioVisualSeparator", + ["", 0, "fixed", 32], + ), + } + for filename, (node_type, widget_values) in expected.items(): + workflow = json.loads((root / "examples" / filename).read_text()) + sampler = next(node for node in workflow["nodes"] if node["type"] == node_type) + assert sampler["widgets_values"] == widget_values + + def test_upstream_installer_uses_no_deps(): root = Path(__file__).parents[1] installer = (root / "install.py").read_text().lower() diff --git a/tests/test_runtime.py b/tests/test_runtime.py index cbec812..2836a18 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -141,6 +141,44 @@ def test_mirror_checksum_mismatch_is_rejected(tmp_path, monkeypatch): runtime._download_official_model("facebook/sam-audio-small", tmp_path) +def test_local_checkpoint_falls_back_for_legacy_hub_signature(tmp_path): + calls = [] + + class FakeSAMAudio: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + calls.append(("public", model_id, kwargs)) + raise TypeError( + "BaseModel._from_pretrained() missing 2 required keyword-only " + "arguments: 'proxies' and 'resume_download'" + ) + + @classmethod + def _from_pretrained(cls, **kwargs): + calls.append(("legacy", kwargs)) + return "model" + + assert runtime._load_sam_audio_checkpoint(FakeSAMAudio, tmp_path) == "model" + assert calls[1][0] == "legacy" + assert calls[1][1]["model_id"] == str(tmp_path) + assert calls[1][1]["local_files_only"] is True + assert calls[1][1]["proxies"] is None + assert calls[1][1]["resume_download"] is False + assert calls[1][1]["strict"] is False + assert calls[1][1]["text_ranker"] is None + assert calls[1][1]["visual_ranker"] is None + + +def test_local_checkpoint_does_not_hide_unrelated_type_errors(tmp_path): + class FakeSAMAudio: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + raise TypeError("bad model config") + + with pytest.raises(TypeError, match="bad model config"): + runtime._load_sam_audio_checkpoint(FakeSAMAudio, tmp_path) + + def test_ode_options_match_upstream_default(): assert ode_options(32) == { "method": "midpoint",