1.modify split_lyrics regex 2.define raw_output outside loop

YuE/inference/infer.py

@@ -131,7 +131,7 @@ def encode_audio(codec_model, audio_prompt, device, target_bw=0.5):
     return raw_codes
 
 def split_lyrics(lyrics):
-    pattern = r"\[(\w+)\](.*?)(?=\[|\Z)"
+    pattern = r"\[([^]]+)\](.*?)(?=\[|\Z)"
     segments = re.findall(pattern, lyrics, re.DOTALL)
     structured_lyrics = [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
     return structured_lyrics
@@ -162,6 +162,7 @@ start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
 end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
 # Format text prompt
 run_n_segments = min(args.run_n_segments+1, len(lyrics))
+raw_output = None
 for i, p in enumerate(tqdm(prompt_texts[:run_n_segments], desc="Stage1 inference...")):
     section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
     guidance_scale = 1.5 if i <=1 else 1.2
@@ -224,6 +225,8 @@ for i, p in enumerate(tqdm(prompt_texts[:run_n_segments], desc="Stage1 inference
         raw_output = output_seq
 
 # save raw output and check sanity
+if raw_output is None:
+    raise ValueError("No valid segments were processed. Check your lyrics format and run_n_segments parameter.")
 ids = raw_output[0].cpu().numpy()
 soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
 eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()