Update app.py

app.py

@@ -1,15 +1,24 @@
 import gradio as gr
-import spaces
+import jax
 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 
+import spaces
+
+# Load the DAC model with padding set to False for chunking
+model, variables = dac_jax.load_model(model_type="44khz", padding=False)
 
-# Load the DAC model
-model, variables = dac_jax.load_model(model_type="44khz")
-model = model.bind(variables)
+# Jit-compile the chunk processing functions for efficiency
+@jax.jit
+def compress_chunk(x):
+    return model.apply(variables, x, method='compress_chunk')
+
+@jax.jit
+def decompress_chunk(c):
+    return model.apply(variables, c, method='decompress_chunk')
 
 @spaces.GPU
 def encode(audio_file_path):
@@ -21,40 +30,31 @@ def encode(audio_file_path):
         while signal.ndim < 3:
             signal = jnp.expand_dims(signal, axis=0)
 
-        target_db = -16  # Normalize audio to -16 dB
-        x, input_db = volume_norm(signal, target_db, sample_rate)
+        # Set chunk duration based on available GPU memory (adjust as needed)
+        win_duration = 0.5  # You might need to experiment with this value
 
-        # Encode audio signal
-        x = model.preprocess(x, sample_rate)
-        z, codes, latents, commitment_loss, codebook_loss = model.encode(x, train=False)
+        # Compress using chunking
+        dac_file = model.compress(compress_chunk, signal, sample_rate, win_duration=win_duration)
 
-        # Save the encoded data (codes and latents)
+        # Save the compressed DAC file to BytesIO
         output = io.BytesIO()
-        torch.save({'codes': codes, 'latents': latents, 'input_db': input_db, 'target_db': target_db}, output)
+        dac_file.save(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(encoded_data_file):
+def decode(compressed_dac_file):
     try:
-        # 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)
+        # Load the compressed DAC file
+        dac_file = dac_jax.DACFile.load(compressed_dac_file)
 
-        # Undo previous loudness normalization
-        y = y * db2linear(input_db - target_db)
+        # Decompress using chunking
+        y = model.decompress(decompress_chunk, dac_file)
 
         # Convert to numpy array and squeeze to remove extra dimensions
         decoded_audio = np.array(y).squeeze()
@@ -74,17 +74,17 @@ with gr.Blocks() as demo:
             audio_input = gr.Audio(type="filepath", label="Input Audio")
             encode_button = gr.Button("Encode", variant="primary")
         with gr.Row():
-            encoded_output = gr.File(label="Encoded Data") 
+            encoded_output = gr.File(label="Compressed Audio (.dac)")
 
         encode_button.click(encode, inputs=audio_input, outputs=encoded_output)
 
     with gr.Tab("Decode"):
         with gr.Row():
-            encoded_input = gr.File(label="Encoded Data")
+            compressed_input = gr.File(label="Compressed Audio (.dac)")
             decode_button = gr.Button("Decode", variant="primary")
         with gr.Row():
             decoded_output = gr.Audio(label="Decompressed Audio")
 
-        decode_button.click(decode, inputs=encoded_input, outputs=decoded_output)
+        decode_button.click(decode, inputs=compressed_input, outputs=decoded_output)
 
 demo.queue().launch()