Update app.py

app.py

@@ -136,35 +136,41 @@ def parse_output(generated_ids):
     return code_lists[0]
 
 def redistribute_codes(code_list, snac_model):
-    device = next(snac_model.parameters()).device
-    
-    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]-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))
-        layer_3.append(code_list[7*i+5]-(5*4096))
-        layer_3.append(code_list[7*i+6]-(6*4096))
-    
-    codes = [
-        torch.tensor(layer_1, device=device).unsqueeze(0),
-        torch.tensor(layer_2, device=device).unsqueeze(0),
-        torch.tensor(layer_3, device=device).unsqueeze(0)
-    ]
-    
-    audio_hat = snac_model.decode(codes)
-    return audio_hat.detach().squeeze().cpu().numpy()
+    try:
+        device = next(snac_model.parameters()).device
+        
+        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]-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))
+            layer_3.append(code_list[7*i+5]-(5*4096))
+            layer_3.append(code_list[7*i+6]-(6*4096))
+        
+        codes = [
+            torch.tensor(layer_1, device=device).unsqueeze(0),
+            torch.tensor(layer_2, device=device).unsqueeze(0),
+            torch.tensor(layer_3, device=device).unsqueeze(0)
+        ]
+        
+        audio_hat = snac_model.decode(codes)
+        return audio_hat.detach().squeeze().cpu().numpy()
+    except Exception as e:
+        logger.error(f"Error in redistribute_codes: {e}", exc_info=True)
+        return None
 
 @spaces.GPU()
 def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new_tokens, progress=gr.Progress()):
     if not text.strip():
+        logger.warning("Empty text input. Skipping speech generation.")
         return None
     
     try:
         progress(0.1, "Processing text...")
         input_ids, attention_mask = process_prompt(text, voice, tokenizer, device)
+        logger.info(f"Input shape: {input_ids.shape}")
         
         progress(0.3, "Generating speech tokens...")
         with torch.no_grad():
@@ -179,16 +185,31 @@ def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new
                 num_return_sequences=1,
                 eos_token_id=128258,
             )
+        logger.info(f"Generated shape: {generated_ids.shape}")
         
         progress(0.6, "Processing speech tokens...")
         code_list = parse_output(generated_ids)
+        logger.info(f"Code list length: {len(code_list)}")
+        
+        if not code_list:
+            logger.warning("No valid code list generated. Skipping audio conversion.")
+            return None
         
         progress(0.8, "Converting to audio...")
         audio_samples = redistribute_codes(code_list, snac_model)
         
+        if audio_samples is None:
+            logger.warning("Audio samples is None.")
+            return None
+        
+        if len(audio_samples) == 0:
+            logger.warning("Audio samples is empty.")
+            return None
+        
+        logger.info(f"Audio samples shape: {audio_samples.shape}")
         return (24000, audio_samples)  # Return sample rate and audio
     except Exception as e:
-        print(f"Error generating speech: {e}")
+        logger.error(f"Error generating speech: {e}", exc_info=True)
         return None
 
 @spaces.GPU()