342 lines
11 KiB
Python
342 lines
11 KiB
Python
import hashlib
|
|
import types
|
|
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_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_support_snapshot_uses_local_cache_without_network(monkeypatch):
|
|
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,
|
|
}
|
|
]
|
|
|
|
|
|
def test_support_snapshot_downloads_only_when_cache_is_absent(monkeypatch):
|
|
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/t5-base",
|
|
"dim": 768,
|
|
}
|
|
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
|
|
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",
|
|
"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_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)
|