573 lines
20 KiB
Python
573 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import hashlib
|
|
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
|
|
|
|
import torch
|
|
|
|
from .attention import validate_attention_backend
|
|
from .audiotools_compat import audiotools_import_compatibility
|
|
from .torchcodec_compat import torchcodec_import_compatibility
|
|
from .xformers_compat import xformers_import_compatibility
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
OFFICIAL_MODELS = (
|
|
"facebook/sam-audio-small",
|
|
"facebook/sam-audio-base",
|
|
"facebook/sam-audio-large",
|
|
"facebook/sam-audio-small-tv",
|
|
"facebook/sam-audio-base-tv",
|
|
"facebook/sam-audio-large-tv",
|
|
)
|
|
|
|
MODEL_FOLDER = "sam_audio"
|
|
MODEL_FILES = ("config.json", "checkpoint.pt")
|
|
MIRROR_MARKER = ".verified-mirror-checkpoint"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MirrorSpec:
|
|
repo_id: str
|
|
revision: str
|
|
checkpoint_size: int
|
|
checkpoint_sha256: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SupportModelSpec:
|
|
repo_id: str
|
|
revision: str
|
|
allow_patterns: tuple[str, ...] | None = None
|
|
|
|
|
|
# Each mirror revision is immutable, includes Meta's SAM License, and was checked
|
|
# against the corresponding facebook checkpoint.pt LFS SHA-256 on 2026-08-23.
|
|
MIRROR_MODELS = {
|
|
"facebook/sam-audio-small": MirrorSpec(
|
|
"mrfakename/sam-audio-small",
|
|
"682824171fd5cfec47b88687f583a11050c1bf1d",
|
|
5_100_547_943,
|
|
"8c44fda9821fd9f2ec8977304e3c0f55290d9eacb6bbf25b4b8fb1f69c2a8c06",
|
|
),
|
|
"facebook/sam-audio-base": MirrorSpec(
|
|
"mrfakename/sam-audio-base",
|
|
"351f6666d46357edea2149a338a680fcd0d7afdd",
|
|
7_725_405_659,
|
|
"b5f3e29ea7a9e80e90a00da495a8aafe890571f371c4bfb88c052c65a5636839",
|
|
),
|
|
"facebook/sam-audio-large": MirrorSpec(
|
|
"mrfakename/sam-audio-large",
|
|
"497809f191dd4f673a2c93557176387cf2a06cf2",
|
|
14_861_356_211,
|
|
"ca55418b1d23e8c8a4dcc55f259d9801c8f79da0131a66e525d862c1289e3c4f",
|
|
),
|
|
"facebook/sam-audio-small-tv": MirrorSpec(
|
|
"mrfakename/sam-audio-small-tv",
|
|
"2fee00b2836ede05fce503726ac9f908b1fc4aa9",
|
|
5_100_547_943,
|
|
"1a9693b235efc3176986664dac349eb64022c7468cf8d4195954490a2457e6a9",
|
|
),
|
|
"facebook/sam-audio-base-tv": MirrorSpec(
|
|
"mrfakename/sam-audio-base-tv",
|
|
"ed6a64b2d3c6151276d81783becb8881d75efe4f",
|
|
7_725_405_659,
|
|
"569e3fecdefe267047b02d117acb5a67c03d2f73c4943d3d8bd43df9e9b9148d",
|
|
),
|
|
"facebook/sam-audio-large-tv": MirrorSpec(
|
|
"mrfakename/sam-audio-large-tv",
|
|
"30eea7c915f43349f3b3f5c7e33e8658a01d7253",
|
|
14_861_356_211,
|
|
"90e047269238c498c5abe0da6e6ba40859d152111d4a09f582a209027a33b72f",
|
|
),
|
|
}
|
|
|
|
# SAM-Audio constructs these support models while loading. Resolve immutable
|
|
# snapshots once, then pass local paths upstream to prevent repeated Hub checks.
|
|
SUPPORT_MODELS = {
|
|
"t5-base": SupportModelSpec(
|
|
"t5-base",
|
|
"a9723ea7f1b39c1eae772870f3b547bf6ef7e6c1",
|
|
("config.json", "model.safetensors", "spiece.model", "tokenizer.json"),
|
|
),
|
|
"pe-a-frame-large": SupportModelSpec(
|
|
"facebook/pe-a-frame-large",
|
|
"40187271298f84e2966d4518c88dde698540c9ad",
|
|
),
|
|
}
|
|
|
|
|
|
class SpanPrompt(NamedTuple):
|
|
token: Literal["+", "-"]
|
|
start: float
|
|
end: float
|
|
|
|
|
|
@dataclass
|
|
class SAMAudioPipeline:
|
|
model: torch.nn.Module
|
|
processor: Any
|
|
patcher: Any
|
|
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:
|
|
model_dir = Path(folder_paths.models_dir) / MODEL_FOLDER
|
|
try:
|
|
folder_paths.add_model_folder_path(
|
|
MODEL_FOLDER, str(model_dir), is_default=True
|
|
)
|
|
except TypeError:
|
|
# ComfyUI versions before is_default support still accept the two-argument form.
|
|
folder_paths.add_model_folder_path(MODEL_FOLDER, str(model_dir))
|
|
return model_dir
|
|
|
|
|
|
def _is_model_directory(path: Path) -> bool:
|
|
return path.is_dir() and all((path / name).is_file() for name in MODEL_FILES)
|
|
|
|
|
|
def discover_local_models(folder_paths: Any) -> dict[str, Path]:
|
|
"""Return stable UI labels for valid model directories in every configured root."""
|
|
discovered: list[tuple[str, Path]] = []
|
|
official_directory_names = {name.rsplit("/", 1)[-1] for name in OFFICIAL_MODELS}
|
|
for raw_root in folder_paths.get_folder_paths(MODEL_FOLDER):
|
|
root = Path(raw_root)
|
|
if not root.is_dir():
|
|
continue
|
|
for current, directories, files in os.walk(root):
|
|
current_path = Path(current)
|
|
if all(name in files for name in MODEL_FILES):
|
|
if current_path.parent == root and current_path.name in official_directory_names:
|
|
directories[:] = []
|
|
continue
|
|
relative = current_path.relative_to(root).as_posix()
|
|
label = relative if relative != "." else current_path.name
|
|
discovered.append((f"local/{label}", current_path))
|
|
directories[:] = []
|
|
continue
|
|
directories[:] = [
|
|
name
|
|
for name in directories
|
|
if not name.startswith(".") and name != "__pycache__"
|
|
]
|
|
|
|
counts: dict[str, int] = {}
|
|
result: dict[str, Path] = {}
|
|
for label, path in sorted(discovered, key=lambda item: (item[0].lower(), str(item[1]))):
|
|
counts[label] = counts.get(label, 0) + 1
|
|
unique_label = label if counts[label] == 1 else f"{label} [{counts[label]}]"
|
|
result[unique_label] = path
|
|
return result
|
|
|
|
|
|
def model_choices(folder_paths: Any) -> list[str]:
|
|
return [*OFFICIAL_MODELS, *discover_local_models(folder_paths).keys()]
|
|
|
|
|
|
def _is_gate_error(error: Exception) -> bool:
|
|
"""Recognize Hugging Face gated/authorization failures across hub versions."""
|
|
current: BaseException | None = error
|
|
visited: set[int] = set()
|
|
while current is not None and id(current) not in visited:
|
|
visited.add(id(current))
|
|
if type(current).__name__ == "GatedRepoError":
|
|
return True
|
|
response = getattr(current, "response", None)
|
|
if getattr(response, "status_code", None) in (401, 403):
|
|
return True
|
|
current = current.__cause__ or current.__context__
|
|
return False
|
|
|
|
|
|
def _checkpoint_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as checkpoint:
|
|
for chunk in iter(lambda: checkpoint.read(8 * 1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _verify_mirror_checkpoint(destination: Path, spec: MirrorSpec) -> None:
|
|
checkpoint = destination / "checkpoint.pt"
|
|
marker = destination / MIRROR_MARKER
|
|
expected_marker = f"{spec.checkpoint_size} {spec.checkpoint_sha256}\n"
|
|
|
|
if (
|
|
marker.is_file()
|
|
and marker.read_text(encoding="utf-8") == expected_marker
|
|
and checkpoint.stat().st_size == spec.checkpoint_size
|
|
):
|
|
return
|
|
|
|
actual_size = checkpoint.stat().st_size
|
|
if actual_size != spec.checkpoint_size:
|
|
raise RuntimeError(
|
|
f"Mirror checkpoint size mismatch for {spec.repo_id}: expected "
|
|
f"{spec.checkpoint_size} bytes, received {actual_size}"
|
|
)
|
|
|
|
actual_hash = _checkpoint_sha256(checkpoint)
|
|
if actual_hash != spec.checkpoint_sha256:
|
|
raise RuntimeError(
|
|
f"Mirror checkpoint checksum mismatch for {spec.repo_id}; refusing to load it"
|
|
)
|
|
marker.write_text(expected_marker, encoding="utf-8")
|
|
|
|
|
|
def _download_mirror_model(
|
|
huggingface_hub: Any,
|
|
destination: Path,
|
|
spec: MirrorSpec,
|
|
) -> Path:
|
|
marker = destination / MIRROR_MARKER
|
|
marker.write_text(f"pending {spec.checkpoint_sha256}\n", encoding="utf-8")
|
|
snapshot = huggingface_hub.snapshot_download(
|
|
repo_id=spec.repo_id,
|
|
revision=spec.revision,
|
|
local_dir=str(destination),
|
|
allow_patterns=[*MODEL_FILES, "LICENSE"],
|
|
)
|
|
path = Path(snapshot)
|
|
if not _is_model_directory(path) or not (path / "LICENSE").is_file():
|
|
raise RuntimeError(
|
|
f"The mirror snapshot for {spec.repo_id} is incomplete; expected "
|
|
"config.json, checkpoint.pt, and LICENSE"
|
|
)
|
|
_verify_mirror_checkpoint(path, spec)
|
|
return path
|
|
|
|
|
|
def _download_official_model(repo_id: str, model_root: Path) -> Path:
|
|
destination = model_root / repo_id.rsplit("/", 1)[-1]
|
|
if _is_model_directory(destination):
|
|
marker = destination / MIRROR_MARKER
|
|
if marker.exists() and (destination / "LICENSE").is_file():
|
|
_verify_mirror_checkpoint(destination, MIRROR_MODELS[repo_id])
|
|
return destination
|
|
if not marker.exists():
|
|
return destination
|
|
|
|
destination.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
huggingface_hub = importlib.import_module("huggingface_hub")
|
|
except Exception as error:
|
|
raise RuntimeError(
|
|
"huggingface-hub is required to download SAM-Audio checkpoints; "
|
|
"run this node's install.py and restart ComfyUI"
|
|
) from error
|
|
|
|
try:
|
|
snapshot = huggingface_hub.snapshot_download(
|
|
repo_id=repo_id,
|
|
local_dir=str(destination),
|
|
allow_patterns=list(MODEL_FILES),
|
|
)
|
|
except Exception as error:
|
|
if _is_gate_error(error):
|
|
spec = MIRROR_MODELS[repo_id]
|
|
LOGGER.warning(
|
|
"Access to %s was gated; downloading checksum-verified mirror %s",
|
|
repo_id,
|
|
spec.repo_id,
|
|
)
|
|
try:
|
|
return _download_mirror_model(huggingface_hub, destination, spec)
|
|
except Exception as mirror_error:
|
|
raise RuntimeError(
|
|
f"Access to {repo_id} was denied and its verified public mirror "
|
|
f"{spec.repo_id} could not be downloaded"
|
|
) from mirror_error
|
|
raise RuntimeError(
|
|
f"Could not download {repo_id}. Accept the model license at "
|
|
f"https://huggingface.co/{repo_id}, run `hf auth login` in ComfyUI's "
|
|
"Python environment, and retry."
|
|
) from error
|
|
|
|
path = Path(snapshot)
|
|
if not _is_model_directory(path):
|
|
raise RuntimeError(
|
|
f"The downloaded snapshot for {repo_id} is incomplete; expected "
|
|
f"{', '.join(MODEL_FILES)} in {path}"
|
|
)
|
|
(path / MIRROR_MARKER).unlink(missing_ok=True)
|
|
return path
|
|
|
|
|
|
def resolve_model(model_name: str, folder_paths: Any) -> Path:
|
|
local_models = discover_local_models(folder_paths)
|
|
if model_name in local_models:
|
|
return local_models[model_name]
|
|
if model_name not in OFFICIAL_MODELS:
|
|
raise ValueError(
|
|
f"Unknown SAM-Audio model {model_name!r}. Refresh the browser after adding local models."
|
|
)
|
|
model_root = register_model_folder(folder_paths)
|
|
return _download_official_model(model_name, model_root)
|
|
|
|
|
|
def import_sam_audio() -> tuple[Any, Any]:
|
|
try:
|
|
LOGGER.info(
|
|
"Loading SAM-Audio through isolated inference-only compatibility layers"
|
|
)
|
|
with (
|
|
xformers_import_compatibility(),
|
|
audiotools_import_compatibility(),
|
|
torchcodec_import_compatibility(),
|
|
):
|
|
module = importlib.import_module("sam_audio")
|
|
return module.SAMAudio, module.SAMAudioProcessor
|
|
except Exception as error:
|
|
raise RuntimeError(
|
|
"SAM-Audio could not be imported. Install this node's requirements with "
|
|
"ComfyUI's Python interpreter, then restart ComfyUI."
|
|
) from error
|
|
|
|
|
|
def _cached_support_snapshot(spec: SupportModelSpec) -> str:
|
|
try:
|
|
huggingface_hub = importlib.import_module("huggingface_hub")
|
|
except Exception as error:
|
|
raise RuntimeError(
|
|
"huggingface-hub is required to resolve SAM-Audio support models"
|
|
) from error
|
|
|
|
options: dict[str, Any] = {
|
|
"repo_id": spec.repo_id,
|
|
"revision": spec.revision,
|
|
}
|
|
if spec.allow_patterns is not None:
|
|
options["allow_patterns"] = list(spec.allow_patterns)
|
|
|
|
try:
|
|
return huggingface_hub.snapshot_download(
|
|
**options,
|
|
local_files_only=True,
|
|
)
|
|
except Exception:
|
|
LOGGER.info(
|
|
"Downloading one-time SAM-Audio support model %s at %s",
|
|
spec.repo_id,
|
|
spec.revision,
|
|
)
|
|
return huggingface_hub.snapshot_download(**options)
|
|
|
|
|
|
def _sam_audio_model_kwargs(model_path: Path) -> dict[str, Any]:
|
|
model_kwargs: dict[str, Any] = {
|
|
"map_location": "cpu",
|
|
# Upstream's override only calls torch.nn.Module.load_state_dict when this
|
|
# is true. Passing false silently leaves the complete model uninitialized.
|
|
"strict": True,
|
|
"text_ranker": None,
|
|
"visual_ranker": None,
|
|
}
|
|
config_path = model_path / "config.json"
|
|
if not config_path.is_file():
|
|
return model_kwargs
|
|
|
|
config = json.loads(config_path.read_text(encoding="utf-8"))
|
|
text_encoder = config.get("text_encoder")
|
|
if isinstance(text_encoder, dict) and text_encoder.get("name") == "t5-base":
|
|
text_encoder = dict(text_encoder)
|
|
text_encoder["name"] = _cached_support_snapshot(SUPPORT_MODELS["t5-base"])
|
|
model_kwargs["text_encoder"] = text_encoder
|
|
|
|
# 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)
|
|
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,
|
|
model_management: Any,
|
|
model_patcher: Any,
|
|
attention_backend: str = "pytorch",
|
|
) -> SAMAudioPipeline:
|
|
load_device = model_management.get_torch_device()
|
|
validate_attention_backend(attention_backend, load_device)
|
|
model_path = resolve_model(model_name, folder_paths)
|
|
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(
|
|
f"Failed to load SAM-Audio from {model_path}. Upstream may also need its "
|
|
"T5, Perception Encoder, and ranker assets available through Hugging Face."
|
|
) from error
|
|
|
|
model = model.eval()
|
|
# BaseModel exposes device() as a method, while ComfyUI's ModelPatcher tracks
|
|
# residency through a mutable model.device attribute. Upstream does not call
|
|
# that helper internally, so adopting ComfyUI's convention is safe here.
|
|
model.device = torch.device("cpu")
|
|
offload_device = model_management.unet_offload_device()
|
|
patcher = model_patcher.ModelPatcher(
|
|
model, load_device=load_device, offload_device=offload_device
|
|
)
|
|
return SAMAudioPipeline(
|
|
model=model,
|
|
processor=processor,
|
|
patcher=patcher,
|
|
device=load_device,
|
|
source=model_name,
|
|
attention_backend=attention_backend,
|
|
span_predictor_source=_configured_span_predictor(model_path),
|
|
)
|
|
|
|
|
|
def ode_options(inference_steps: int) -> dict[str, Any]:
|
|
if inference_steps < 2 or inference_steps % 2:
|
|
raise ValueError("inference_steps must be an even integer of at least 2")
|
|
return {"method": "midpoint", "options": {"step_size": 2.0 / inference_steps}}
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def seeded_inference(seed: int, device: torch.device) -> Iterator[None]:
|
|
cuda_devices: list[int] = []
|
|
if device.type == "cuda" and torch.cuda.is_available():
|
|
cuda_devices = [device.index if device.index is not None else torch.cuda.current_device()]
|
|
with torch.random.fork_rng(devices=cuda_devices):
|
|
torch.manual_seed(seed)
|
|
if device.type == "cuda" and torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
yield
|
|
|
|
|
|
def validate_spans(spans: tuple[SpanPrompt, ...], duration: float) -> None:
|
|
if not spans:
|
|
raise ValueError("At least one span prompt is required")
|
|
if not any(span.token == "+" for span in spans):
|
|
raise ValueError("At least one positive span is required")
|
|
for span in spans:
|
|
if span.start < 0 or span.end <= span.start:
|
|
raise ValueError(
|
|
f"Invalid span {span.start:g}-{span.end:g}: end must be after start"
|
|
)
|
|
if span.end > duration + 1e-6:
|
|
raise ValueError(
|
|
f"Span ends at {span.end:g}s but the input audio is only {duration:.3f}s"
|
|
)
|