Update app.py

app.py

@@ -1,52 +1,62 @@
 import gradio as gr
-import spaces
-import torch
-import torchaudio
-from encodec import EncodecModel
-from encodec.utils import convert_audio
-from encodec.compress import compress_to_file, decompress_from_file
+import jax.numpy as jnp
+import librosa
+import dac_jax
+from dac_jax.audio_utils import volume_norm, db2linear
 import io
+import soundfile as sf 
 
-# Check for CUDA availability and set device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+# Load the DAC model
+model, variables = dac_jax.load_model(model_type="44khz")
+model = model.bind(variables)
 
-# Load the Encodec model and move it to the selected device
-model = EncodecModel.encodec_model_48khz().to(device)
-model.set_target_bandwidth(6.0)
-
-@spaces.GPU  # Indicate GPU usage for Spaces environment (if applicable)
+@spaces.GPU
 def encode(audio_file_path):
     try:
-        # Load and pre-process the audio waveform
-        wav, sr = torchaudio.load(audio_file_path)
+        # Load a mono audio file
+        signal, sample_rate = librosa.load(audio_file_path, sr=44100, mono=True) 
+
+        signal = jnp.array(signal, dtype=jnp.float32)
+        while signal.ndim < 3:
+            signal = jnp.expand_dims(signal, axis=0)
 
-        # Convert to mono if necessary
-        if wav.shape[0] > 1:
-            wav = torch.mean(wav, dim=0, keepdim=True)
+        target_db = -16  # Normalize audio to -16 dB
+        x, input_db = volume_norm(signal, target_db, sample_rate)
 
-        wav = convert_audio(wav, sr, model.sample_rate, model.channels)
-        wav = wav.to(device)  # Move the input audio to the selected device
+        # Encode audio signal
+        x = model.preprocess(x, sample_rate)
+        z, codes, latents, commitment_loss, codebook_loss = model.encode(x, train=False)
 
-        # Compress to ecdc file in memory
+        # Save the encoded data (codes and latents)
         output = io.BytesIO()
-        compress_to_file(model, wav, output) 
+        torch.save({'codes': codes, 'latents': latents, 'input_db': input_db, 'target_db': target_db}, output)
         output.seek(0)
 
-        return output
+        return output 
 
     except Exception as e:
         gr.Warning(f"An error occurred during encoding: {e}")
         return None
 
-@spaces.GPU 
-def decode(compressed_audio_file):
+@spaces.GPU
+def decode(encoded_data_file):
     try:
-        # Decompress audio
-        wav, sr = decompress_from_file(compressed_audio_file, device=device) # Pass the device to decompress_from_file
+        # Load the encoded data
+        encoded_data = torch.load(encoded_data_file)
+        codes = encoded_data['codes']
+        latents = encoded_data['latents']
+        input_db = encoded_data['input_db']
+        target_db = encoded_data['target_db']
+
+        # Decode audio signal
+        z = model.quantizer.decode(codes, latents)
+        y = model.decode(z)
+
+        # Undo previous loudness normalization
+        y = y * db2linear(input_db - target_db)
 
-        # Convert the decoded audio to a numpy array for Gradio output
-        decoded_audio = wav.cpu().numpy() 
+        # Convert to numpy array and squeeze to remove extra dimensions
+        decoded_audio = np.array(y).squeeze()
 
         return decoded_audio
 
@@ -56,24 +66,24 @@ def decode(compressed_audio_file):
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("<h1 style='text-align: center;'>Audio Compression with Encodec</h1>")
+    gr.Markdown("<h1 style='text-align: center;'>Audio Compression with DAC-JAX</h1>")
 
     with gr.Tab("Encode"):
         with gr.Row():
             audio_input = gr.Audio(type="filepath", label="Input Audio")
             encode_button = gr.Button("Encode", variant="primary")
         with gr.Row():
-            encoded_output = gr.File(label="Compressed Audio (.ecdc)")
+            encoded_output = gr.File(label="Encoded Data") 
 
         encode_button.click(encode, inputs=audio_input, outputs=encoded_output)
 
     with gr.Tab("Decode"):
         with gr.Row():
-            compressed_input = gr.File(label="Compressed Audio (.ecdc)")
+            encoded_input = gr.File(label="Encoded Data")
             decode_button = gr.Button("Decode", variant="primary")
         with gr.Row():
             decoded_output = gr.Audio(label="Decompressed Audio")
 
-        decode_button.click(decode, inputs=compressed_input, outputs=decoded_output)
+        decode_button.click(decode, inputs=encoded_input, outputs=decoded_output)
 
 demo.queue().launch()