テキスト・視覚・時間範囲のプロンプトを使用して音声から特定の音を分離する「sam-audio」をuv + Windowsで動かす

初めに

SAM-Audio(Segment Anything Model for Audio)は、Meta(Facebook Research)が開発した音声分離のためのファウンデーションモデルです。テキスト、視覚、または時間範囲のプロンプトを使用して、複雑な音声ミックスから特定の音を分離できます。

github.com

以下に uvで環境構築をしたリポジトリを公開しています

github.com

開発環境

環境構築

以下に対応するためにpyproject.toml の修正を行います

  • PyTorchのCUDA版を取得
  • Git依存関係のuv形式指定
  • perception-models が Python 3.11以上を要求
  • torchcodecのWindows対応
[project]
name = "sam_audio"
version = "0.1.0"
description = "Segment Anything Audio"
authors = [
    { name="Andros Tjandra", email="androstj@meta.com" },
    { name="Ann Lee", email="annl@meta.com" },
    { name="Bowen Shi", email="bshi@meta.com" },
    { name="Julius Richter", email="jrichter@meta.com" },
    { name="Matt Le", email="mattle@meta.com" },
    { name="Yi-Chiao Wu", email="yichiaowu@meta.com" },
]

readme = "README.md"
license = { file="LICENSE" }
requires-python = ">=3.11"
dependencies = [
    "dacvae",
    "audiobox_aesthetics",
    "einops",
    "imagebind",
    "laion-clap",
    "numpy",
    "perception-models",
    "pydub",
    "torch>=2.5.0",
    "torchaudio>=2.5.0",
    "torchcodec",
    "torchdiffeq",
    "torchvision>=0.20.0",
    "transformers>=4.54.0",
]

[tool.setuptools.packages.find]
include = ["sam_audio*"]

[tool.ruff]

target-version = "py311"

lint.select=[
    "B",
    "C",
    "E",
    "W",
    "F",
    "I",
]
lint.ignore = [
    "E501",
    "E731",
    "C901",
    "B006",
]

[project.urls]
Homepage = "https://github.com/facebookresearch/sam-audio"
Repository = "https://github.com/facebookresearch/sam-audio"

[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"

# ========== uv設定(CUDA 12.6対応) ==========

[[tool.uv.index]]
name = "pytorch-cu126"
url = "https://download.pytorch.org/whl/cu126"
explicit = true

[tool.uv.sources]
# Git依存関係
dacvae = { git = "https://github.com/facebookresearch/dacvae.git" }
imagebind = { git = "https://github.com/facebookresearch/ImageBind.git" }
laion-clap = { git = "https://github.com/LAION-AI/CLAP.git" }
perception-models = { git = "https://github.com/facebookresearch/perception_models.git", branch = "unpin-deps" }

# PyTorchパッケージ(CUDA 12.6インデックスから取得)
torch = { index = "pytorch-cu126" }
torchvision = { index = "pytorch-cu126" }
torchaudio = { index = "pytorch-cu126" }
# torchcodecはWindows用CUDA wheelがないためPyPIから取得

修正後のライブラリをインストールするために以下を実行します

uv sync

torchcodecが動作するためにFFmpeg 6のDLL(共有ライブラリ)が必要です。そのため FFmpeg Buildsから「shared」ビルドをダウンロードし、DLLをPATHの通ったディレクトリに配置。

実行

推論をするための以下のようなサンプルスクリプトを作成します

"""SAM-Audio サンプル実行スクリプト"""

import os
import torch
import torchaudio
import gc
from sam_audio import SAMAudio, SAMAudioProcessor

# メモリ最適化設定
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

def main():
    print("=== SAM-Audio Sample ===")

    # デバイス設定
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # GPUメモリをクリア
    if device.type == "cuda":
        torch.cuda.empty_cache()
        gc.collect()
        print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")

    # モデルとプロセッサをロード(baseモデルを使用してメモリ節約)
    model_name = "facebook/sam-audio-base"
    print(f"Loading model: {model_name}")
    model = SAMAudio.from_pretrained(model_name).to(device).eval()
    processor = SAMAudioProcessor.from_pretrained(model_name)
    print("Model loaded!")

    # 入力ファイル
    video_file = "examples/assets/office.mp4"
    description = "A man speaking"

    print(f"Input: {video_file}")
    print(f"Description: {description}")

    # 処理
    print("Processing...")
    inputs = processor(audios=[video_file], descriptions=[description]).to(device)

    with torch.inference_mode():
        result = model.separate(inputs)

    # 保存
    sample_rate = processor.audio_sampling_rate
    torchaudio.save("target.wav", result.target[0].cpu(), sample_rate)
    torchaudio.save("residual.wav", result.residual[0].cpu(), sample_rate)

    print("=== Done! ===")
    print(f"Saved: target.wav (分離された音声)")
    print(f"Saved: residual.wav (残りの音声)")

if __name__ == "__main__":
    main()

このときに VRAM 16GBだとlargeモデルは動かなかったので、baseモデルで推論を行っています。入力時に以下のようなプロンプトを指定することで、指定した男性のみの音声を取得することができます

    description = "A man speaking"

時間指定する場合は、Inputの部分に anchors=[[["+", 6.3, 7.0]]] を追加して実行します