Upload codec/audio_processing/dune_codec.py with huggingface_hub

codec/audio_processing/dune_codec.py

@@ -275,6 +275,30 @@ class DuneAudioTokenizer(nn.Module):
 
         return tokens[0]
 
+    def _post_filter(self, audio):
+        """Spectral post-filter over the reconstructed waveform.
+
+        Applied per item at the output rate. A failure here must not cost the
+        caller their audio, so it degrades to the unfiltered signal.
+        """
+        try:
+            from ._postfilter import get_post_filter
+
+            pf = get_post_filter(device="cpu")
+        except Exception:
+            return audio
+
+        out = np.array(audio, dtype=np.float32, copy=True)
+        flat = out.reshape(-1, out.shape[-1]) if out.ndim > 1 else out[None]
+        for i in range(flat.shape[0]):
+            try:
+                filtered = pf(flat[i], self.output_sample_rate)
+            except Exception:
+                continue
+            n = min(filtered.size, flat.shape[1])
+            flat[i, :n] = filtered[:n]
+        return flat.reshape(out.shape) if out.ndim > 1 else flat[0]
+
     def decode(self, vq_code):
         tokens = vq_code if vq_code.dim() == 3 else vq_code.unsqueeze(0)
         tokens = tokens.to(self.device)
@@ -290,7 +314,7 @@ class DuneAudioTokenizer(nn.Module):
             else:
                 audio, _ = self.codec.decode(tokens=tokens, tokens_len=tokens_len)
 
-        return audio.cpu().numpy()
+        return self._post_filter(audio.cpu().numpy())
 
 
 def _state_dict_from(ckpt):