Add exception handling to YuEGP generation loop

- Add try-catch around generation loop to capture actual errors
- Provide detailed error messages and stack traces for debugging
- Handle partial output recovery when possible
- Prevent undefined raw_output variable error

YuEGP/inference/infer.py

@@ -236,10 +236,11 @@ end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
 # Format text prompt
 run_n_segments = min(args.run_n_segments + 1, len(lyrics))
 for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
-    section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
-    guidance_scale = 1.5 if i <= 1 else 1.2
-    if i == 0:
-        continue
+    try:
+        section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
+        guidance_scale = 1.5 if i <= 1 else 1.2
+        if i == 0:
+            continue
     if i == 1:
         if args.use_dual_tracks_prompt or args.use_audio_prompt:
             if args.use_dual_tracks_prompt:
@@ -298,10 +299,21 @@ for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
         if output_seq[0][-1].item() != mmtokenizer.eoa:
             tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
             output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
-    if i > 1:
-        raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1)
-    else:
-        raw_output = output_seq
+        if i > 1:
+            raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1)
+        else:
+            raw_output = output_seq
+    except Exception as e:
+        print(f"❌ Error during generation for segment {i}: {e}")
+        print(f"Section text: {section_text}")
+        import traceback
+        traceback.print_exc()
+        # If this is the first successful generation, we still have something
+        if i > 1 and 'raw_output' in locals():
+            print("⚠️ Using partial output from previous segments")
+            break
+        else:
+            raise RuntimeError(f"Generation failed at segment {i}: {e}")
 
 # save raw output and check sanity
 ids = raw_output[0].cpu().numpy()