Update diffrhythm2/utils.py

diffrhythm2/utils.py

@@ -1,6 +1,7 @@
 import torch
 import torchaudio
 import os
+import re
 import json
 import random
 import io
@@ -103,17 +104,31 @@ def prepare_model(repo_id, device, dtype):
 
     return diffrhythm2, mulan, lrc_tokenizer, decoder
 
+STRUCT_PATTERN = re.compile(r'^\[.*?\]$')
+    
 def parse_lyrics(lrc_tokenizer, lyrics: str):
     lyrics_with_time = []
     lyrics = lyrics.split("\n")
+    get_start = False
     for line in lyrics:
-        struct_idx = STRUCT_INFO.get(line, None)
-        if struct_idx is not None:
-            lyrics_with_time.append([struct_idx, STRUCT_INFO['[stop]']])
+        line = line.strip()
+        if not line:
+            continue
+        struct_flag = STRUCT_PATTERN.match(line)
+        if struct_flag:
+            struct_idx = STRUCT_INFO.get(line.lower(), None)
+            if struct_idx is not None:
+                if struct_idx == STRUCT_INFO['[start]']:
+                    get_start = True
+                lyrics_with_time.append([struct_idx, STRUCT_INFO['[stop]']])
+            else:
+                continue
         else:
             tokens = lrc_tokenizer.encode(line.strip())
             tokens = tokens + [STRUCT_INFO['[stop]']]
             lyrics_with_time.append(tokens)
+    if len(lyrics_with_time) != 0 and not get_start:
+        lyrics_with_time = [STRUCT_INFO['[start]'], STRUCT_INFO['[stop]']] + lyrics_with_time
     return lyrics_with_time
 
 @torch.no_grad()