Update app.py

app.py

@@ -3,94 +3,60 @@ import jax
 import jax.numpy as jnp
 import librosa
 import dac_jax
-from dac_jax.audio_utils import volume_norm, db2linear
-import soundfile as sf 
 import spaces
 import tempfile
 import os
 import numpy as np
 
-# Check for CUDA availability and set device
-try:
-    import jax.tools.colab_tpu
-    jax.tools.colab_tpu.setup_tpu()
-    print("Connected to TPU")
-except:
-    print("No TPU detected, using GPU or CPU.")
-
 # Load the DAC model with padding set to False for chunking
 model, variables = dac_jax.load_model(model_type="44khz", padding=False)
 
-# GPU-accelerated and jit-compiled chunk processing functions
-@spaces.GPU
+# Jit-compile these functions because they're used inside a loop over chunks.
 @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):
     try:
-        # Load and ensure mono audio
-        signal, sample_rate = librosa.load(audio_file_path, sr=44100)
-        signal, sample_rate = ensure_mono(signal, sample_rate)
+        # Load audio with librosa, specifying duration
+        signal, sample_rate = librosa.load(audio_file_path, sr=44100, mono=True, duration=.5)  # Set duration as needed
 
         signal = jnp.array(signal, dtype=jnp.float32)
-        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 = 2.0  # Reduced win_duration for testing
+        while signal.ndim < 3:
+            signal = jnp.expand_dims(signal, axis=0)
 
-        # Compress using chunking
+        win_duration = 0.5  # Adjust based on your GPU's memory size
         dac_file = model.compress(compress_chunk, signal, sample_rate, win_duration=win_duration)
 
-        # Save the compressed DAC file to a temporary file
+        # Save to a temporary file
         with tempfile.NamedTemporaryFile(delete=False, suffix=".dac") as temp_file:
             dac_file.save(temp_file.name)
 
-        # Check file size
-        file_size = os.path.getsize(temp_file.name)
-        if file_size == 0:
-            print("WARNING: Compressed file size is 0 bytes!")
-
-        # Return the temporary file path
         return temp_file.name
 
     except Exception as e:
-        print(f"ERROR during encoding: {e}")
         gr.Warning(f"An error occurred during encoding: {e}")
         return None
 
 @spaces.GPU
 def decode(compressed_dac_file):
     try:
-        # Create a temporary file to save the uploaded content
+        # Load from the uploaded file
         with tempfile.NamedTemporaryFile(delete=False, suffix=".dac") as temp_file:
             temp_file.write(compressed_dac_file.file.read())
             temp_file_path = temp_file.name
 
-        # Load the compressed DAC file
         dac_file = dac_jax.DACFile.load(temp_file_path)
-
-        # 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()
 
-        # Delete the temporary file
         os.unlink(temp_file_path)
-
-        return (44100, decoded_audio)  # Return sample rate and audio data
+        return (44100, decoded_audio)
 
     except Exception as e:
         gr.Warning(f"An error occurred during decoding: {e}")