Update app.py

app.py

@@ -69,7 +69,11 @@ def decode_audio(encoded_file_path):
         tokens = torch.from_numpy(tokens_numpy).to(torch_device)
 
     # Ensure tokens has the right dimensions
-    tokens = tokens.unsqueeze(0) if tokens.ndimension() == 1 else tokens
+    if tokens.ndimension() == 2:  # If tokens have only 2 dimensions
+        tokens = tokens.unsqueeze(0)  # Add batch dimension
+
+    # Debugging prints to check tensor shapes
+    print(f"Tokens shape: {tokens.shape}, dtype: {tokens.dtype}")
 
     # Decode the audio
     with torch.no_grad():
@@ -93,7 +97,8 @@ async def stream_decode_audio(encoded_file_path) -> Generator[tuple, None, None]
             tokens = torch.from_numpy(tokens_numpy).to(torch_device)
 
         # Ensure tokens has the right dimensions
-        tokens = tokens.unsqueeze(0) if tokens.ndimension() == 1 else tokens
+        if tokens.ndimension() == 2:  # If tokens have only 2 dimensions
+            tokens = tokens.unsqueeze(0)  # Add batch dimension
 
         # Decode the audio in chunks
         chunk_size = sample_rate  # Use the stored sample rate as chunk size