Update app.py

app.py

@@ -4,16 +4,12 @@ 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
 import tempfile
 import os
 import numpy as np
 
-# Global variable to store the temporary file path
-temp_file_path = None
-
 # Check for CUDA availability and set device
 try:
     import jax.tools.colab_tpu
@@ -25,41 +21,43 @@ except:
 # Load the DAC model with padding set to False for chunking
 model, variables = dac_jax.load_model(model_type="44khz", padding=False)
 
-# Jit-compile the chunk processing functions for efficiency
+# GPU-accelerated and jit-compiled chunk processing functions
+@spaces.GPU
 @jax.jit
 def compress_chunk(x):
     return model.apply(variables, x, method='compress_chunk')
 
+@spaces.GPU
 @jax.jit
 def decompress_chunk(c):
     return model.apply(variables, c, method='decompress_chunk')
 
+def ensure_mono(audio, sr):
+    if audio.ndim > 1:
+        return librosa.to_mono(audio.T), sr
+    return audio, sr
+
 @spaces.GPU
 def encode(audio_file_path):
-    global temp_file_path
     try:
-        # Load a mono audio file directly from the file path
-        signal, sample_rate = librosa.load(audio_file_path, sr=44100, mono=True) 
+        # Load and ensure mono audio
+        signal, sample_rate = librosa.load(audio_file_path, sr=44100)
+        signal, sample_rate = ensure_mono(signal, sample_rate)
 
         signal = jnp.array(signal, dtype=jnp.float32)
-        while signal.ndim < 3:
-            signal = jnp.expand_dims(signal, axis=0)
+        signal = jnp.expand_dims(signal, axis=(0, 1))  # Add batch and channel dimensions
 
         # Set chunk duration based on available GPU memory (adjust as needed)
-        win_duration = 10.0
+        win_duration = 1.0
 
         # Compress using chunking
         dac_file = model.compress(compress_chunk, signal, sample_rate, win_duration=win_duration)
 
-        # Save the compressed DAC file to a temporary file
-        if temp_file_path:
-            os.remove(temp_file_path)  # Remove the previous temporary file if it exists
-        
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".dac") as temp_file:
-            dac_file.save(temp_file.name)
-            temp_file_path = temp_file.name
+        # Save the compressed DAC file to a file in the current directory
+        output_path = "compressed_audio.dac"
+        dac_file.save(output_path)
         
-        return temp_file_path
+        return output_path
 
     except Exception as e:
         gr.Warning(f"An error occurred during encoding: {e}")
@@ -83,13 +81,6 @@ def decode(compressed_dac_file):
         gr.Warning(f"An error occurred during decoding: {e}")
         return None
 
-def cleanup(audio_file_path):
-    global temp_file_path
-    if temp_file_path and os.path.exists(temp_file_path):
-        os.remove(temp_file_path)
-        temp_file_path = None
-    return None
-
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("<h1 style='text-align: center;'>Audio Compression with DAC-JAX</h1>")
@@ -102,7 +93,6 @@ with gr.Blocks() as demo:
             encoded_output = gr.File(label="Compressed Audio (.dac)")
 
         encode_button.click(encode, inputs=audio_input, outputs=encoded_output)
-        encoded_output.change(cleanup, inputs=[audio_input], outputs=None)
 
     with gr.Tab("Decode"):
         with gr.Row():