stem-splitter 0.0.1__py3-none-any.whl

stem_splitter/band_split_roformer.py

@@ -0,0 +1,418 @@
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+import librosa
+import numpy as np
+import einops
+from .transformer import Transformer, RMSNorm
+
+from typing import List, Tuple
+
+
+def checkpoint_bypass(f, *args):
+    return f(*args)
+
+
+def build_mel_band_indices(
+    sample_rate: int,
+    n_fft: int,
+    n_bands: int,
+) -> list[np.ndarray[int]]:
+    """
+    各 Mel バンドがカバーする「周波数ビン番号だけ」を返す。
+    返り値: list[np.ndarray]  (len = n_bands)
+    """
+    # --- Mel フィルタ → 0/1 マスク --------------------------------
+    mel_fb = librosa.filters.mel(sr=sample_rate, n_fft=n_fft, n_mels=n_bands)
+    mel_fb[0, 0] = 1.0
+    mel_fb[-1, -1] = 1.0
+    mask = mel_fb > 0  # (B, F) bool
+
+    return [np.where(mask[b])[0].astype(np.int32) for b in range(n_bands)]
+
+
+def build_band_indices(sample_rate: int = 44_100, n_fft: int = 2_048) -> Tuple[np.ndarray, List[np.ndarray]]:
+    """
+    BSRoformer 用バンド分割 (各バンドの bin 数が必ず閾値以上)。
+
+    Returns
+    -------
+    freqs : np.ndarray
+        rfft の周波数 (長さ n_fft//2 + 1) [Hz]
+    bands : List[np.ndarray]
+        各要素が「そのバンドに含まれるビン番号」の 1‑D 配列
+    """
+    freqs = np.fft.rfftfreq(n_fft, d=1 / sample_rate)
+
+    # バンドが始まるビンの周波数から必要 bin 数を決定
+    def bins_per_band(freq_hz: float) -> int:
+        if freq_hz < 1_000:  # 0‑1 kHz
+            return 2
+        elif freq_hz < 2_000:  # 1‑2 kHz
+            return 4
+        elif freq_hz < 4_000:  # 2‑4 kHz
+            return 12
+        elif freq_hz < 8_000:  # 4‑8 kHz
+            return 24
+        elif freq_hz < 16_000:  # 8‑16 kHz
+            return 48
+        else:  # ≥16 kHz → 後で 2 分割
+            return -1
+
+    freq_indices: List[np.ndarray] = []
+    idx = 0
+    n_bins = len(freqs)
+
+    while idx < n_bins:
+        # 16 kHz 以上は残りを 2 分割して終了
+        if freqs[idx] >= 16_000:
+            remaining = np.arange(idx, n_bins)
+            half = len(remaining) // 2
+            freq_indices.append(remaining[:half])
+            freq_indices.append(remaining[half:])
+            break
+
+        size = bins_per_band(freqs[idx])
+        band = np.arange(idx, min(idx + size, n_bins))
+        freq_indices.append(band)
+        idx += size  # bin数 単位で次へ
+
+    return freqs, freq_indices
+
+
+class BandSplit(torch.nn.Module):
+    def __init__(self, hidden_size: int, band_indices, num_channels: int, extra_windows: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.band_indices = band_indices
+
+        for i, idx in enumerate(band_indices):
+            self.register_buffer(f"band_idx_{i}", torch.tensor(idx, dtype=torch.long), persistent=False)
+
+        self.num_bands = len(band_indices)
+        self.to_features = torch.nn.ModuleList([])
+        for i in range(self.num_bands):
+            sub_band_freqs = len(band_indices[i]) * num_channels * 2 * (extra_windows + 1)
+            self.to_features.append(
+                torch.nn.Sequential(
+                    RMSNorm(sub_band_freqs),
+                    torch.nn.Linear(sub_band_freqs, hidden_size),
+                )
+            )
+
+    def forward(self, x):
+        # x: (B, T, C, F)
+        sub_band_list = []
+        for i, proj_layer in enumerate(self.to_features):
+            band_indices = getattr(self, f"band_idx_{i}")
+            # サブバンドインデックスで周波数軸から抜きだす
+            sub_band = x[..., band_indices]  # (B, T, C, sub_band_freqs)
+            sub_band = einops.rearrange(sub_band, "b t c f -> b t (f c)")
+            sub_band = proj_layer(sub_band)  # (B, T, hidden_size)
+            sub_band_list.append(sub_band)
+
+        return torch.stack(sub_band_list, dim=-2)  # (B, T, num_bands, hidden_size)
+
+
+def band_mask_mlp(input_dim, band_dim, hidden_dim, depth):
+    layers = []
+    # 入力→隠れ→…→隠れ→(band_dim*2) の構造
+    dims = [input_dim] + [hidden_dim] * depth + [band_dim * 2]
+    for i in range(len(dims) - 1):
+        layers.append(nn.Linear(dims[i], dims[i + 1]))
+        # 最終層の後に活性化は入れない
+        if i < len(dims) - 2:
+            layers.append(nn.Tanh())
+    return nn.Sequential(*layers)
+
+
+class MaskEstimator(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        band_indices,
+        num_channels: int,
+        mlp_expansion_factor: int = 4,
+        depth: int = 2,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.num_channels = num_channels
+        self.band_indices = band_indices
+
+        self.num_bands = len(band_indices)
+        self.to_freqs = torch.nn.ModuleList([])
+        for i in range(self.num_bands):
+            sub_band_freqs = len(band_indices[i]) * num_channels * 2
+            self.to_freqs.append(
+                nn.Sequential(
+                    band_mask_mlp(
+                        input_dim=input_dim,
+                        band_dim=sub_band_freqs,
+                        hidden_dim=input_dim * mlp_expansion_factor,
+                        depth=depth,
+                    ),
+                    nn.GLU(dim=-1),
+                )
+            )
+
+    def forward(self, x):
+        # x: (B, T, K, D)
+        out = []
+        sub_band_list = torch.unbind(x, dim=2)  # list([B, T, D])
+        assert len(sub_band_list) == len(self.to_freqs)
+
+        for i, proj_layer in enumerate(self.to_freqs):
+            sub_band = proj_layer(sub_band_list[i])  # (B, T, sub_band_freqs)
+            sub_band = einops.rearrange(
+                sub_band,
+                "b t (f c) -> b t c f",
+                c=self.num_channels * 2,
+                f=len(self.band_indices[i]),
+            )
+            out.append(sub_band)
+
+        return torch.concat(out, dim=-1)  # (B, T, C, F)
+
+
+class BSRoformer(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_layers: int = 12,
+        sample_rate: int = 44100,
+        num_channels: int = 2,
+        head_dim: int = 64,
+        num_heads: int = 8,
+        n_fft: int = 2048,
+        hop_length: int = 512,
+        n_bands: int = 60,
+        num_stems: int = 2,
+        band_split_type: str = "bs",
+        ffn_hidden_size_factor: int = 4,
+        dropout: float = 0.0,
+        use_shared_bias: bool = True,
+        use_gradient_checkpoint=True,
+    ):
+        super().__init__()
+        self.n_fft = n_fft
+        self.window_size = n_fft
+        self.hop_length = hop_length
+        self.num_channels = num_channels
+        self.num_stems = num_stems
+        self.band_split_type = band_split_type
+        self.use_gradient_checkpoint = use_gradient_checkpoint
+
+        if band_split_type == "bs":
+            _, self.band_indices = build_band_indices(sample_rate=sample_rate, n_fft=n_fft)
+        elif band_split_type == "mel":
+            self.band_indices = build_mel_band_indices(sampling_rate=sample_rate, n_fft=n_fft, n_bands=n_bands)
+        else:
+            raise NotImplementedError("サポートしていない band_split_type です")
+
+        self.register_buffer("hann_window", torch.hann_window(self.window_size), persistent=False)
+        self.num_bands = len(self.band_indices)
+        self.band_split = BandSplit(
+            hidden_size=dim,
+            band_indices=self.band_indices,
+            num_channels=num_channels,
+            extra_windows=0,
+        )
+
+        self.layers = nn.ModuleList([])
+        if use_shared_bias:
+            hidden_size = head_dim * num_heads
+            self.shared_qkv_bias = nn.Parameter(torch.zeros(hidden_size * 3))  # QKV
+            self.shared_out_bias = nn.Parameter(torch.zeros(dim))  # OUT
+
+        for _ in range(num_layers):
+            time_roformer = Transformer(
+                input_dim=dim,
+                head_dim=head_dim,
+                num_layers=1,
+                num_heads=num_heads,
+                ffn_hidden_size_factor=ffn_hidden_size_factor,
+                dropout=dropout,
+                shared_qkv_bias=self.shared_qkv_bias,
+                shared_out_bias=self.shared_out_bias,
+            )
+            band_roformer = Transformer(
+                input_dim=dim,
+                head_dim=head_dim,
+                num_layers=1,
+                num_heads=num_heads,
+                ffn_hidden_size_factor=ffn_hidden_size_factor,
+                dropout=dropout,
+                shared_qkv_bias=self.shared_qkv_bias,
+                shared_out_bias=self.shared_out_bias,
+            )
+            self.layers.append(nn.ModuleList([time_roformer, band_roformer]))
+
+        self.final_norm = RMSNorm(dim)
+
+        self.register_buffer(
+            "freq_indices",
+            torch.tensor(np.concatenate(self.band_indices), dtype=torch.long),
+            persistent=False,
+        )
+        # そのビンが何本のバンドに含まれるかを数える
+        counts = np.zeros(n_fft // 2 + 1, dtype=np.float32)
+        for idx in self.band_indices:
+            counts[idx] += 1
+        self.register_buffer(
+            "num_bands_per_freq",
+            torch.tensor(counts, dtype=torch.float32),
+            persistent=False,
+        )  # shape (F_total,)
+        self.mask_estimators = nn.ModuleList(
+            [
+                MaskEstimator(
+                    input_dim=dim,
+                    band_indices=self.band_indices,
+                    num_channels=num_channels,
+                    depth=1,
+                    mlp_expansion_factor=4,
+                )
+                for _ in range(num_stems)
+            ]
+        )
+
+    def to_spectrogram(self, inputs: torch.Tensor) -> torch.Tensor:
+        # (B, C, T)
+        inputs = einops.rearrange(inputs, "b c t -> (b c) t")
+
+        windows: list[torch.Tensor] = [self.hann_window]
+
+        spectrograms = []
+        for win in windows:
+            spectrogram = torch.stft(
+                input=inputs,
+                n_fft=self.n_fft,
+                hop_length=self.hop_length,
+                win_length=self.window_size,
+                window=win.to(inputs.device),
+                center=True,
+                return_complex=True,
+            )
+            spectrogram = einops.rearrange(spectrogram, "(b c) f t -> b c f t", c=self.num_channels)
+            spectrograms.append(spectrogram)  # list of [B, C, F, T]
+
+        spectrogram = torch.concat(spectrograms, dim=1)
+        spectrogram = torch.view_as_real(spectrogram)
+        spectrogram = einops.rearrange(spectrogram, "b c f t s -> b (c s) t f", s=2)
+        return spectrogram.to(inputs.dtype), spectrograms[0]
+
+    def mixture_consistency_projection(
+        self,
+        source_estimates: torch.Tensor,  # [B, C, N, F, T] complex
+        mixture: torch.Tensor,  # [B, C, F, T] complex
+        weights: torch.Tensor | None = None,  # [B, C, N, F, T] real, 任意
+    ) -> torch.Tensor:
+        residual = mixture[:, :, None] - source_estimates.sum(dim=2, keepdim=True)
+        if weights is None:
+            correction = residual / source_estimates.shape[2]
+        else:
+            weights = weights / (weights.sum(dim=2, keepdim=True) + 1e-8)
+            correction = weights * residual
+        return source_estimates + correction
+
+    def to_recon_audio(
+        self,
+        mask: torch.Tensor,
+        original_spec_complex: torch.Tensor,
+        original_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        # mask: [B, C, N, F, T]
+        source_estimates = original_spec_complex[:, :, None] * mask
+
+        # mixture consistency
+        magnitude = source_estimates.abs().clamp_min(1e-8)
+        weights = magnitude**2
+        separated_spectrogram = self.mixture_consistency_projection(
+            source_estimates=source_estimates,
+            mixture=original_spec_complex,
+            weights=weights,
+        )
+
+        separated_spectrogram = einops.rearrange(separated_spectrogram, "b c n f t -> (b c n) f t")
+        recon_audio = torch.istft(
+            separated_spectrogram,
+            n_fft=self.window_size,
+            hop_length=self.hop_length,
+            win_length=self.window_size,
+            window=self.hann_window.to(device),
+            center=True,
+            return_complex=False,
+            length=original_length,
+        )  # [B*C, T]
+        recon_audio = einops.rearrange(recon_audio, "(b c n) t -> b c n t", c=self.num_channels, n=self.num_stems)
+        recon_audio = einops.rearrange(recon_audio, "b c n t -> b n c t")
+        return recon_audio
+
+    def forward(self, x):
+        # x: (B, C, T)
+        if self.use_gradient_checkpoint or self.training:
+            checkpoint = torch.utils.checkpoint.checkpoint
+        else:
+            checkpoint = checkpoint_bypass
+
+        istft_length = x.shape[-1]
+        x, original_spec_complex = self.to_spectrogram(x)  # x: [B, C*S, T, F]
+        x = einops.rearrange(x, "b c t f -> b t c f")
+
+        # 勾配が流れない問題に対処
+        x = checkpoint(self.band_split, x, use_reentrant=False)  # (B, T, K, hidden_size)
+
+        for time_roformer, band_roformer in self.layers:
+            B, T, K, F = x.shape
+            # 時間軸Transformer
+            x = einops.rearrange(x, "b t k f -> (b k) t f")  # [B*K, T, F]
+            x = checkpoint(time_roformer, x, use_reentrant=False)
+            x = einops.rearrange(x, "(b k) t f -> b t k f", k=K)  # [B, T, K, F]
+
+            # バンド軸Transformer
+            x = x.reshape(B * T, K, F)  # [B*T, K, F]
+            x = checkpoint(band_roformer, x, use_reentrant=False)
+            x = x.reshape(B, T, K, F)
+
+        x = self.final_norm(x)
+
+        # 音源分離マスク
+        mask_list = []
+        for mask_estimator in self.mask_estimators:
+            pred_mask = checkpoint(mask_estimator, x, use_reentrant=False)  # [B, T, C*S, F]
+            pred_mask = einops.rearrange(pred_mask, "b t (c s) f -> b t c f s", s=2).contiguous()
+            pred_mask = torch.view_as_complex(pred_mask)  # [B, T, C, F]
+            mask_list.append(pred_mask)
+
+        mask_all = torch.stack(mask_list, dim=-2)  # [B, T, C, N, F]
+        if self.band_split_type == "mel":
+            # 周波数ビンで重なる部分を加算する
+            B, T, C, N, F_concat = mask_all.shape
+            F_total = self.n_fft // 2 + 1
+            mask_sum = torch.zeros((B, T, C, N, F_total), dtype=mask_all.dtype, device=mask_all.device)
+            freq_idx = self.freq_indices.view(1, 1, 1, 1, -1).expand(B, T, C, N, -1)  # [B, T, C, N, F_concat]
+            mask_sum.scatter_add_(
+                dim=-1,
+                index=freq_idx,
+                src=mask_all,
+            )
+
+            # 重なった本数で割って平均
+            denom = self.num_bands_per_freq.clamp(min=1e-8)  # [F_total]
+            denom = denom.view(1, 1, 1, 1, -1)  # broadcast
+            mask_avg = mask_sum / denom  # [B, T, C, N, F_total]
+            mask_avg = einops.rearrange(mask_avg, "b t c n f -> b c n f t")
+        else:
+            mask_avg = einops.rearrange(mask_all, "b t c n f -> b c n f t")
+
+        recon_audio = self.to_recon_audio(
+            mask=mask_avg,
+            original_spec_complex=original_spec_complex,
+            original_length=istft_length,
+            dtype=x.dtype,
+            device=x.device,
+        )
+        return recon_audio

stem_splitter/cli.py

@@ -0,0 +1,4 @@
+from .inference import main
+
+if __name__ == "__main__":
+    main()

stem_splitter/inference.py

@@ -0,0 +1,470 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import sys
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Optional, Tuple, Union
+
+import librosa
+import numpy as np
+import math
+import torch
+import torchaudio
+import warnings
+from tqdm import tqdm
+
+from .band_split_roformer import BSRoformer
+
+
+@dataclass
+class SeparationConfig:
+    # 推論関連
+    target_sample_rate: int = 44100
+    device_preference: Optional[str] = None  # "cuda" / "cpu" / "mps"
+    use_half_precision: bool = False
+
+    chunk_size: int = 588_800  # 約 13.35 秒 @ 44.1kHz
+    hop_size: Optional[int] = None  # 既定: chunk_size // 2（50% overlap）
+    window_type: str = "hann"  # "hann" 推奨
+
+    stem_names: Tuple[str, ...] = (
+        "bass",
+        "drums",
+        "other",
+        "vocals",
+        "guitar",
+        "piano",
+    )
+
+    model_name: str = "bs_roformer"
+    hf_repo_id: Optional[str] = None
+    hf_filename: Optional[str] = None
+    hf_revision: str = "main"
+    expected_sha256: Optional[str] = None  # 任意：完全性検証に使用
+
+    # キャッシュ先（未指定ならユーザホーム配下）
+    cache_dir: Optional[Path] = None
+
+
+# モデル名→Hugging Face 上の既定情報（必要に応じて上書き）
+MODEL_REGISTRY: Dict[str, Dict[str, str]] = {
+    "bs_roformer": {
+        "repo_id": "anime-song/stem-splitter",
+        "filename": "stem_splitter.pt",
+        "revision": "main",
+    }
+}
+
+
+def _save_wav_no_warning(out_path: Path, samples: torch.Tensor, sample_rate: int) -> None:
+    """
+    警告を出さずに WAV 保存する。
+    優先度: TorchCodec > torchaudio.save_with_torchcodec > torchaudio.save(最後の手段)
+    入力: samples (C, T), float32 [-1, 1] 推奨
+    """
+    if samples.device.type != "cpu":
+        samples = samples.cpu()
+    samples = samples.contiguous().to(torch.float32)
+
+    # 可能なら TorchCodec を直接使用（推奨）
+    try:
+        from torchcodec.encoders import AudioEncoder  # pip install torchcodec
+
+        AudioEncoder(samples, sample_rate=sample_rate).to_file(str(out_path))
+        return
+    except Exception:
+        pass
+
+    # torchaudio の TorchCodec 経由 API（存在する場合）
+    if hasattr(torchaudio, "save_with_torchcodec"):
+        torchaudio.save_with_torchcodec(str(out_path), samples, sample_rate, channels_first=True)
+        return
+
+    warnings.filterwarnings(
+        "ignore",
+        message="In 2.9, this function's implementation will be changed to use torchaudio.save_with_torchcodec",
+        category=UserWarning,
+        module=r"torchaudio\._backend\.utils",
+    )
+    warnings.filterwarnings(
+        "ignore",
+        message="StreamWriter has been deprecated",
+        category=UserWarning,
+        module=r"torchaudio\._backend\.ffmpeg",
+    )
+    torchaudio.save(str(out_path), samples, sample_rate)
+
+
+def resolve_device(device_preference: Optional[str]) -> torch.device:
+    if device_preference is not None:
+        return torch.device(device_preference)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def default_cache_dir() -> Path:
+    return Path(os.environ.get("STEM_SPLITTER_CACHE", Path.home() / ".cache" / "stem_splitter" / "weights"))
+
+
+def get_model_info(config: SeparationConfig) -> Tuple[str, str, str]:
+    """
+    config から Hugging Face の repo / filename / revision を解決。
+    未指定なら MODEL_REGISTRY を参照。
+    """
+    repo_id = config.hf_repo_id
+    filename = config.hf_filename
+    revision = config.hf_revision
+
+    if (repo_id is None or filename is None) and config.model_name in MODEL_REGISTRY:
+        fallback = MODEL_REGISTRY[config.model_name]
+        repo_id = repo_id or fallback["repo_id"]
+        filename = filename or fallback["filename"]
+        revision = revision or fallback.get("revision", "main")
+
+    if repo_id is None or filename is None:
+        raise ValueError(
+            "Hugging Face の重みファイル情報が不足しています。"
+            " SeparationConfig(hf_repo_id=..., hf_filename=...) を指定するか、MODEL_REGISTRY を更新してください。"
+        )
+    return repo_id, filename, revision
+
+
+def build_weight_path(config: SeparationConfig, filename: str) -> Path:
+    cache_root = config.cache_dir or default_cache_dir()
+    cache_root.mkdir(parents=True, exist_ok=True)
+    return cache_root / filename
+
+
+def sha256sum(file_path: Path) -> str:
+    h = hashlib.sha256()
+    with file_path.open("rb") as f:
+        for chunk in iter(lambda: f.read(8192), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def download_from_huggingface(
+    repo_id: str,
+    filename: str,
+    revision: str,
+    destination: Path,
+) -> None:
+    base_url = f"https://huggingface.co/{repo_id}/resolve/{revision}/{filename}?download=true"
+
+    with urllib.request.urlopen(base_url) as response:
+        total_size = int(response.headers.get("Content-Length", "0"))
+        block_size = 1024 * 1024  # 1MB
+        tmp_path = destination.with_suffix(".tmp")
+
+        with (
+            tmp_path.open("wb") as out_file,
+            tqdm(
+                total=total_size if total_size > 0 else None,
+                unit="B",
+                unit_scale=True,
+                desc=f"Downloading {filename}",
+            ) as progress_bar,
+        ):
+            while True:
+                data = response.read(block_size)
+                if not data:
+                    break
+                out_file.write(data)
+                progress_bar.update(len(data))
+
+        tmp_path.replace(destination)
+
+
+def ensure_weight_file(config: SeparationConfig) -> Path:
+    repo_id, filename, revision = get_model_info(config)
+    weight_path = build_weight_path(config, filename)
+
+    if weight_path.exists():
+        if config.expected_sha256:
+            actual = sha256sum(weight_path)
+            if actual.lower() != config.expected_sha256.lower():
+                print(
+                    f"[WARN] 既存の重みファイルの SHA256 が一致しません。再取得します。\n"
+                    f"  expected={config.expected_sha256}\n  actual  ={actual}"
+                )
+                weight_path.unlink(missing_ok=True)
+        else:
+            # 検証なしでそのまま使用
+            return weight_path
+
+    # ダウンロード
+    print(f"[INFO] 重みファイルが見つかりません。Hugging Face から取得します: {repo_id}/{filename}@{revision}")
+    download_from_huggingface(repo_id, filename, revision, weight_path)
+
+    if config.expected_sha256:
+        actual = sha256sum(weight_path)
+        if actual.lower() != config.expected_sha256.lower():
+            weight_path.unlink(missing_ok=True)
+            raise RuntimeError("ダウンロードした重みファイルの SHA256 が一致しません。")
+
+    return weight_path
+
+
+def load_state_dict_flex(path: Path) -> Dict[str, torch.Tensor]:
+    """
+    checkpoint 形式の違いに寛容なロード（state_dict 直/ラップの両対応）。
+    """
+    try:
+        checkpoint = torch.load(str(path), map_location="cpu", weights_only=False)
+    except TypeError:
+        # PyTorch の引数差異対策（古い/新しい両対応）
+        checkpoint = torch.load(str(path), map_location="cpu")
+
+    if isinstance(checkpoint, dict):
+        # 典型例: {"state_dict": ..., ...} / {"model": ..., ...} / そのまま state_dict
+        for key in ("state_dict", "model", "ema_state_dict"):
+            if key in checkpoint and isinstance(checkpoint[key], dict):
+                return checkpoint[key]
+        return checkpoint  # そのまま state_dict とみなす
+    raise RuntimeError("未知のチェックポイント形式です。dict ではありません。")
+
+
+def load_mss_model(config: SeparationConfig, device: torch.device) -> torch.nn.Module:
+    """
+    BSRoformer を作成し、重みを読み込んで返します。
+    """
+    weight_path = ensure_weight_file(config)
+
+    model = BSRoformer(
+        dim=256,
+        num_layers=12,
+        sample_rate=config.target_sample_rate,
+        num_channels=2,
+        head_dim=64,
+        num_heads=8,
+        n_fft=2048,
+        hop_length=512,
+        num_stems=len(config.stem_names),
+    )
+
+    # 重みを読み込み
+    state_dict = load_state_dict_flex(weight_path)
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if missing:
+        print(f"[WARN] 欠落している重みキー: {sorted(missing)[:5]}{' ...' if len(missing) > 5 else ''}")
+    if unexpected:
+        print(f"[WARN] 予期しない重みキー: {sorted(unexpected)[:5]}{' ...' if len(unexpected) > 5 else ''}")
+
+    model.to(device)
+    if config.use_half_precision and device.type == "cuda":
+        model.half()
+    model.eval()
+    return model
+
+
+def _separate_one_file(
+    input_wav_path: Path,
+    output_directory: Path,
+    config: SeparationConfig,
+    model: torch.nn.Module,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> Dict[str, Path]:
+    """
+    単一の WAV を処理し、ステムごとの出力パスを返す。
+    入力: (C, T) / モデル入出力: (B, C, T) -> (B, N, C, T)
+    OLA: chunk_size, hop_size, Hann 窓（sum-of-windows 正規化）
+    """
+    # 読み込み -> (C, T) に正規化
+    y, _ = librosa.load(str(input_wav_path), sr=config.target_sample_rate, mono=False)
+    if y.ndim == 1:
+        y = y[None, :]  # (1, T)
+    elif y.ndim == 2 and y.shape[0] > y.shape[1]:
+        # librosa は (T, C) を返すことがあるので (C, T) へ
+        y = y.T
+    y = y.astype(np.float32, copy=False)  # (C, T)
+
+    channels, total_length = y.shape
+
+    # チャンク条件
+    chunk_size = int(config.chunk_size)
+    hop_size = int(config.hop_size) if config.hop_size is not None else chunk_size // 2
+    if hop_size <= 0 or hop_size > chunk_size:
+        raise ValueError("hop_size は 1..chunk_size の範囲で指定してください")
+
+    # padding 後の長さ（最後が半端でも必ず 1 チャンクぶん確保）
+    if total_length <= chunk_size:
+        padded_length = chunk_size
+    else:
+        steps = math.ceil((total_length - chunk_size) / hop_size)
+        padded_length = steps * hop_size + chunk_size
+
+    if padded_length > total_length:
+        pad_amount = padded_length - total_length
+        y = np.pad(y, ((0, 0), (0, pad_amount)), mode="constant")
+
+    # 窓・蓄積（sum-of-windows で正規化）
+    if config.window_type.lower() != "hann":
+        raise ValueError(f"未対応の window_type: {config.window_type}")
+    base_window = torch.hann_window(chunk_size, periodic=False, dtype=dtype, device=device)
+
+    num_stems = len(config.stem_names)
+    accum = np.zeros((num_stems, channels, padded_length), dtype=np.float32)
+    weight_sum = np.zeros(padded_length, dtype=np.float32)  # ← 窓の和を積む
+
+    # チャンク推論ループ
+    for start in range(0, padded_length - chunk_size + 1, hop_size):
+        end = start + chunk_size
+
+        # 末尾の短い区間はここでゼロパディング（上の pad と二重になっても影響なし）
+        input_chunk_np = y[:, start:end]
+        if input_chunk_np.shape[1] < chunk_size:
+            pad = chunk_size - input_chunk_np.shape[1]
+            input_chunk_np = np.pad(input_chunk_np, ((0, 0), (0, pad)), mode="constant")
+
+        # (1, C, T)
+        input_chunk = torch.from_numpy(input_chunk_np).to(device=device, dtype=dtype).unsqueeze(0)
+
+        with torch.no_grad():
+            output_chunk = model(input_chunk)  # (1, N, C, T_out)
+
+        if not isinstance(output_chunk, torch.Tensor) or output_chunk.ndim != 4:
+            raise RuntimeError("モデル出力は (B, N, C, T) の Tensor を想定しています。")
+
+        _, _, _, t_out = output_chunk.shape
+        window = (
+            base_window if t_out == chunk_size else torch.hann_window(t_out, periodic=False, dtype=dtype, device=device)
+        )
+
+        # 出力に 1 回だけ窓掛け → 合成は窓の「和」で割る
+        windowed = (output_chunk * window.view(1, 1, 1, -1)).squeeze(0)  # (N, C, T)
+        out_np = windowed.to(torch.float32).cpu().numpy()
+
+        accum[:, :, start : start + t_out] += out_np
+        weight_sum[start : start + t_out] += window.to(torch.float32).cpu().numpy()
+
+        del input_chunk, output_chunk, windowed  # メモリ節約
+
+    # sum-of-windows で正規化（端も自然に補正される）
+    eps = 1e-8
+    weight_sum = np.maximum(weight_sum, eps)
+    accum /= weight_sum[None, None, :]
+
+    # 元の長さにトリムして保存
+    trim_length = total_length
+    base_name = input_wav_path.stem
+    saved_paths: Dict[str, Path] = {}
+
+    for i, stem_name in enumerate(config.stem_names):
+        stem_array = accum[i, :, :trim_length]  # (C, T)
+        stem_tensor = torch.from_numpy(stem_array)  # float32
+        out_path = output_directory / base_name / f"{base_name}_{stem_name}.wav"
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        _save_wav_no_warning(out_path, stem_tensor, config.target_sample_rate)
+        saved_paths[stem_name] = out_path
+
+    return saved_paths
+
+
+def separate_stems(
+    input_audio_path: Union[str, Path],
+    output_directory: Union[str, Path],
+    config: Optional[SeparationConfig] = None,
+) -> Union[Dict[str, Path], Dict[Path, Dict[str, Path]]]:
+    if config is None:
+        config = SeparationConfig()
+
+    input_path = Path(input_audio_path)
+    output_dir = Path(output_directory)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # デバイス・dtype・モデルはここで1回だけ
+    device = resolve_device(config.device_preference)
+    dtype = torch.float16 if (config.use_half_precision and device.type == "cuda") else torch.float32
+    model = load_mss_model(config, device=device)
+
+    if input_path.is_file():
+        if input_path.suffix.lower() != ".wav":
+            raise ValueError("WAV 以外は許可していません。入力は .wav にしてください。")
+        return _separate_one_file(input_path, output_dir, config, model, device, dtype)
+
+    if input_path.is_dir():
+        wav_files = sorted([p for p in input_path.rglob("*") if p.is_file() and p.suffix.lower() == ".wav"])
+        if not wav_files:
+            raise FileNotFoundError("指定ディレクトリに .wav ファイルが見つかりません。")
+
+        results: Dict[Path, Dict[str, Path]] = {}
+        for wav_path in wav_files:
+            current_out_dir = output_dir  # フラットに保存する場合はこちら
+            result = _separate_one_file(wav_path, current_out_dir, config, model, device, dtype)
+            results[wav_path] = result
+        return results
+
+    raise FileNotFoundError(f"入力パスが存在しません: {input_path}")
+
+
+def _build_arg_parser():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Stem Splitter Inference (BSRoformer, OLA)")
+    parser.add_argument("input_audio_path", type=Path, help="入力音声ファイルのパス")
+    parser.add_argument("--out-dir", type=Path, required=True, help="出力ディレクトリ")
+    parser.add_argument("--device", default=None, help="cuda / cpu / mps（未指定で自動判定）")
+    parser.add_argument("--sr", type=int, default=44100, help="サンプルレート")
+    parser.add_argument("--no-half", action="store_true", help="半精度を無効化（既定はCUDA上で有効）")
+
+    parser.add_argument("--chunk-size", type=int, default=588_800, help="チャンク長（サンプル数）")
+    parser.add_argument("--hop-size", type=int, default=None, help="ホップ長（未指定で chunk_size//2）")
+    parser.add_argument("--window", type=str, default="hann", help="ウィンドウ種別（hann のみ対応）")
+
+    # 重み指定の上書き
+    parser.add_argument("--model-name", default="bs_roformer", help="MODEL_REGISTRY のキー")
+    parser.add_argument("--hf-repo", default=None, help="Hugging Face repo_id")
+    parser.add_argument("--hf-file", default=None, help="Hugging Face ファイル名")
+    parser.add_argument("--hf-rev", default="main", help="Hugging Face リビジョン")
+    parser.add_argument("--weights-cache", type=Path, default=None, help="重みキャッシュディレクトリ")
+    parser.add_argument("--sha256", default=None, help="重みファイル SHA256（任意）")
+    return parser
+
+
+def main():
+    parser = _build_arg_parser()
+    args = parser.parse_args()
+
+    config = SeparationConfig(
+        target_sample_rate=args.sr,
+        device_preference=args.device,
+        use_half_precision=not args.no_half,
+        chunk_size=args.chunk_size,
+        hop_size=args.hop_size,
+        window_type=args.window,
+        model_name=args.model_name,
+        hf_repo_id=args.hf_repo,
+        hf_filename=args.hf_file,
+        hf_revision=args.hf_rev,
+        expected_sha256=args.sha256,
+        cache_dir=args.weights_cache,
+    )
+
+    try:
+        result = separate_stems(args.input_audio_path, args.out_dir, config)
+    except Exception as exc:
+        print(f"[ERROR] 推論に失敗しました: {exc}", file=sys.stderr)
+        sys.exit(1)
+
+    # 表示
+    if isinstance(result, dict) and result and isinstance(next(iter(result.values())), Path):
+        # 単一ファイルケース: {stem: path}
+        for stem_name, out_path in result.items():
+            print(f"{stem_name}: {out_path}")
+    else:
+        # ディレクトリケース: {input_wav: {stem: path}}
+        for input_wav, stems_dict in result.items():
+            print(f"[{input_wav}]")
+            for stem_name, out_path in stems_dict.items():
+                print(f"  {stem_name}: {out_path}")
+
+
+if __name__ == "__main__":
+    main()

stem_splitter/transformer.py

@@ -0,0 +1,257 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+import einops
+from typing import Tuple
+
+
+def choose_low_precision_dtype() -> torch.dtype:
+    """
+    GPU の機能を調べて BF16 → FP16 → FP32 の順に
+    最も高速な演算 dtype を返す。
+    """
+    if not torch.cuda.is_available():
+        return torch.float32  # CPU 実行なら FP32 一択
+
+    # Ampere (sm80) 以降ならほぼ BF16 演算に対応
+    if torch.cuda.is_bf16_supported():  # PyTorch 2.1+
+        return torch.bfloat16
+
+    major_cc, _ = torch.cuda.get_device_capability()
+    # Pascal (sm60) 以降なら FP16 演算ユニットあり
+    if major_cc >= 6:
+        return torch.float16
+
+    return torch.float32  # それ以前の Maxwell など
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps: float = 5.960464477539063e-08):  # 0x1p-24
+        super().__init__()
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        l2_norm = torch.linalg.norm(x, dim=-1, keepdim=True)
+        denom = torch.maximum(l2_norm, torch.full_like(l2_norm, self.eps))
+        normalized_x = x / denom
+        return normalized_x * self.scale * self.gamma
+
+
+class RotaryEmbeddings(torch.nn.Module):
+    """
+    RoPE 用の sin・cos テーブルをキャッシュし、必要に応じて伸張／切り詰めるクラス。
+
+    Args:
+        head_dim: 1 ヘッドあたりの埋め込み次元 (必ず偶数にする)
+        max_seq_len: 事前に準備しておく最大シーケンス長
+        base_theta: 周波数スケーリング係数 (多くの論文では 10000.0)
+        learned: True にすると sin, cos をパラメータとして学習させられる
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        max_seq_len: int = 2048,
+        base_theta: float = 10000.0,
+        learned: bool = False,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+
+        if head_dim % 2 != 0:
+            raise ValueError("head_dim (＝1 ヘッドの次元数) は偶数にしてください。")
+
+        self.head_dim = head_dim
+        self.base_theta = base_theta
+        self.max_seq_len = max_seq_len
+
+        # 角周波数: θ_k = (θ_base)^(2k / d)
+        freqs = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base_theta ** (freqs / head_dim))  # (head_dim/2,)
+
+        # time 方向へアウター積 → (max_seq_len, head_dim/2)
+        t = torch.arange(max_seq_len, device=device).float()
+        sinusoid_inp = torch.einsum("i,j->ij", t, inv_freq)
+
+        sin, cos = (
+            sinusoid_inp.sin(),
+            sinusoid_inp.cos(),
+        )  # 各が (max_seq_len, head_dim/2)
+
+        if learned:
+            self.register_parameter("sin_cached", torch.nn.Parameter(sin))
+            self.register_parameter("cos_cached", torch.nn.Parameter(cos))
+        else:
+            self.register_buffer("sin_cached", sin, persistent=False)
+            self.register_buffer("cos_cached", cos, persistent=False)
+
+    def forward(
+        self,
+        seq_len: int,
+        dtype: torch.dtype = torch.float32,
+        device: torch.device | None = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        指定長の sin, cos を返す。
+
+        Returns:
+            cos: (seq_len, head_dim/2)
+            sin: (seq_len, head_dim/2)
+        """
+        if seq_len > self.max_seq_len:
+            raise ValueError(f"要求シーケンス長 {seq_len} は max_seq_len={self.max_seq_len} を超えています。")
+        cos = self.cos_cached[:seq_len].to(dtype=dtype, device=device)
+        sin = self.sin_cached[:seq_len].to(dtype=dtype, device=device)
+        return cos, sin
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """
+    偶数次元のテンソルを (…, 2i, 2i+1) → (…, -2i+1, 2i) のように 90° 回転させる。
+    具体的には (x_even, x_odd) → (-x_odd, x_even)。
+    """
+    x_even, x_odd = x[..., 0::2], x[..., 1::2]
+    return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)
+
+
+def apply_rotary_embedding(
+    query: torch.Tensor,  # (batch, num_heads, seq_len, head_dim)
+    key: torch.Tensor,  # (batch, num_heads, seq_len, head_dim)
+    cos: torch.Tensor,  # (seq_len, head_dim/2)
+    sin: torch.Tensor,  # (seq_len, head_dim/2)
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    cos = cos[None, None, :, :]
+    sin = sin[None, None, :, :]
+    cos = torch.repeat_interleave(cos, 2, dim=-1)
+    sin = torch.repeat_interleave(sin, 2, dim=-1)
+
+    q_rot = query * cos + rotate_half(query) * sin
+    k_rot = key * cos + rotate_half(key) * sin
+    return q_rot, k_rot
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, ffn_hidden_size_factor=4, dropout=0.0):
+        super().__init__()
+        dim_inner = int(dim * ffn_hidden_size_factor)
+        self.net = nn.Sequential(
+            RMSNorm(dim),
+            nn.Linear(dim, dim_inner),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim_inner, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        num_heads=8,
+        head_dim=64,
+        shared_qkv_bias=None,
+        shared_out_bias=None,
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.hidden_size = head_dim * num_heads
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+        self.norm = RMSNorm(input_dim)
+        self.to_qkv = nn.Linear(input_dim, self.hidden_size * 3, bias=(shared_qkv_bias is not None))
+        if shared_qkv_bias is not None:
+            self.to_qkv.bias = shared_qkv_bias
+
+        self.to_gates = nn.Linear(input_dim, num_heads)
+        self.to_out = nn.Sequential(
+            nn.Linear(self.hidden_size, input_dim, bias=(shared_out_bias is not None)),
+            nn.Dropout(dropout),
+        )
+        if shared_out_bias is not None:
+            self.to_out[0].bias = shared_out_bias
+
+        self.rope = RotaryEmbeddings(
+            head_dim=self.head_dim,
+            learned=False,
+        )
+        self.lowp_dtype = choose_low_precision_dtype()
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        q, k, v = einops.rearrange(self.to_qkv(x), "b t (qkv h d) -> qkv b h t d", qkv=3, h=self.num_heads)
+
+        cos, sin = self.rope(q.shape[-2], dtype=q.dtype, device=q.device)
+        q, k = apply_rotary_embedding(q, k, cos, sin)
+
+        q = q.to(self.lowp_dtype)
+        k = k.to(self.lowp_dtype)
+        v = v.to(self.lowp_dtype)
+        with sdpa_kernel(
+            [
+                SDPBackend.FLASH_ATTENTION,
+                SDPBackend.EFFICIENT_ATTENTION,
+                SDPBackend.MATH,
+            ]
+        ):
+            fetched = F.scaled_dot_product_attention(q, k, v)
+
+        gates = self.to_gates(x)
+        gates = gates.sigmoid()
+
+        out = fetched.to(x.dtype) * einops.rearrange(gates, "b n h -> b h n 1")
+        out = einops.rearrange(out, "b h t d -> b t (h d)")
+        return self.to_out(out)
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        head_dim: int,
+        num_heads: int,
+        num_layers: int,
+        ffn_hidden_size_factor: int = 4,
+        dropout: float = 0.0,
+        shared_qkv_bias=None,
+        shared_out_bias=None,
+        output_norm: bool = False,
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+
+        for _ in range(num_layers):
+            attention = MultiHeadAttention(
+                input_dim=input_dim,
+                head_dim=head_dim,
+                num_heads=num_heads,
+                shared_qkv_bias=shared_qkv_bias,
+                shared_out_bias=shared_out_bias,
+            )
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        attention,
+                        FeedForward(dim=input_dim, ffn_hidden_size_factor=ffn_hidden_size_factor),
+                    ]
+                )
+            )
+
+        self.norm = RMSNorm(input_dim) if output_norm else nn.Identity()
+
+    def forward(self, x):
+        # x: [B, T, F]
+        for attention, ffn in self.layers:
+            x = attention(x) + x
+            x = ffn(x) + x
+
+        x = self.norm(x)
+        return x

stem_splitter-0.0.1.dist-info/METADATA

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: stem-splitter
+Version: 0.0.1
+Summary: Simple, readable audio stem separation library
+Author: anime-song
+License: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.24
+Requires-Dist: tqdm>=4.66
+Requires-Dist: librosa>=0.10
+Requires-Dist: einops>=0.7
+Requires-Dist: torch
+Requires-Dist: torchaudio
+Requires-Dist: torchcodec

stem_splitter-0.0.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+stem_splitter/band_split_roformer.py,sha256=-Ni4F0DBitLhWjmFupjb46pb8UjE3zSEY8gMvcuUjjM,15810
+stem_splitter/cli.py,sha256=afs4cil4jXAR35sO26RcPWfRLSs7VRwvOC9Fn65e1JY,67
+stem_splitter/inference.py,sha256=aXVYczaXpWiKsHDQpMSY1lZtqsIyEI-nEX-W892TQRM,18396
+stem_splitter/transformer.py,sha256=JHlqbMFBI_8zSXhyp5G8bpXAvdldhDvIWMLlg0m0ZT8,8696
+stem_splitter-0.0.1.dist-info/METADATA,sha256=Izh4xvehfyAu-tL7tKbBUPuVsCcbtLpe871suBr3IHo,391
+stem_splitter-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+stem_splitter-0.0.1.dist-info/entry_points.txt,sha256=rMaUzno9dQgUZT0GH6j1NTieeDjQIOW8Nrri-EDojYg,57
+stem_splitter-0.0.1.dist-info/top_level.txt,sha256=d-BKlmkJUV_uluTwOpHhfFhwyzD_9hYPI4VCmxiElwU,14
+stem_splitter-0.0.1.dist-info/RECORD,,

stem_splitter-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

stem_splitter-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+stem-splitter = stem_splitter.cli:main

stem_splitter-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+stem_splitter