Update app.py

app.py

@@ -6,6 +6,7 @@ from scipy.signal import resample
 from dataclasses import dataclass
 from huggingface_hub import hf_hub_download
 
+
 # =============================
 # SIMPLE DACVAE WRAPPER
 # =============================
@@ -30,11 +31,13 @@ class SimpleDACCodec:
 
     @torch.inference_mode()
     def encode(self, audio):
-        z = self.model.encode(audio)      # (B, D, T)
-        return z.transpose(1, 2)          # (B, T, D)
+        # 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)
 
@@ -52,6 +55,7 @@ codec = SimpleDACCodec.load(device=DEVICE)
 def load_audio(path):
     audio, sr = sf.read(path, dtype="float32")
 
+    # mono
     if audio.ndim > 1:
         audio = np.mean(audio, axis=1)
 
@@ -75,63 +79,72 @@ def to_tensor(audio):
 # =============================
 def encode_audio(file):
     if file is None:
-        raise gr.Error("Please upload an audio file first.")
+        raise ValueError("Please upload an audio file first.")
 
     audio, sr = load_audio(file)
     audio = resample_audio(audio, sr, codec.sample_rate)
-
     wav = to_tensor(audio).to(DEVICE)
-    latent = codec.encode(wav)
 
-    # keep as numpy (NOT list → avoids huge lag)
-    return latent.cpu().numpy()
+    latent = codec.encode(wav)  # (B, T, D)
+
+    latent_list = latent.detach().cpu().numpy().tolist()
+    return latent_list, latent_list  # one for display, one for hidden state
 
 
 # =============================
 # DECODE
 # =============================
-def decode_audio(latent):
-    if latent is None:
-        raise gr.Error("No latent available. Click Encode first.")
+def decode_audio(latent_list):
+    if latent_list is None:
+        raise ValueError("No latent found. Click Encode first.")
 
-    latent = torch.tensor(latent, dtype=torch.float32).to(DEVICE)
+    # Convert nested list to tensor safely
+    try:
+        latent = torch.tensor(latent_list, dtype=torch.float32, device=DEVICE)
+    except Exception as e:
+        raise ValueError(f"Invalid latent data: {e}")
 
     if latent.ndim == 2:
         latent = latent.unsqueeze(0)
 
     audio = codec.decode(latent)
-    audio = audio.squeeze().cpu().numpy()
+    audio = audio.squeeze().detach().cpu().numpy()
+
+    # clip just in case
+    audio = np.nan_to_num(audio)
+    audio = np.clip(audio, -1.0, 1.0)
 
     return (codec.sample_rate, audio)
 
 
 # =============================
-# UI (SINGLE WINDOW)
+# UI
 # =============================
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎧 Simple DAC Audio Codec")
-
-    audio_in = gr.Audio(type="filepath", label="Upload Audio")
+    gr.Markdown("## 🎧 Simple DAC Audio Codec (Single Window)")
 
-    encode_btn = gr.Button("Encode")
-    decode_btn = gr.Button("Decode")
+    latent_state = gr.State()
 
-    latent_state = gr.State()  # 🔥 hidden storage (best practice)
+    with gr.Row():
+        with gr.Column(scale=1):
+            audio_in = gr.Audio(type="filepath", label="Upload Audio")
+            encode_btn = gr.Button("Encode")
+            decode_btn = gr.Button("Decode")
 
-    audio_out = gr.Audio(label="Reconstructed Audio")
+        with gr.Column(scale=1):
+            latent_out = gr.JSON(label="Latent")
+            audio_out = gr.Audio(label="Reconstructed Audio")
 
-    # Encode → store in state
     encode_btn.click(
         fn=encode_audio,
         inputs=audio_in,
-        outputs=latent_state
+        outputs=[latent_out, latent_state],
     )
 
-    # Decode from state
     decode_btn.click(
         fn=decode_audio,
         inputs=latent_state,
-        outputs=audio_out
+        outputs=audio_out,
     )