From 57568a511c4ac9993eb9eb26919efcd866dffd86 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Mon, 24 Feb 2025 11:31:43 +0800
Subject: [PATCH] init

---
 .../descriptaudiocodec/dac/compare/encodec.py | 54 +++++++++++++++++++
 1 file changed, 54 insertions(+)
 create mode 100644 inference/xcodec_mini_infer/descriptaudiocodec/dac/compare/encodec.py

diff --git a/inference/xcodec_mini_infer/descriptaudiocodec/dac/compare/encodec.py b/inference/xcodec_mini_infer/descriptaudiocodec/dac/compare/encodec.py
new file mode 100644
index 0000000..42877de
--- /dev/null
+++ b/inference/xcodec_mini_infer/descriptaudiocodec/dac/compare/encodec.py
@@ -0,0 +1,54 @@
+import torch
+from audiotools import AudioSignal
+from audiotools.ml import BaseModel
+from encodec import EncodecModel
+
+
+class Encodec(BaseModel):
+    def __init__(self, sample_rate: int = 24000, bandwidth: float = 24.0):
+        super().__init__()
+
+        if sample_rate == 24000:
+            self.model = EncodecModel.encodec_model_24khz()
+        else:
+            self.model = EncodecModel.encodec_model_48khz()
+        self.model.set_target_bandwidth(bandwidth)
+        self.sample_rate = 44100
+
+    def forward(
+        self,
+        audio_data: torch.Tensor,
+        sample_rate: int = 44100,
+        n_quantizers: int = None,
+    ):
+        signal = AudioSignal(audio_data, sample_rate)
+        signal.resample(self.model.sample_rate)
+        recons = self.model(signal.audio_data)
+        recons = AudioSignal(recons, self.model.sample_rate)
+        recons.resample(sample_rate)
+        return {"audio": recons.audio_data}
+
+
+if __name__ == "__main__":
+    import numpy as np
+    from functools import partial
+
+    model = Encodec()
+
+    for n, m in model.named_modules():
+        o = m.extra_repr()
+        p = sum([np.prod(p.size()) for p in m.parameters()])
+        fn = lambda o, p: o + f" {p/1e6:<.3f}M params."
+        setattr(m, "extra_repr", partial(fn, o=o, p=p))
+    print(model)
+    print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()]))
+
+    length = 88200 * 2
+    x = torch.randn(1, 1, length).to(model.device)
+    x.requires_grad_(True)
+    x.retain_grad()
+
+    # Make a forward pass
+    out = model(x)["audio"]
+
+    print(x.shape, out.shape)