545 lines
18 KiB
Python
545 lines
18 KiB
Python
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)
|