import hashlib import sys import types from dataclasses import dataclass from pathlib import Path import pytest import torch from sam_audio_comfy import runtime from sam_audio_comfy.runtime import ( SpanPrompt, discover_local_models, ode_options, register_model_folder, validate_spans, ) class FakeFolderPaths: def __init__(self, root: Path): self.models_dir = str(root) self.paths = [] def add_model_folder_path(self, name, path, is_default=False): assert name == "sam_audio" if path not in self.paths: self.paths.append(path) def get_folder_paths(self, name): assert name == "sam_audio" return self.paths def _make_model(path: Path): path.mkdir(parents=True) (path / "config.json").write_text("{}") (path / "checkpoint.pt").write_bytes(b"checkpoint") def test_register_and_discover_nested_local_models(tmp_path): folders = FakeFolderPaths(tmp_path) model_root = register_model_folder(folders) _make_model(model_root / "custom" / "my-model") assert discover_local_models(folders) == { "local/custom/my-model": model_root / "custom" / "my-model" } def test_official_download_directories_are_not_duplicated_as_local(tmp_path): folders = FakeFolderPaths(tmp_path) model_root = register_model_folder(folders) _make_model(model_root / "sam-audio-large") 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 checkpoint = b"verified mirror checkpoint" spec = runtime.MirrorSpec( repo_id="mirror/sam-audio-small", revision="pinned-revision", checkpoint_size=len(checkpoint), checkpoint_sha256=hashlib.sha256(checkpoint).hexdigest(), ) monkeypatch.setitem(runtime.MIRROR_MODELS, "facebook/sam-audio-small", spec) calls = [] class FakeHub: @staticmethod def snapshot_download(**kwargs): calls.append(kwargs) if kwargs["repo_id"] == "facebook/sam-audio-small": raise GatedRepoError("access denied") destination = Path(kwargs["local_dir"]) (destination / "config.json").write_text("{}") (destination / "checkpoint.pt").write_bytes(checkpoint) (destination / "LICENSE").write_text("SAM License") return str(destination) monkeypatch.setattr(runtime.importlib, "import_module", lambda _: FakeHub) result = runtime._download_official_model( "facebook/sam-audio-small", tmp_path ) assert result == tmp_path / "sam-audio-small" assert calls[1]["repo_id"] == spec.repo_id assert calls[1]["revision"] == spec.revision assert calls[1]["allow_patterns"] == [ "config.json", "checkpoint.pt", "LICENSE", ] assert (result / runtime.MIRROR_MARKER).read_text() == ( f"{len(checkpoint)} {spec.checkpoint_sha256}\n" ) def test_non_gate_download_error_does_not_use_mirror(tmp_path, monkeypatch): calls = [] class FakeHub: @staticmethod def snapshot_download(**kwargs): calls.append(kwargs) raise ConnectionError("offline") monkeypatch.setattr(runtime.importlib, "import_module", lambda _: FakeHub) with pytest.raises(RuntimeError, match="Could not download"): runtime._download_official_model("facebook/sam-audio-small", tmp_path) assert [call["repo_id"] for call in calls] == ["facebook/sam-audio-small"] def test_mirror_checksum_mismatch_is_rejected(tmp_path, monkeypatch): class GatedRepoError(Exception): pass spec = runtime.MirrorSpec( repo_id="mirror/sam-audio-small", revision="pinned-revision", checkpoint_size=3, checkpoint_sha256=hashlib.sha256(b"good").hexdigest(), ) monkeypatch.setitem(runtime.MIRROR_MODELS, "facebook/sam-audio-small", spec) class FakeHub: @staticmethod def snapshot_download(**kwargs): if kwargs["repo_id"] == "facebook/sam-audio-small": raise GatedRepoError("access denied") destination = Path(kwargs["local_dir"]) (destination / "config.json").write_text("{}") (destination / "checkpoint.pt").write_bytes(b"bad") (destination / "LICENSE").write_text("SAM License") return str(destination) monkeypatch.setattr(runtime.importlib, "import_module", lambda _: FakeHub) with pytest.raises(RuntimeError, match="verified public mirror"): 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 True assert calls[1][1]["text_ranker"] is None assert calls[1][1]["visual_ranker"] is None def test_checkpoint_loader_uses_upstream_weight_loading_path(tmp_path): class UpstreamStyleModel(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.nn.Parameter(torch.zeros(1)) def load_state_dict(self, state_dict, strict=True): # SAM-Audio's pinned implementation has no false branch. if strict: return super().load_state_dict(state_dict, strict=False) return None @classmethod def from_pretrained(cls, model_id, **kwargs): model = cls() model.load_state_dict({"weight": torch.ones(1)}, strict=kwargs["strict"]) return model model = runtime._load_sam_audio_checkpoint(UpstreamStyleModel, 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: @staticmethod def snapshot_download(**kwargs): calls.append(kwargs) return "/cached/support-model" monkeypatch.setattr(runtime.importlib, "import_module", lambda _: FakeHub) spec = runtime.SupportModelSpec("owner/model", "pinned", ("config.json",)) assert runtime._cached_support_snapshot(spec) == "/cached/support-model" assert calls == [ { "repo_id": "owner/model", "revision": "pinned", "allow_patterns": ["config.json"], "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: @staticmethod def snapshot_download(**kwargs): calls.append(kwargs) if kwargs.get("local_files_only"): raise FileNotFoundError("not cached") return "/downloaded/support-model" monkeypatch.setattr(runtime.importlib, "import_module", lambda _: FakeHub) spec = runtime.SupportModelSpec("owner/model", "pinned") assert runtime._cached_support_snapshot(spec) == "/downloaded/support-model" assert [call.get("local_files_only", False) for call in calls] == [True, False] 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"}' ) monkeypatch.setattr( runtime, "_cached_support_snapshot", lambda spec: f"/cache/{spec.repo_id.replace('/', '--')}", ) kwargs = runtime._sam_audio_model_kwargs(tmp_path) assert kwargs["text_encoder"] == { "name": "/cache/google-t5--t5-base", "dim": 768, } assert runtime.SUPPORT_MODELS["t5-base"].repo_id == "google-t5/t5-base" 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_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 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_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", "options": {"step_size": 2.0 / 32}, } with pytest.raises(ValueError, match="even"): ode_options(31) def test_span_validation_requires_positive_and_checks_duration(): validate_spans((SpanPrompt("+", 0.25, 0.75),), duration=1.0) with pytest.raises(ValueError, match="positive"): validate_spans((SpanPrompt("-", 0.25, 0.75),), duration=1.0) with pytest.raises(ValueError, match="only 1.000s"): 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 torch.manual_seed(123) expected_first = torch.rand(1) expected_second = torch.rand(1) torch.manual_seed(123) actual_first = torch.rand(1) with seeded_inference(999, torch.device("cpu")): inside_first = torch.rand(1) actual_second = torch.rand(1) torch.manual_seed(999) assert torch.equal(inside_first, torch.rand(1)) assert torch.equal(actual_first, expected_first) assert torch.equal(actual_second, expected_second)