Update app.py

app.py

@@ -104,10 +104,14 @@ def parse_output(generated_ids):
     return code_lists[0]
 
 def redistribute_codes(code_list, snac_model):
-    layer_1, layer_2, layer_3 = [], [], []
-    for i in range(len(code_list) // 7):
+    device = next(snac_model.parameters()).device  # Get the device of SNAC model
+    
+    layer_1 = []
+    layer_2 = []
+    layer_3 = []
+    for i in range((len(code_list)+1)//7):
         layer_1.append(code_list[7*i])
-        layer_2.append(code_list[7*i+1]-(1*4096))
+        layer_2.append(code_list[7*i+1]-4096)
         layer_3.append(code_list[7*i+2]-(2*4096))
         layer_3.append(code_list[7*i+3]-(3*4096))
         layer_2.append(code_list[7*i+4]-(4*4096))
@@ -120,8 +124,8 @@ def redistribute_codes(code_list, snac_model):
         torch.tensor(layer_3, device=device).unsqueeze(0)
     ]
     
-    audio = snac_model.decode(codes)
-    return audio.cpu().numpy().flatten()
+    audio_hat = snac_model.decode(codes)
+    return audio_hat.detach().squeeze().cpu().numpy()  # Always return CPU numpy array
 
 @spaces.GPU()
 def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new_tokens, progress=gr.Progress()):