Update app.py

app.py

@@ -11,14 +11,6 @@ from typing import Generator
 import asyncio  # Import asyncio for cancellation
 import traceback  # Import traceback for error handling
 
-# Attempt to use GPU, fallback to CPU
-try:
-    torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"Using device: {torch_device}")
-except Exception as e:
-    print(f"Error detecting GPU. Using CPU. Error: {e}")
-    torch_device = torch.device("cpu")
-
 # Load the SemantiCodec model without specifying a device
 semanticodec = SemantiCodec(token_rate=100, semantic_vocab_size=32768)
 
@@ -103,10 +95,16 @@ def decode_audio(encoded_file_path):
             tokens_numpy_bytes = lz4.frame.decompress(compressed_data)
             tokens_numpy = np.frombuffer(tokens_numpy_bytes, dtype=np.int64).reshape(shape)
 
-        # Move the tensor to the same device as the model
+        # Create a tensor from the numpy array
         tokens = torch.from_numpy(tokens_numpy)
+        
+        # Determine the device of the model
+        model_device = next(semanticodec.parameters()).device
+        print(f"Model device: {model_device}")
+
+        # Move the tokens to the same device as the model
+        tokens = tokens.to(model_device)
         print(f"Tokens device: {tokens.device}")
-        print(f"Model device: {next(semanticodec.parameters()).device}")
 
         # Decode the audio
         with torch.no_grad():
@@ -145,10 +143,16 @@ async def stream_decode_audio(encoded_file_path) -> Generator[tuple, None, None]
             tokens_numpy_bytes = lz4.frame.decompress(compressed_data)
             tokens_numpy = np.frombuffer(tokens_numpy_bytes, dtype=np.int64).reshape(shape)
             
-        # Move the tensor to the same device as the model
+        # Create a tensor from the numpy array
         tokens = torch.from_numpy(tokens_numpy)
+        
+        # Determine the device of the model
+        model_device = next(semanticodec.parameters()).device
+        print(f"Model device: {model_device}")
+
+        # Move the tokens to the same device as the model
+        tokens = tokens.to(model_device)
         print(f"Streaming tokens device: {tokens.device}")
-        print(f"Model device: {next(semanticodec.parameters()).device}")
 
         # Decode the audio in chunks
         chunk_size = sample_rate * 2  # Adjust chunk size as needed