add argument --repetition_penalty, default value 1.1

README.md

@@ -119,7 +119,8 @@ python infer.py \
     --run_n_segments 2 \
     --stage2_batch_size 4 \
     --output_dir ../output \
-    --max_new_tokens 3000
+    --max_new_tokens 3000 \
+    --repetition_penalty 1.1
 ```
 
 We also support music in-context-learning (provide a reference song), there are 2 types: single-track (mix/vocal/instrumental) and dual-track. 
@@ -152,6 +153,7 @@ python infer.py \
     --stage2_batch_size 4 \
     --output_dir ../output \
     --max_new_tokens 3000 \
+    --repetition_penalty 1.1 \
     --use_dual_tracks_prompt \
     --vocal_track_prompt_path ../prompt_egs/pop.00001.Vocals.mp3 \
     --instrumental_track_prompt_path ../prompt_egs/pop.00001.Instrumental.mp3 \
@@ -175,6 +177,7 @@ python infer.py \
     --stage2_batch_size 4 \
     --output_dir ../output \
     --max_new_tokens 3000 \
+    --repetition_penalty 1.1 \
     --use_audio_prompt \
     --audio_prompt_path ../prompt_egs/pop.00001.mp3 \
     --prompt_start_time 0 \

inference/codecmanipulator.py

@@ -132,7 +132,8 @@ class CodecManipulator(object):
         return einops.rearrange(x, 'K T -> (T K)')
 
     def unflatten(self, x, n_quantizer=None):
-        x = x.squeeze()
+        if x.ndim > 1 and x.shape[0] == 1:
+            x = x.squeeze(0)
         assert len(x.shape) == 1
         assert x.shape[0] % self.num_codebooks == 0 or x.shape[0] % self.n_quantizer == 0, \
             f"x.shape[0]={x.shape[0]}, num_codebooks={self.num_codebooks}, n_quantizer={self.n_quantizer}"

inference/infer.py

@@ -29,6 +29,7 @@ parser = argparse.ArgumentParser()
 parser.add_argument("--stage1_model", type=str, default="m-a-p/YuE-s1-7B-anneal-en-cot", help="The model checkpoint path or identifier for the Stage 1 model.")
 parser.add_argument("--stage2_model", type=str, default="m-a-p/YuE-s2-1B-general", help="The model checkpoint path or identifier for the Stage 2 model.")
 parser.add_argument("--max_new_tokens", type=int, default=3000, help="The maximum number of new tokens to generate in one pass during text generation.")
+parser.add_argument("--repetition_penalty", type=float, default=1.1, help="repetition_penalty ranges from 1.0 to 2.0 (or higher in some cases). It controls the diversity and coherence of the audio tokens generated. The higher the value, the greater the discouragement of repetition. Setting value to 1.0 means no penalty.")
 parser.add_argument("--run_n_segments", type=int, default=2, help="The number of segments to process during the generation.")
 parser.add_argument("--stage2_batch_size", type=int, default=4, help="The batch size used in Stage 2 inference.")
 # Prompt
@@ -155,7 +156,7 @@ output_seq = None
 # Here is suggested decoding config
 top_p = 0.93
 temperature = 1.0
-repetition_penalty = 1.0
+repetition_penalty = args.repetition_penalty
 # special tokens
 start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
 end_of_segment = mmtokenizer.tokenize('[end_of_segment]')