Update app.py

app.py

@@ -67,8 +67,6 @@ import copy
 from collections import Counter
 from models.soundstream_hubert_new import SoundStream
 
-import threading
-
 device = "cuda:0"
 
 # Load model and tokenizer outside the generation function (load once)
@@ -171,12 +169,11 @@ def generate_music(
 
         print(list(enumerate(tqdm(prompt_texts[:run_n_segments]))))
 
-        # Helper function to process each segment
-        def process_segment(i, p, raw_output):
+        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  # Guidance scale adjusted based on segment index
             if i == 0:
-                return raw_output
+                continue
             if i == 1:
                 if use_audio_prompt:
                     audio_prompt = load_audio_mono(audio_prompt_path)
@@ -230,28 +227,10 @@ def generate_music(
                     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)
+                raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1)
             else:
-              raw_output = output_seq
+                raw_output = output_seq
             print(len(raw_output))
-            return raw_output
-        
-        # Create threads for each segment
-        threads = []
-        segment_outputs = {}
-
-        for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
-           thread = threading.Thread(target=lambda i=i, p=p: segment_outputs.update({i:process_segment(i,p, raw_output)}))
-           threads.append(thread)
-           thread.start()
-        
-        for thread in threads:
-          thread.join()
-        
-        
-        raw_output = segment_outputs[0]
-        for i in range(1,len(segment_outputs)):
-          raw_output = segment_outputs[i]
 
         # save raw output and check sanity
         ids = raw_output[0].cpu().numpy()