Update app.py

app.py

@@ -3,29 +3,62 @@ import torch
 import numpy as np
 import soundfile as sf
 from scipy.signal import resample
+from dataclasses import dataclass
+from huggingface_hub import hf_hub_download
 
-# import your codec
-from irodori_tts.codec import DACVAECodec
+
+# =============================
+# SIMPLE DACVAE WRAPPER
+# =============================
+@dataclass
+class SimpleDACCodec:
+    model: torch.nn.Module
+    sample_rate: int
+    device: torch.device
+
+    @classmethod
+    def load(cls, repo_id="Aratako/Semantic-DACVAE-Japanese-32dim", device="cpu"):
+        # lazy import (no local repo needed)
+        from dacvae import DACVAE
+
+        # download weights
+        weights_path = hf_hub_download(repo_id=repo_id, filename="weights.pth")
+
+        model = DACVAE.load(weights_path).eval().to(device)
+
+        return cls(
+            model=model,
+            sample_rate=int(model.sample_rate),
+            device=torch.device(device),
+        )
+
+    @torch.inference_mode()
+    def encode(self, audio):
+        # audio: (1,1,T)
+        z = self.model.encode(audio)  # (B, D, T)
+        return z.transpose(1, 2)      # (B, T, D)
+
+    @torch.inference_mode()
+    def decode(self, latent):
+        # latent: (B, T, D)
+        z = latent.transpose(1, 2)
+        return self.model.decode(z)
 
 
 # =============================
-# LOAD MODEL
+# INIT
 # =============================
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-codec = DACVAECodec.load(
-    repo_id="Aratako/Semantic-DACVAE-Japanese-32dim",
-    device=DEVICE,
-)
+codec = SimpleDACCodec.load(device=DEVICE)
 
 
 # =============================
-# AUDIO UTILS (NO TORCHAUDIO)
+# AUDIO UTILS (soundfile only)
 # =============================
 def load_audio(path):
     audio, sr = sf.read(path, dtype="float32")
 
-    # convert to mono
+    # mono
     if audio.ndim > 1:
         audio = np.mean(audio, axis=1)
 
@@ -41,7 +74,7 @@ def resample_audio(audio, orig_sr, target_sr):
 
 
 def to_tensor(audio):
-    return torch.from_numpy(audio).unsqueeze(0).unsqueeze(0)  # (1,1,T)
+    return torch.from_numpy(audio).unsqueeze(0).unsqueeze(0)
 
 
 # =============================
@@ -50,26 +83,24 @@ def to_tensor(audio):
 def encode_audio(file):
     audio, sr = load_audio(file)
 
-    # resample
     audio = resample_audio(audio, sr, codec.sample_rate)
-
     wav = to_tensor(audio).to(DEVICE)
 
-    latent = codec.encode_waveform(wav, codec.sample_rate)
+    latent = codec.encode(wav)
 
-    return latent.cpu().numpy()
+    return latent.cpu().numpy().tolist()
 
 
 # =============================
 # DECODE
 # =============================
-def decode_audio(latent_np):
-    latent = torch.tensor(latent_np).to(DEVICE)
+def decode_audio(latent_list):
+    latent = torch.tensor(latent_list, dtype=torch.float32).to(DEVICE)
 
     if latent.ndim == 2:
         latent = latent.unsqueeze(0)
 
-    audio = codec.decode_latent(latent)
+    audio = codec.decode(latent)
 
     audio = audio.squeeze().cpu().numpy()
 
@@ -77,28 +108,30 @@ def decode_audio(latent_np):
 
 
 # =============================
-# GRADIO UI
+# UI
 # =============================
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎧 DACVAE Audio Codec (SoundFile Version)")
+    gr.Markdown("## 🎧 Simple DAC Audio Codec (No torchaudio)")
 
     with gr.Tab("Encode"):
         audio_in = gr.Audio(type="filepath")
-        latent_out = gr.Textbox(label="Latent (numpy array)")
+        latent_out = gr.JSON(label="Latent")
 
-        btn_encode = gr.Button("Encode")
-        btn_encode.click(encode_audio, inputs=audio_in, outputs=latent_out)
+        gr.Button("Encode").click(
+            encode_audio,
+            inputs=audio_in,
+            outputs=latent_out
+        )
 
     with gr.Tab("Decode"):
-        latent_in = gr.Textbox(label="Paste latent numpy array")
+        latent_in = gr.JSON(label="Latent")
         audio_out = gr.Audio()
 
-        def decode_from_text(text):
-            latent = np.array(eval(text))
-            return decode_audio(latent)
-
-        btn_decode = gr.Button("Decode")
-        btn_decode.click(decode_from_text, inputs=latent_in, outputs=audio_out)
+        gr.Button("Decode").click(
+            decode_audio,
+            inputs=latent_in,
+            outputs=audio_out
+        )
 
 
 # =============================