Update app.py

app.py

@@ -127,6 +127,14 @@ def generate_music(genre_txt, lyrics_txt, max_new_tokens=5, run_n_segments=2, us
     raw_output = None
     stage1_output_set = []
 
+    class BlockTokenRangeProcessor(LogitsProcessor):
+        def __init__(self, start_id, end_id):
+            self.blocked_token_ids = list(range(start_id, end_id))
+
+        def __call__(self, input_ids, scores):
+            scores[:, self.blocked_token_ids] = -float("inf")
+            return scores
+    
     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