Update app.py

app.py

@@ -11,9 +11,10 @@ from typing import Generator
 import asyncio  # Import asyncio for cancellation
 import traceback  # Import traceback for error handling
 import pickle
+import soundfile as sf
 
-# Load the SemantiCodec model without specifying a device
-semanticodec = SemantiCodec(token_rate=100, semantic_vocab_size=16384)
+# Initialize model with the specified parameters
+semanticodec = SemantiCodec(token_rate=100, semantic_vocab_size=32768)  # 1.40 kbps
 
 # Move the entire model to GPU if available
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -27,27 +28,18 @@ cancel_stream = False
 
 @spaces.GPU(duration=30)
 def encode_audio(filepath):
-    """Encode audio file and save tokens"""
+    """Encode and decode audio file"""
     try:
-        # Encode the audio
+        # Encode and decode directly as in the example
         tokens = semanticodec.encode(filepath)
+        waveform = semanticodec.decode(tokens)
         
-        # Get the original audio's sample rate
-        waveform, sample_rate = torchaudio.load(filepath)
-        
-        # Save tokens and sample rate
-        data = {
-            'tokens': tokens.cpu().detach(),
-            'sample_rate': sample_rate
-        }
-        
-        with open('encoded.pkl', 'wb') as f:
-            pickle.dump(data, f)
-            
-        return 'encoded.pkl'
+        # Save using soundfile
+        sf.write("output.wav", waveform[0,0], 16000)
+        return "output.wav"
         
     except Exception as e:
-        print(f"Encoding error: {e}")
+        print(f"Error: {e}")
         traceback.print_exc()
         return None
 
@@ -173,4 +165,8 @@ with gr.Blocks() as demo:
         stream_button.click(stream_decode_audio, inputs=input_encoded_stream, outputs=audio_output)
         cancel_stream_button.click(lambda: globals().update(cancel_stream=True), outputs=None)
 
-demo.queue().launch()
+demo.queue().launch()
+
+if __name__ == "__main__":
+    filepath = "test/test.wav"
+    encode_audio(filepath)