Update app.py

app.py

@@ -66,11 +66,12 @@ def decode_audio(encoded_file_path):
         compressed_data = temp_file.read()
         tokens_numpy_bytes = lz4.frame.decompress(compressed_data)
         tokens_numpy = np.frombuffer(tokens_numpy_bytes, dtype=np.int64)  # Ensure dtype matches encoder's output
-        tokens = torch.from_numpy(tokens_numpy).to(torch_device)
 
-    # Ensure tokens has the right dimensions
-    if tokens.ndimension() == 2:  # If tokens have only 2 dimensions
-        tokens = tokens.unsqueeze(0)  # Add batch dimension
+        # Check if tokens are 1D and reshape to 3D
+        if tokens_numpy.ndim == 1:
+            tokens_numpy = tokens_numpy.reshape(1, -1, 1)  # Reshape to [batch_size, token_length, 1]
+
+        tokens = torch.from_numpy(tokens_numpy).to(torch_device)
 
     # Debugging prints to check tensor shapes
     print(f"Tokens shape: {tokens.shape}, dtype: {tokens.dtype}")
@@ -94,6 +95,11 @@ async def stream_decode_audio(encoded_file_path) -> Generator[tuple, None, None]
             compressed_data = temp_file.read()
             tokens_numpy_bytes = lz4.frame.decompress(compressed_data)
             tokens_numpy = np.frombuffer(tokens_numpy_bytes, dtype=np.int64)  # Ensure dtype matches encoder's output
+
+            # Check if tokens are 1D and reshape to 3D
+            if tokens_numpy.ndim == 1:
+                tokens_numpy = tokens_numpy.reshape(1, -1, 1)  # Reshape to [batch_size, token_length, 1]
+
             tokens = torch.from_numpy(tokens_numpy).to(torch_device)
 
         # Ensure tokens has the right dimensions