Upload inference.py with huggingface_hub

inference.py

@@ -0,0 +1,891 @@
+#!/usr/bin/env python3
+"""
+PS4 — Target Speaker Extraction Inference Script
+=================================================
+
+Self-contained inference script for the PS4 TSE model.
+No external dependencies beyond torch, torchaudio, and numpy.
+
+Usage:
+    # Basic inference
+    python inference.py \\
+        --checkpoint checkpoint_epoch037.pt \\
+        --mix mix.wav \\
+        --enroll target_speaker.wav \\
+        --output result.wav
+
+    # Use GPU
+    python inference.py \\
+        --checkpoint checkpoint_epoch037.pt \\
+        --mix mix.wav \\
+        --enroll target.wav \\
+        --output result.wav \\
+        --device cuda
+
+    # Batch mode (process a directory of mixtures with one enrollment per file)
+    python inference.py \\
+        --checkpoint checkpoint_epoch037.pt \\
+        --mix-dir ./mixtures/ \\
+        --enroll-dir ./enrollments/ \\
+        --output-dir ./results/ \\
+        --device cuda
+
+    # List available CUDA devices
+    python inference.py --list-devices
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+
+
+# ============================================================================
+# Helper: LinearLayer (used by SpeakerFuseLayer)
+# ============================================================================
+
+class LinearLayer(nn.Module):
+    """Simple linear layer with a dummy second argument for compatibility."""
+
+    def __init__(self, in_features, out_features, bias=True):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias)
+
+    def forward(self, x, dummy: Optional[torch.Tensor] = None):
+        return self.linear(x)
+
+
+# ============================================================================
+# Speaker helper modules
+# ============================================================================
+
+class PreEmphasis(nn.Module):
+    """Pre-emphasis filter: y(t) = x(t) - coef * x(t-1)."""
+
+    def __init__(self, coef: float = 0.97):
+        super().__init__()
+        self.coef = coef
+        self.register_buffer(
+            "flipped_filter",
+            torch.FloatTensor([-self.coef, 1.0]).unsqueeze(0).unsqueeze(0),
+        )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        input = input.unsqueeze(1)
+        input = F.pad(input, (1, 0), "reflect")
+        return F.conv1d(input, self.flipped_filter).squeeze(1)
+
+
+class SpeakerTransform(nn.Module):
+    """Transform speaker embeddings through a series of 1x1 conv layers."""
+
+    def __init__(self, embed_dim=256, num_layers=3, hid_dim=128):
+        super().__init__()
+        layers = []
+        layers.append(nn.Conv1d(embed_dim, hid_dim, 1))
+        for _ in range(num_layers - 2):
+            layers.append(nn.Conv1d(hid_dim, hid_dim, 1))
+            layers.append(nn.Tanh())
+        layers.append(nn.Conv1d(hid_dim, embed_dim, 1))
+        self.transforms = nn.Sequential(*layers)
+
+    def forward(self, x):
+        if len(x.size()) == 2:
+            return self.transforms(x.unsqueeze(-1)).squeeze(-1)
+        return self.transforms(x)
+
+
+class SpeakerFuseLayer(nn.Module):
+    """Fuse speaker embedding with audio features via various fusion strategies."""
+
+    def __init__(self, embed_dim=256, feat_dim=512, fuse_type="concat"):
+        super().__init__()
+        assert fuse_type in ["concat", "additive", "multiply", "FiLM", "None"]
+        self.fuse_type = fuse_type
+        if fuse_type == "concat":
+            self.fc = LinearLayer(embed_dim + feat_dim, feat_dim)
+        elif fuse_type in ("additive", "multiply"):
+            self.fc = LinearLayer(embed_dim, feat_dim)
+        elif fuse_type == "FiLM":
+            raise NotImplementedError("FiLM not supported in this standalone script")
+        else:
+            raise ValueError(f"Fuse type not defined: {fuse_type}")
+
+    def forward(self, x, embed):
+        if self.fuse_type == "concat":
+            if len(x.size()) == 3:
+                embed_t = embed.expand(-1, -1, x.size(2))
+                y = torch.cat([x, embed_t], 1)
+                y = torch.transpose(y, 1, 2)
+                x = torch.transpose(self.fc(y), 1, 2)
+            else:
+                embed_t = embed.expand(-1, x.size(1), -1, x.size(3))
+                y = torch.cat([x, embed_t], 2)
+                y = torch.transpose(y, 2, 3)
+                x = torch.transpose(self.fc(y), 2, 3).contiguous()
+        elif self.fuse_type == "additive":
+            if len(x.size()) == 3:
+                embed_t = embed.expand(-1, -1, x.size(2))
+                embed_t = torch.transpose(embed_t, 1, 2)
+                x = x + torch.transpose(self.fc(embed_t), 1, 2)
+            else:
+                embed_t = embed.expand(-1, x.size(1), -1, x.size(3))
+                embed_t = torch.transpose(embed_t, 2, 3)
+                x = x + torch.transpose(self.fc(embed_t), 2, 3)
+        elif self.fuse_type == "multiply":
+            if len(x.size()) == 3:
+                embed_t = embed.expand(-1, -1, x.size(2))
+                embed_t = torch.transpose(embed_t, 1, 2)
+                x = x * torch.transpose(self.fc(embed_t), 1, 2)
+            else:
+                embed_t = embed.expand(-1, x.size(1), -1, x.size(3))
+                embed_t = torch.transpose(embed_t, 2, 3)
+                x = x * torch.transpose(self.fc(embed_t), 2, 3)
+        else:
+            embed = embed.squeeze(-1)
+            x = self.fc(embed, x)
+        return x
+
+
+# ============================================================================
+# ECAPA-TDNN Speaker Encoder (for joint speaker embedding extraction)
+# ============================================================================
+
+class Conv1dReluBn(nn.Module):
+    """Conv1d + BatchNorm1d + ReLU."""
+
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1,
+                 padding=0, dilation=1, bias=True):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size,
+                              stride, padding, dilation, bias=bias)
+        self.bn = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        return self.bn(F.relu(self.conv(x)))
+
+
+class Res2Conv1dReluBn(nn.Module):
+    """Res2Conv1d + BatchNorm1d + ReLU."""
+
+    def __init__(self, channels, kernel_size=1, stride=1, padding=0,
+                 dilation=1, bias=True, scale=4):
+        super().__init__()
+        assert channels % scale == 0, f"{channels} % {scale} != 0"
+        self.scale = scale
+        self.width = channels // scale
+        self.nums = scale if scale == 1 else scale - 1
+
+        self.convs = nn.ModuleList()
+        self.bns = nn.ModuleList()
+        for _ in range(self.nums):
+            self.convs.append(
+                nn.Conv1d(self.width, self.width, kernel_size,
+                          stride, padding, dilation, bias=bias))
+            self.bns.append(nn.BatchNorm1d(self.width))
+
+    def forward(self, x):
+        out = []
+        spx = torch.split(x, self.width, 1)
+        sp = spx[0]
+        for i, (conv, bn) in enumerate(zip(self.convs, self.bns)):
+            if i >= 1:
+                sp = sp + spx[i]
+            sp = conv(sp)
+            sp = bn(F.relu(sp))
+            out.append(sp)
+        if self.scale != 1:
+            out.append(spx[self.nums])
+        return torch.cat(out, dim=1)
+
+
+class SE_Connect(nn.Module):
+    """Squeeze-Excitation block for 1D."""
+
+    def __init__(self, channels, se_bottleneck_dim=128):
+        super().__init__()
+        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
+        self.linear2 = nn.Linear(se_bottleneck_dim, channels)
+
+    def forward(self, x):
+        out = x.mean(dim=2)
+        out = F.relu(self.linear1(out))
+        out = torch.sigmoid(self.linear2(out))
+        return x * out.unsqueeze(2)
+
+
+class SE_Res2Block(nn.Module):
+    """SE-Res2Block of the ECAPA-TDNN architecture."""
+
+    def __init__(self, channels, kernel_size, stride, padding, dilation, scale):
+        super().__init__()
+        self.se_res2block = nn.Sequential(
+            Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0),
+            Res2Conv1dReluBn(channels, kernel_size, stride, padding, dilation, scale=scale),
+            Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0),
+            SE_Connect(channels),
+        )
+
+    def forward(self, x):
+        return x + self.se_res2block(x)
+
+
+class ASTP(nn.Module):
+    """Attentive statistics pooling: first used in ECAPA-TDNN."""
+
+    def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False, **kwargs):
+        super().__init__()
+        self.in_dim = in_dim
+        self.global_context_att = global_context_att
+        if global_context_att:
+            self.linear1 = nn.Conv1d(in_dim * 3, bottleneck_dim, kernel_size=1)
+        else:
+            self.linear1 = nn.Conv1d(in_dim, bottleneck_dim, kernel_size=1)
+        self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, kernel_size=1)
+
+    def forward(self, x):
+        if len(x.shape) == 4:
+            x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])
+        assert len(x.shape) == 3
+
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-7).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+
+        alpha = torch.tanh(self.linear1(x_in))
+        alpha = torch.softmax(self.linear2(alpha), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        var = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
+        std = torch.sqrt(var.clamp(min=1e-7))
+        return torch.cat([mean, std], dim=1)
+
+    def get_out_dim(self):
+        return 2 * self.in_dim
+
+
+class ECAPA_TDNN(nn.Module):
+    """ECAPA-TDNN speaker encoder."""
+
+    def __init__(self, channels=512, feat_dim=80, embed_dim=192,
+                 pooling_func="ASTP", global_context_att=False, emb_bn=False):
+        super().__init__()
+        self.layer1 = Conv1dReluBn(feat_dim, channels, kernel_size=5, padding=2)
+        self.layer2 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=2, dilation=2, scale=8)
+        self.layer3 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=3, dilation=3, scale=8)
+        self.layer4 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=4, dilation=4, scale=8)
+
+        cat_channels = channels * 3
+        out_channels = 512 * 3
+        self.conv = nn.Conv1d(cat_channels, out_channels, kernel_size=1)
+        self.pool = ASTP(in_dim=out_channels, global_context_att=global_context_att)
+        self.pool_out_dim = self.pool.get_out_dim()
+        self.bn = nn.BatchNorm1d(self.pool_out_dim)
+        self.linear = nn.Linear(self.pool_out_dim, embed_dim)
+        self.emb_bn = emb_bn
+        self.bn2 = nn.BatchNorm1d(embed_dim) if emb_bn else nn.Identity()
+
+    def forward(self, x):
+        x = x.permute(0, 2, 1)  # (B, T, F) -> (B, F, T)
+        out1 = self.layer1(x)
+        out2 = self.layer2(out1)
+        out3 = self.layer3(out2)
+        out4 = self.layer4(out3)
+        out = torch.cat([out2, out3, out4], dim=1)
+        out = self.conv(out)
+        out = F.relu(out)
+        out = self.bn(self.pool(out))
+        out = self.linear(out)
+        if self.emb_bn:
+            out = self.bn2(out)
+        return out4, out  # returns (frame_level, segment_level)
+
+
+# ============================================================================
+# BSRNN Legacy Model
+# ============================================================================
+
+class ResRNN(nn.Module):
+    """Residual LSTM with GroupNorm + projection."""
+
+    def __init__(self, input_size, hidden_size, bidirectional=True):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.eps = torch.finfo(torch.float32).eps
+        self.norm = nn.GroupNorm(1, input_size, self.eps)
+        self.rnn = nn.LSTM(input_size, hidden_size, 1, batch_first=True,
+                           bidirectional=bidirectional)
+        self.proj = nn.Linear(hidden_size * 2, input_size)
+
+    def forward(self, input):
+        rnn_output, _ = self.rnn(self.norm(input).transpose(1, 2).contiguous())
+        rnn_output = self.proj(
+            rnn_output.contiguous().view(-1, rnn_output.shape[2])
+        ).view(input.shape[0], input.shape[2], input.shape[1])
+        return input + rnn_output.transpose(1, 2).contiguous()
+
+
+class BSNet(nn.Module):
+    """Band-split network with intra-band and inter-band RNN."""
+
+    def __init__(self, in_channel, nband=7, bidirectional=True):
+        super().__init__()
+        self.nband = nband
+        self.feature_dim = in_channel // nband
+        self.band_rnn = ResRNN(self.feature_dim, self.feature_dim * 2,
+                               bidirectional=bidirectional)
+        self.band_comm = ResRNN(self.feature_dim, self.feature_dim * 2,
+                                bidirectional=bidirectional)
+
+    def forward(self, input, dummy: Optional[torch.Tensor] = None):
+        B, N, T = input.shape
+        band_output = self.band_rnn(
+            input.view(B * self.nband, self.feature_dim, -1)
+        ).view(B, self.nband, -1, T)
+        band_output = band_output.permute(0, 3, 2, 1).contiguous().view(
+            B * T, -1, self.nband)
+        output = self.band_comm(band_output).view(
+            B, T, -1, self.nband).permute(0, 3, 2, 1).contiguous()
+        return output.view(B, N, T)
+
+
+class FuseSeparation(nn.Module):
+    """Separation module with speaker fusion at each repeat."""
+
+    def __init__(self, nband=7, num_repeat=6, feature_dim=128,
+                 spk_emb_dim=256, spk_fuse_type="concat", multi_fuse=True):
+        super().__init__()
+        self.multi_fuse = multi_fuse
+        self.nband = nband
+        self.feature_dim = feature_dim
+        self.separation = nn.ModuleList([])
+        if self.multi_fuse:
+            for _ in range(num_repeat):
+                self.separation.append(
+                    SpeakerFuseLayer(embed_dim=spk_emb_dim,
+                                     feat_dim=feature_dim,
+                                     fuse_type=spk_fuse_type))
+                self.separation.append(BSNet(nband * feature_dim, nband))
+        else:
+            self.separation.append(
+                SpeakerFuseLayer(embed_dim=spk_emb_dim,
+                                 feat_dim=feature_dim,
+                                 fuse_type=spk_fuse_type))
+            for _ in range(num_repeat):
+                self.separation.append(BSNet(nband * feature_dim, nband))
+
+    def forward(self, x, spk_embedding, nch: torch.Tensor = torch.tensor(1)):
+        batch_size = x.shape[0]
+        if self.multi_fuse:
+            for i, sep_func in enumerate(self.separation):
+                x = sep_func(x, spk_embedding)
+                if i % 2 == 0:
+                    x = x.view(batch_size * nch,
+                               self.nband * self.feature_dim, -1)
+                else:
+                    x = x.view(batch_size * nch, self.nband,
+                               self.feature_dim, -1)
+        else:
+            x = self.separation[0](x, spk_embedding)
+            x = x.view(batch_size * nch, self.nband * self.feature_dim, -1)
+            for idx, sep in enumerate(self.separation):
+                if idx > 0:
+                    x = sep(x, spk_embedding)
+            x = x.view(batch_size * nch, self.nband, self.feature_dim, -1)
+        return x
+
+
+class BSRNN(nn.Module):
+    """Legacy BSRNN with joint speaker encoder (flat-config format).
+
+    This is the exact model architecture used to train the PS4 checkpoint.
+    State dict keys: separator.separation.*, spk_model.layer*, mask.*, BN.*,
+    spk_encoder.*, preEmphasis.*
+    """
+
+    def __init__(
+        self,
+        spk_emb_dim=256,
+        sr=16000,
+        win=512,
+        stride=128,
+        feature_dim=128,
+        num_repeat=6,
+        use_spk_transform=True,
+        use_bidirectional=True,
+        spk_fuse_type="concat",
+        multi_fuse=True,
+        joint_training=True,
+        multi_task=False,
+        spksInTrain=251,
+        spk_model=None,
+        spk_model_init=None,
+        spk_model_freeze=False,
+        spk_args=None,
+        spk_feat=False,
+        feat_type="consistent",
+    ):
+        super().__init__()
+        self.sr = sr
+        self.win = win
+        self.stride = stride
+        self.group = self.win // 2
+        self.enc_dim = self.win // 2 + 1
+        self.feature_dim = feature_dim
+        self.eps = torch.finfo(torch.float32).eps
+        self.spk_emb_dim = spk_emb_dim
+        self.joint_training = joint_training
+        self.spk_feat = spk_feat
+        self.feat_type = feat_type
+        self.spk_model_freeze = spk_model_freeze
+        self.multi_task = multi_task
+
+        # Band split: 100Hz bins → 200Hz bins → 500Hz bins → 2kHz bins → rest
+        bandwidth_100 = int(np.floor(100 / (sr / 2.0) * self.enc_dim))
+        bandwidth_200 = int(np.floor(200 / (sr / 2.0) * self.enc_dim))
+        bandwidth_500 = int(np.floor(500 / (sr / 2.0) * self.enc_dim))
+        bandwidth_2k = int(np.floor(2000 / (sr / 2.0) * self.enc_dim))
+        self.band_width = [bandwidth_100] * 15
+        self.band_width += [bandwidth_200] * 10
+        self.band_width += [bandwidth_500] * 5
+        self.band_width += [bandwidth_2k] * 1
+        self.band_width.append(self.enc_dim - int(np.sum(self.band_width)))
+        self.nband = len(self.band_width)
+
+        # Speaker embedding transform
+        if use_spk_transform:
+            self.spk_transform = SpeakerTransform()
+        else:
+            self.spk_transform = nn.Identity()
+
+        # Joint speaker encoder
+        if joint_training:
+            spk_args = spk_args or {}
+            self.spk_model = ECAPA_TDNN_GLOB_c512(
+                feat_dim=spk_args.get("feat_dim", 80),
+                embed_dim=spk_args.get("embed_dim", 192),
+                pooling_func=spk_args.get("pooling_func", "ASTP"),
+            )
+            if spk_model_freeze:
+                for param in self.spk_model.parameters():
+                    param.requires_grad = False
+            if not spk_feat:
+                if feat_type == "consistent":
+                    self.preEmphasis = PreEmphasis()
+                    self.spk_encoder = torchaudio.transforms.MelSpectrogram(
+                        sample_rate=sr,
+                        n_fft=win,
+                        win_length=win,
+                        hop_length=stride,
+                        f_min=20,
+                        window_fn=torch.hamming_window,
+                        n_mels=spk_args.get("feat_dim", 80),
+                    )
+            else:
+                self.preEmphasis = nn.Identity()
+                self.spk_encoder = nn.Identity()
+
+            if multi_task:
+                self.pred_linear = nn.Linear(spk_emb_dim, spksInTrain)
+            else:
+                self.pred_linear = nn.Identity()
+
+        # Band normalization
+        self.BN = nn.ModuleList([])
+        for i in range(self.nband):
+            self.BN.append(
+                nn.Sequential(
+                    nn.GroupNorm(1, self.band_width[i] * 2, self.eps),
+                    nn.Conv1d(self.band_width[i] * 2, self.feature_dim, 1),
+                )
+            )
+
+        # Separator
+        self.separator = FuseSeparation(
+            nband=self.nband,
+            num_repeat=num_repeat,
+            feature_dim=feature_dim,
+            spk_emb_dim=spk_emb_dim,
+            spk_fuse_type=spk_fuse_type,
+            multi_fuse=multi_fuse,
+        )
+
+        # Mask estimation
+        self.mask = nn.ModuleList([])
+        for i in range(self.nband):
+            self.mask.append(
+                nn.Sequential(
+                    nn.GroupNorm(1, self.feature_dim, torch.finfo(torch.float32).eps),
+                    nn.Conv1d(self.feature_dim, self.feature_dim * 4, 1),
+                    nn.Tanh(),
+                    nn.Conv1d(self.feature_dim * 4, self.feature_dim * 4, 1),
+                    nn.Tanh(),
+                    nn.Conv1d(self.feature_dim * 4, self.band_width[i] * 4, 1),
+                )
+            )
+
+    def train(self, mode: bool = True):
+        """Override train(): keep spk_model in eval mode when frozen."""
+        super().train(mode)
+        if self.spk_model_freeze and hasattr(self, "spk_model"):
+            self.spk_model.eval()
+        return self
+
+    def forward(self, input, embeddings):
+        """
+        Args:
+            input: (B, T) mixture waveform
+            embeddings: (B, T_enroll) enrollment waveform (will be processed
+                        by the internal speaker encoder, or (B, D) pre-extracted
+                        speaker embedding if spk_feat=True)
+        Returns:
+            s: (B, T) extracted target speaker waveform
+            _: dummy speaker label prediction (ignored at inference)
+        """
+        wav_input = input
+        spk_emb_input = embeddings
+        batch_size, nsample = wav_input.shape
+        nch = 1
+
+        # STFT
+        spec = torch.stft(
+            wav_input,
+            n_fft=self.win,
+            hop_length=self.stride,
+            window=torch.hann_window(self.win).to(wav_input.device).type(
+                wav_input.type()),
+            return_complex=True,
+        )
+        spec_RI = torch.stack([spec.real, spec.imag], 1)
+
+        # Band split
+        subband_spec = []
+        subband_mix_spec = []
+        band_idx = 0
+        for i in range(len(self.band_width)):
+            subband_spec.append(
+                spec_RI[:, :, band_idx:band_idx + self.band_width[i]].contiguous())
+            subband_mix_spec.append(
+                spec[:, band_idx:band_idx + self.band_width[i]])
+            band_idx += self.band_width[i]
+
+        # Band normalization
+        subband_feature = []
+        for i, bn_func in enumerate(self.BN):
+            subband_feature.append(
+                bn_func(subband_spec[i].view(batch_size * nch,
+                                             self.band_width[i] * 2, -1)))
+        subband_feature = torch.stack(subband_feature, 1)
+
+        predict_speaker_lable = torch.tensor(0.0).to(spk_emb_input.device)
+
+        # Joint speaker encoder
+        if self.joint_training:
+            if not self.spk_feat:
+                if self.feat_type == "consistent":
+                    with torch.no_grad():
+                        spk_emb_input = self.preEmphasis(spk_emb_input)
+                        spk_emb_input = self.spk_encoder(spk_emb_input) + 1e-8
+                        spk_emb_input = spk_emb_input.log()
+                        spk_emb_input = spk_emb_input - torch.mean(
+                            spk_emb_input, dim=-1, keepdim=True)
+                        spk_emb_input = spk_emb_input.permute(0, 2, 1)
+
+            tmp_spk_emb_input = self.spk_model(spk_emb_input)
+            if isinstance(tmp_spk_emb_input, tuple):
+                spk_emb_input = tmp_spk_emb_input[-1]
+            else:
+                spk_emb_input = tmp_spk_emb_input
+            predict_speaker_lable = self.pred_linear(spk_emb_input)
+
+        spk_embedding = self.spk_transform(spk_emb_input)
+        spk_embedding = spk_embedding.unsqueeze(1).unsqueeze(3)
+
+        # Separation
+        sep_output = self.separator(subband_feature, spk_embedding,
+                                    torch.tensor(nch))
+
+        # Mask estimation and complex mask application
+        sep_subband_spec = []
+        for i, mask_func in enumerate(self.mask):
+            this_output = mask_func(sep_output[:, i]).view(
+                batch_size * nch, 2, 2, self.band_width[i], -1)
+            this_mask = this_output[:, 0] * torch.sigmoid(this_output[:, 1])
+            this_mask_real = this_mask[:, 0]
+            this_mask_imag = this_mask[:, 1]
+            est_spec_real = (subband_mix_spec[i].real * this_mask_real
+                             - subband_mix_spec[i].imag * this_mask_imag)
+            est_spec_imag = (subband_mix_spec[i].real * this_mask_imag
+                             + subband_mix_spec[i].imag * this_mask_real)
+            sep_subband_spec.append(
+                torch.complex(est_spec_real, est_spec_imag))
+
+        # iSTFT
+        est_spec = torch.cat(sep_subband_spec, 1)
+        output = torch.istft(
+            est_spec.view(batch_size * nch, self.enc_dim, -1),
+            n_fft=self.win,
+            hop_length=self.stride,
+            window=torch.hann_window(self.win).to(wav_input.device).type(
+                wav_input.type()),
+            length=nsample,
+        )
+        output = output.view(batch_size, nch, -1)
+        s = torch.squeeze(output, dim=1)
+        return s, predict_speaker_lable
+
+
+def ECAPA_TDNN_GLOB_c512(feat_dim, embed_dim, pooling_func="ASTP", emb_bn=False):
+    """Factory function for ECAPA-TDNN with global context attention and 512 channels."""
+    return ECAPA_TDNN(
+        channels=512,
+        feat_dim=feat_dim,
+        embed_dim=embed_dim,
+        pooling_func=pooling_func,
+        global_context_att=True,
+        emb_bn=emb_bn,
+    )
+
+
+# ============================================================================
+# Checkpoint Loading
+# ============================================================================
+
+def build_model(device: torch.device) -> BSRNN:
+    """Build the PS4 BSRNN model with the exact training config parameters.
+
+    Returns:
+        BSRNN model in eval mode, moved to the specified device.
+    """
+    model = BSRNN(
+        feat_type="consistent",
+        feature_dim=128,
+        num_repeat=6,
+        spk_emb_dim=192,
+        spk_fuse_type="multiply",
+        multi_fuse=False,
+        spk_model="ECAPA_TDNN_GLOB_c512",
+        sr=16000,
+        win=512,
+        stride=128,
+        spk_args={"feat_dim": 80, "embed_dim": 192, "pooling_func": "ASTP"},
+        spk_model_freeze=True,
+        use_spk_transform=False,
+        joint_training=True,
+        multi_task=False,
+        spk_feat=False,
+    )
+    model.eval()
+    return model.to(device)
+
+
+def load_checkpoint(path: str, model: nn.Module, device: torch.device):
+    """Load PS4 checkpoint weights into the model.
+
+    The checkpoint is saved by train.py as:
+        {"model": state_dict, "optimizer": ..., "scheduler": ..., ...}
+    """
+    print(f"[PS4] Loading checkpoint: {path}")
+    ckpt = torch.load(path, map_location=device)
+
+    # Handle various checkpoint formats
+    if isinstance(ckpt, dict) and "model" in ckpt:
+        state_dict = ckpt["model"]
+    elif isinstance(ckpt, dict) and "state_dict" in ckpt:
+        state_dict = ckpt["state_dict"]
+    else:
+        state_dict = ckpt
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    if missing:
+        print(f"[PS4]  WARNING: missing keys ({len(missing)}): {missing[:5]}...")
+    if unexpected:
+        print(f"[PS4]  WARNING: unexpected keys ({len(unexpected)}): {unexpected[:5]}...")
+    print(f"[PS4]  Loaded successfully. "
+          f"Epoch: {ckpt.get('epoch', 'N/A')}, "
+          f"Step: {ckpt.get('step', 'N/A')}")
+    return model
+
+
+# ============================================================================
+# Inference
+# ============================================================================
+
+def load_audio(path: str, target_sr: int = 16000) -> torch.Tensor:
+    """Load audio at target sample rate.
+
+    Returns:
+        Tensor of shape (1, T) — mono, normalized to [-1, 1].
+    """
+    wav, sr = torchaudio.load(path)
+    if wav.size(0) > 1:
+        wav = wav.mean(dim=0, keepdim=True)  # mono
+    if sr != target_sr:
+        wav = torchaudio.functional.resample(wav, sr, target_sr)
+    # Normalize
+    peak = wav.abs().max()
+    if peak > 0:
+        wav = wav / peak
+    return wav
+
+
+def save_audio(path: str, wav: torch.Tensor, sr: int = 16000):
+    """Save audio tensor to file."""
+    torchaudio.save(path, wav.cpu(), sr)
+    print(f"[PS4] Saved: {path}")
+
+
+def extract_speaker(
+    model: BSRNN,
+    mixture: torch.Tensor,
+    enrollment: torch.Tensor,
+    device: torch.device,
+) -> torch.Tensor:
+    """Run target speaker extraction.
+
+    Args:
+        model: Loaded BSRNN model.
+        mixture: (1, T_mix) mixture waveform, 16 kHz.
+        enrollment: (1, T_enroll) enrollment waveform, 16 kHz.
+        device: Computation device.
+
+    Returns:
+        (1, T_mix) extracted target speaker waveform.
+    """
+    with torch.no_grad():
+        mixture = mixture.to(device)
+        enrollment = enrollment.to(device)
+        extracted, _ = model(mixture, enrollment)
+    return extracted.cpu()
+
+
+# ============================================================================
+# CLI
+# ============================================================================
+
+def list_devices():
+    """Print available CUDA devices."""
+    print("Available devices:")
+    print(f"  cpu")
+    if torch.cuda.is_available():
+        for i in range(torch.cuda.device_count()):
+            print(f"  cuda:{i}  {torch.cuda.get_device_name(i)}")
+    else:
+        print("  (no CUDA devices found)")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="PS4 Target Speaker Extraction — Inference",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Single file
+  python inference.py --checkpoint checkpoint_epoch037.pt \\
+      --mix mix.wav --enroll target.wav --output result.wav
+
+  # Directory batch
+  python inference.py --checkpoint checkpoint_epoch037.pt \\
+      --mix-dir ./mixtures/ --enroll-dir ./enrollments/ --output-dir ./results/
+
+  # List devices
+  python inference.py --list-devices
+        """,
+    )
+    parser.add_argument("--checkpoint", type=str,
+                        default="checkpoint_epoch037.pt",
+                        help="Path to PS4 checkpoint (.pt)")
+    parser.add_argument("--mix", type=str, default=None,
+                        help="Path to mixture audio (16 kHz mono WAV)")
+    parser.add_argument("--enroll", type=str, default=None,
+                        help="Path to enrollment audio (16 kHz mono WAV)")
+    parser.add_argument("--output", type=str, default="output.wav",
+                        help="Path to save extracted audio")
+    parser.add_argument("--mix-dir", type=str, default=None,
+                        help="Directory of mixture audio files (batch mode)")
+    parser.add_argument("--enroll-dir", type=str, default=None,
+                        help="Directory of enrollment audio files (batch mode, "
+                             "must match mixture filenames)")
+    parser.add_argument("--output-dir", type=str, default=None,
+                        help="Output directory (batch mode)")
+    parser.add_argument("--device", type=str, default="auto",
+                        help="Device: 'auto', 'cpu', or 'cuda:N'")
+    parser.add_argument("--list-devices", action="store_true",
+                        help="List available devices and exit")
+
+    args = parser.parse_args()
+
+    if args.list_devices:
+        list_devices()
+        return
+
+    # Device selection
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+    print(f"[PS4] Using device: {device}")
+
+    # Build model
+    print("[PS4] Building model...")
+    model = build_model(device)
+    load_checkpoint(args.checkpoint, model, device)
+    print(f"[PS4] Model parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Single file mode
+    if args.mix is not None and args.enroll is not None:
+        print(f"[PS4] Loading mixture: {args.mix}")
+        mix = load_audio(args.mix)
+        print(f"[PS4] Loading enrollment: {args.enroll}")
+        enroll = load_audio(args.enroll)
+        print(f"[PS4] Running extraction (mix: {mix.shape[-1]/16000:.1f}s, "
+              f"enroll: {enroll.shape[-1]/16000:.1f}s)...")
+        extracted = extract_speaker(model, mix, enroll, device)
+        save_audio(args.output, extracted)
+        return
+
+    # Batch mode
+    if args.mix_dir is not None and args.enroll_dir is not None and args.output_dir is not None:
+        mix_dir = Path(args.mix_dir)
+        enroll_dir = Path(args.enroll_dir)
+        output_dir = Path(args.output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        mix_files = sorted(mix_dir.glob("*.wav"))
+        if not mix_files:
+            print(f"[PS4] No .wav files found in {mix_dir}")
+            return
+
+        print(f"[PS4] Batch mode: {len(mix_files)} files")
+        for mix_path in mix_files:
+            enroll_path = enroll_dir / mix_path.name
+            if not enroll_path.exists():
+                print(f"[PS4]  Skipping {mix_path.name}: no matching enrollment")
+                continue
+            out_path = output_dir / mix_path.name
+            print(f"[PS4]  Processing {mix_path.name}...", end=" ", flush=True)
+            mix = load_audio(str(mix_path))
+            enroll = load_audio(str(enroll_path))
+            extracted = extract_speaker(model, mix, enroll, device)
+            save_audio(str(out_path), extracted)
+            print("done")
+        return
+
+    # If neither mode is specified
+    parser.print_help()
+    print("\n[PS4] ERROR: Specify either --mix/--enroll (single) or "
+          "--mix-dir/--enroll-dir/--output-dir (batch).")
+
+
+if __name__ == "__main__":
+    main()