Updated Code to run on Colab Notebook without Errors (#13)

- Updated Code to run on Colab Notebook without Errors (34eb501a93ab974b0b0ff45d6bba22669d51cea4)


Co-authored-by: Shresth Shukla <[email protected]>

README.md

@@ -152,6 +152,9 @@ def decode_snac_tokens(snac_tokens, snac_model):
     if not snac_tokens or len(snac_tokens) % 7 != 0:
         return None
 
+    # Get the device of the SNAC model. Fixed by Shresth to run on colab notebook :)
+    snac_device = next(snac_model.parameters()).device
+
     # De-interleave tokens into 3 hierarchical levels
     codes_lvl = [[] for _ in range(3)]
     llm_codebook_offsets = [AUDIO_CODE_BASE_OFFSET + i * 4096 for i in range(7)]
@@ -171,7 +174,7 @@ def decode_snac_tokens(snac_tokens, snac_model):
     # Convert to tensors for SNAC decoder
     hierarchical_codes = []
     for lvl_codes in codes_lvl:
-        tensor = torch.tensor(lvl_codes, dtype=torch.int32, device=snac_model.device).unsqueeze(0)
+        tensor = torch.tensor(lvl_codes, dtype=torch.int32, device=snac_device).unsqueeze(0)
         if torch.any((tensor < 0) | (tensor > 4095)):
             raise ValueError("Invalid SNAC token values")
         hierarchical_codes.append(tensor)