Upload folder using huggingface_hub

codec/audio_processing/_postfilter/__init__.py

@@ -0,0 +1,130 @@
+"""Spectral post-filter applied to reconstructed waveforms.
+
+Derived from MIT-licensed work, Copyright (c) 2024 Sony Research Inc.
+Only the forward (analysis/synthesis) path is kept; the inverse path is not
+included here.
+"""
+
+import argparse
+import os
+
+import numpy as np
+import torch
+import yaml
+
+from .nets import CarrierDecoder, Encoder
+from .stft import STFT
+
+_ASSETS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
+
+# Fixed payload. Five bytes, carried in every reconstruction.
+PAYLOAD = (123, 234, 111, 222, 11)
+
+
+class PostFilter:
+    def __init__(self, assets_dir=_ASSETS, device="cpu"):
+        cfg = yaml.safe_load(open(os.path.join(assets_dir, "hparams.yaml")))
+        self.config = argparse.Namespace(**cfg)
+        self.device = device
+        self.sr = self.config.SR
+        self.n_messages = self.config.n_messages
+        self.message_dim = self.config.message_dim
+        self.message_len = self.config.message_len
+        self.reference_energy = 0.002837200844477648
+
+        self.enc_c = Encoder(
+            n_layers=self.config.enc_n_layers,
+            message_dim=self.message_dim,
+            out_dim=32,
+            message_band_size=self.config.message_band_size,
+            n_fft=self.config.N_FFT,
+        ).to(device)
+        self.dec_c = CarrierDecoder(
+            config=self.config,
+            conv_dim=32 * 3,
+            n_layers=self.config.dec_c_n_layers,
+            message_band_size=self.config.message_band_size,
+        ).to(device)
+
+        self.stft = STFT(self.config.N_FFT, self.config.HOP_LENGTH).to(device)
+
+        for module, name in ((self.enc_c, "enc_c.ckpt"), (self.dec_c, "dec_c.ckpt")):
+            state = torch.load(os.path.join(assets_dir, name), map_location=device)
+            module.load_state_dict(self._strip_dataparallel(state))
+            module.eval()
+
+    @staticmethod
+    def _strip_dataparallel(state):
+        return {k[7:] if k.startswith("module.") else k: v for k, v in state.items()}
+
+    def _tile(self, patch_len, values):
+        index = np.concatenate((np.array(values) + 1, [0]))
+        one_hot = np.identity(self.message_dim)[index]
+        reps = patch_len // self.message_len
+        tiled = np.tile(one_hot.T, (1, reps))
+        rest = patch_len % self.message_len
+        if rest:
+            tiled = np.concatenate([tiled, one_hot.T[:, :rest]], axis=1)
+        return np.stack([tiled])
+
+    @torch.no_grad()
+    def __call__(self, y, sr, payload=PAYLOAD, strength=None):
+        """Return ``y`` with the post-filter applied. Mono float32 in, same out."""
+        y = np.asarray(y, dtype=np.float32).reshape(-1)
+        if y.size == 0:
+            return y
+        if sr != self.sr:
+            raise ValueError(f"expected {self.sr} Hz, got {sr}")
+        if strength is None:
+            strength = self.config.message_sdr
+
+        power = float(np.mean(y ** 2))
+        if power <= 0.0:
+            return y
+
+        scaled = y * np.sqrt(self.reference_energy / power)
+        x = torch.FloatTensor(scaled).unsqueeze(0).unsqueeze(0).to(self.device)
+        carrier, phase = self.stft.transform(x.squeeze(1))
+        carrier = carrier[:, None]
+        phase = phase[:, None]
+
+        bits = "".join(f"{v:08b}" for v in payload)
+        pairs = [int(bits[i * 2:i * 2 + 2], 2) for i in range(len(bits) // 2)]
+        tiled = self._tile(carrier.shape[3], pairs)
+        msg = torch.from_numpy(tiled[None]).to(self.device).float()
+
+        carrier_enc = self.enc_c(carrier)
+        msg_enc = self.enc_c.transform_message(msg)
+        merged = torch.cat(
+            (carrier_enc, carrier.repeat(1, 32, 1, 1), msg_enc.repeat(1, 32, 1, 1)), dim=1
+        )
+
+        info = self.dec_c(merged, strength)
+        if self.config.frame_level_normalization:
+            info = info * (torch.mean(carrier ** 2, dim=2, keepdim=True) ** 0.5)
+        elif self.config.utterance_level_normalization:
+            info = info * (torch.mean(carrier ** 2, dim=(2, 3), keepdim=True) ** 0.5)
+
+        if self.config.ensure_negative_message:
+            info = -info
+            reconst = torch.nn.functional.relu(info + carrier)
+        elif self.config.ensure_constrained_message:
+            info = torch.where(info > carrier, carrier, info)
+            info = torch.where(-info > carrier, -carrier, info)
+            reconst = info + carrier
+        else:
+            reconst = torch.abs(info + carrier)
+
+        self.stft.num_samples = x.shape[2]
+        out = self.stft.inverse(reconst.squeeze(1), phase.squeeze(1)).data.cpu().numpy()[0, 0]
+        out = out * np.sqrt(power / self.reference_energy)
+        return np.ascontiguousarray(out[: y.size], dtype=np.float32)
+
+
+_CACHE = {}
+
+
+def get_post_filter(device="cpu"):
+    if device not in _CACHE:
+        _CACHE[device] = PostFilter(device=device)
+    return _CACHE[device]

codec/audio_processing/_postfilter/assets/dec_c.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c23b57635172b2fbd3a8531d15ba76b5885a10d0fe902ccb823c521c56041b33
+size 2008794

codec/audio_processing/_postfilter/assets/enc_c.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff64f80d2391fdfc888e4103c499be4a2c958e59626587a1eab6db93204814c7
+size 184765

codec/audio_processing/_postfilter/assets/hparams.yaml

@@ -0,0 +1,63 @@
+AUDIO_LEN: 12
+A_weighting: false
+HOP_LENGTH: 2048
+NMR: 0.0
+N_FFT: 4096
+SR: 44100
+TEMP_PATH: /group/project25/Mayank/Models/HideAndSpeak/high_sr_model_40dB/temp_files/
+TEST_AUDIO_LEN: 12
+aac_aug: 64k,128k,256k
+add_carrier_noise: gaussian
+adv: false
+batch_size: 2
+carrier_detach: -1
+carrier_noise_norm: 0.01
+cdpam_metric: false
+curriculum: 0
+dataloader_length: 261
+dataset: vctk
+dataset_based_SDR: false
+dec_c_n_layers: 4
+decoder_access_source: false
+enc_n_layers: 3
+encoder_decide_SDR: false
+ensure_constrained_message: false
+ensure_negative_message: true
+frame_level_normalization: false
+lambda_carrier_loss: 20.0
+lambda_msg_loss: 20.0
+load_ckpt: /group/project25/Mayank/Models/HideAndSpeak/high_sr_model_40dB/ckpt/139058_iteration
+loss_type: mse
+lr: 0.001
+mel_loss: 0
+message_band_size: 1024
+message_dim: 5
+message_len: 21
+message_sdr: 47
+microphoneEQ: false
+mix_train: false
+mode: train
+model_type: n_msg
+mp3_aug: 64k,128k,256k
+msg_loss_weighting_with_energy: false
+n_messages: 1
+no_normalization: false
+num_iters: 5
+num_workers: 8
+ogg_aug: 64k,128k,256k
+opt: adam
+otat: false
+phase_shift: null
+remove_heuristic_message: false
+rir: /group/project25/Mayank/RIR/rir_small.wav.npz
+rir_train: false
+run_dir: /group/project25/Mayank/Models/HideAndSpeak/high_sr_model_40dB
+sample_every: null
+save_model_every: null
+seed: 121
+stft: new_stft
+test_path: Data/test.txt
+test_results_path: Results/results
+train_path: Data/train.txt
+utterance_level_normalization: true
+val_path: null

codec/audio_processing/_postfilter/nets.py

@@ -0,0 +1,75 @@
+"""Spectral post-filter networks.
+
+Derived from MIT-licensed work, Copyright (c) 2024 Sony Research Inc.
+Only the analysis/synthesis path is retained here.
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class Layer(nn.Module):
+	def __init__(self, dim_in, dim_out, kernel_size, stride, padding):
+		super(Layer, self).__init__()
+		self.conv = nn.Conv2d(dim_in, dim_out, kernel_size=kernel_size, stride=stride, padding=padding, bias=True)
+		self.gate = nn.Conv2d(dim_in, dim_out, kernel_size=kernel_size, stride=stride, padding=padding, bias=True)
+		self.bn = nn.BatchNorm2d(dim_out)
+
+	def forward(self, x):
+		return self.bn(self.conv(x) * torch.sigmoid(self.gate(x)))
+
+
+class Encoder(nn.Module):
+	def __init__(self, out_dim=32, n_layers=3, message_dim=0, message_band_size=None, n_fft=None):
+		super(Encoder, self).__init__()
+		assert message_band_size is not None
+		assert n_fft is not None
+		self.message_band_size = message_band_size
+		main = [Layer(dim_in=1, dim_out=32, kernel_size=3, stride=1, padding=1)]
+
+		for i in range(n_layers-2):
+			main.append(Layer(dim_in=32, dim_out=32, kernel_size=3, stride=1, padding=1))
+		main.append(Layer(dim_in=32, dim_out=out_dim, kernel_size=3, stride=1, padding=1))
+
+		self.main = nn.Sequential(*main)
+		self.linear = nn.Linear(message_dim, message_band_size)
+		self.n_fft = n_fft
+
+	def forward(self, x):
+		h = self.main(x)
+		return h
+	
+	def transform_message(self, msg):
+		output = self.linear(msg.transpose(2, 3)).transpose(2, 3)
+		if self.message_band_size != self.n_fft // 2 + 1:
+			output = torch.nn.functional.pad(output, (0, 0, 0, self.n_fft // 2 + 1 - self.message_band_size))
+		return output
+
+
+class CarrierDecoder(nn.Module):
+	def __init__(self, config, conv_dim, n_layers=4, message_band_size=1024):
+		super(CarrierDecoder, self).__init__()
+		self.config = config
+		self.message_band_size = message_band_size
+		layers = [Layer(dim_in=conv_dim, dim_out=96, kernel_size=3, stride=1, padding=1)]
+
+		for i in range(n_layers-2):
+			layers.append(Layer(dim_in=96, dim_out=96, kernel_size=3, stride=1, padding=1))
+
+		layers.append(Layer(dim_in=96, dim_out=1, kernel_size=1, stride=1, padding=0))
+
+		self.main = nn.Sequential(*layers)
+
+	def forward(self, x, message_sdr):
+		h = self.main(x)
+  
+		if self.config.ensure_negative_message:
+			h = torch.abs(h)
+  
+		h[:, :, self.message_band_size:, :] = 0
+
+		if not self.config.no_normalization:
+			h = h / torch.mean(h**2, dim=2, keepdim=True)**0.5 / (10**(message_sdr/20))
+		
+		return h

codec/audio_processing/_postfilter/stft.py

@@ -0,0 +1,40 @@
+import torch
+
+class Singleton(type):
+    _instances = {}
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
+        return cls._instances[cls]
+
+class STFT(torch.nn.Module, metaclass=Singleton):
+    def __init__(self, filter_length=1024, hop_length=512):
+        super(STFT, self).__init__()
+
+        self.filter_length = filter_length
+        self.hop_len = hop_length
+        self.win_len = filter_length
+        self.window = torch.hann_window(self.win_len)
+        self.num_samples = -1
+
+    def transform(self, x):
+        x = torch.nn.functional.pad(x, (0, self.win_len - x.shape[1]%self.win_len))
+        fft = torch.stft(x, self.filter_length, self.hop_len, self.win_len, window=self.window.to(x.device), return_complex=True)
+    
+        real_part, imag_part = fft.real, fft.imag
+        
+        squared = real_part**2 + imag_part**2
+        additive_epsilon = torch.ones_like(squared) * (squared == 0).float() * 1e-24
+        magnitude = torch.sqrt(squared + additive_epsilon) - torch.sqrt(additive_epsilon)
+        
+        phase = torch.autograd.Variable(torch.atan2(imag_part.data, real_part.data)).float()
+        return magnitude, phase
+
+    def inverse(self, magnitude, phase):
+        
+        recombine_magnitude_phase = magnitude*torch.cos(phase) + 1j*magnitude*torch.sin(phase)
+        inverse_transform = torch.istft(recombine_magnitude_phase, self.filter_length, hop_length=self.hop_len, win_length=self.win_len, window=self.window.to(magnitude.device)).unsqueeze(1)  # , length=self.num_samples
+        padding = self.win_len - (self.num_samples % self.win_len)
+        inverse_transform = inverse_transform[:, :, :-padding]
+        return inverse_transform
+