refactor

.gitignore

@@ -161,5 +161,6 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-inference/output
+output
 inference/xcodec_mini_infer
+inference/run_infer.sh

README.md

@@ -107,15 +107,15 @@ Note:
 # This is the CoT mode.
 cd YuE/inference/
 python infer.py \
+    --cuda_idx 0 \
     --stage1_model m-a-p/YuE-s1-7B-anneal-en-cot \
     --stage2_model m-a-p/YuE-s2-1B-general \
-    --genre_txt ../genre.txt \
-    --lyrics_txt ../lyrics.txt \
+    --genre_txt ../prompt_egs/genre.txt \
+    --lyrics_txt ../prompt_egs/lyrics.txt \
     --run_n_segments 2 \
-    --stage2_batch_size 4 \
-    --output_dir ./output \
-    --cuda_idx 0 \
-    --max_new_tokens 3000 \
+    --stage2_batch_size 16 \
+    --output_dir ../output \
+    --max_new_tokens 3000
 ```
 
 We also support music in-context-learning (provide a reference song), there are 2 types: single-track (mix/vocal/instrumental) and dual-track. 
@@ -139,18 +139,19 @@ Note:
 # The ref audio is taken from GTZAN test set.
 cd YuE/inference/
 python infer.py \
+    --cuda_idx 0 \
     --stage1_model m-a-p/YuE-s1-7B-anneal-en-icl \
     --stage2_model m-a-p/YuE-s2-1B-general \
-    --genre_txt ../genre.txt \
-    --lyrics_txt ../lyrics.txt \
+    --genre_txt ../prompt_egs/genre.txt \
+    --lyrics_txt ../prompt_egs/lyrics.txt \
     --run_n_segments 2 \
-    --stage2_batch_size 4 \
-    --output_dir ./output \
-    --cuda_idx 0 \
+    --stage2_batch_size 16 \
+    --output_dir ../output \
     --max_new_tokens 3000 \
     --use_dual_tracks_prompt \
-    --vocal_track_prompt_path ../pop.00001.Vocals.mp3 \
-    --instrumental_track_prompt_path ../pop.00001.Instrumental.mp3 \
+    --seed 0 \
+    --vocal_track_prompt_path ../prompt_egs/pop.00001.Vocals.mp3 \
+    --instrumental_track_prompt_path ../prompt_egs/pop.00001.Instrumental.mp3 \
     --prompt_start_time 0 \
     --prompt_end_time 30 
 ```
@@ -162,17 +163,17 @@ python infer.py \
 # The ref audio is taken from GTZAN test set.
 cd YuE/inference/
 python infer.py \
+    --cuda_idx 0 \
     --stage1_model m-a-p/YuE-s1-7B-anneal-en-icl \
     --stage2_model m-a-p/YuE-s2-1B-general \
-    --genre_txt ../genre.txt \
-    --lyrics_txt ../lyrics.txt \
+    --genre_txt ../prompt_egs/genre.txt \
+    --lyrics_txt ../prompt_egs/lyrics.txt \
     --run_n_segments 2 \
     --stage2_batch_size 4 \
-    --output_dir ./output \
-    --cuda_idx 0 \
+    --output_dir ../output \
     --max_new_tokens 3000 \
     --use_audio_prompt \
-    --audio_prompt_path ../pop.00001.mp3 \
+    --audio_prompt_path ../prompt_egs/pop.00001.mp3 \
     --prompt_start_time 0 \
     --prompt_end_time 30 
 ```
@@ -182,7 +183,7 @@ python infer.py \
 The prompt consists of three parts: genre tags, lyrics, and ref audio.
 
 ### Genre Tagging Prompt
-1. An example genre tagging prompt can be found [here](inference/prompt_examples/genre.txt).
+1. An example genre tagging prompt can be found [here](prompt_egs/genre.txt).
 
 2. A stable tagging prompt usually consists of five components: genre, instrument, mood, gender, and timbre. All five should be included if possible, separated by space (space delimiter).
 
@@ -193,7 +194,7 @@ The prompt consists of three parts: genre tags, lyrics, and ref audio.
 4. Additionally, we have introduced the "Mandarin" and "Cantonese" tags to distinguish between Mandarin and Cantonese, as their lyrics often share similarities.
 
 ### Lyrics Prompt
-1. An example lyric prompt can be found [here](inference/prompt_examples/lyrics.txt).
+1. An example lyric prompt can be found [here](prompt_egs/lyrics.txt).
 
 2. We support multiple languages, including but not limited to English, Mandarin Chinese, Cantonese, Japanese, and Korean. The default top language distribution during the annealing phase is revealed in [issue 12](https://github.com/multimodal-art-projection/YuE/issues/12#issuecomment-2620845772). A language ID on a specific annealing checkpoint indicates that we have adjusted the mixing ratio to enhance support for that language.
 

inference/infer.py

@@ -84,11 +84,15 @@ model = AutoModelForCausalLM.from_pretrained(
     stage1_model, 
     torch_dtype=torch.bfloat16,
     attn_implementation="flash_attention_2", # To enable flashattn, you have to install flash-attn
+    # device_map="auto",
     )
 # to device, if gpu is available
 model.to(device)
 model.eval()
 
+if torch.__version__ >= "2.0.0":
+    model = torch.compile(model)
+
 codectool = CodecManipulator("xcodec", 0, 1)
 codectool_stage2 = CodecManipulator("xcodec", 0, 8)
 model_config = OmegaConf.load(args.basic_model_config)
@@ -256,12 +260,16 @@ if not args.disable_offload_model:
 print("Stage 2 inference...")
 model_stage2 = AutoModelForCausalLM.from_pretrained(
     stage2_model, 
-    torch_dtype=torch.float16,
-    attn_implementation="flash_attention_2"
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    # device_map="auto",
     )
 model_stage2.to(device)
 model_stage2.eval()
 
+if torch.__version__ >= "2.0.0":
+    model_stage2 = torch.compile(model_stage2)
+
 def stage2_generate(model, prompt, batch_size=16):
     codec_ids = codectool.unflatten(prompt, n_quantizer=1)
     codec_ids = codectool.offset_tok_ids(

inference/prompt_examples/genre.txt

@@ -1 +0,0 @@
-inspiring female uplifting pop airy vocal electronic bright vocal vocal

lyrics.txt

@@ -1,39 +0,0 @@
-[verse]
-Staring at the sunset, colors paint the sky
-Thoughts of you keep swirling, can't deny
-I know I let you down, I made mistakes
-But I'm here to mend the heart I didn't break
-
-[chorus]
-Every road you take, I'll be one step behind
-Every dream you chase, I'm reaching for the light
-You can't fight this feeling now
-I won't back down
-You know you can't deny it now
-I won't back down
-
-[verse]
-They might say I'm foolish, chasing after you
-But they don't feel this love the way we do
-My heart beats only for you, can't you see?
-I won't let you slip away from me
-
-[chorus]
-Every road you take, I'll be one step behind
-Every dream you chase, I'm reaching for the light
-You can't fight this feeling now
-I won't back down
-You know you can't deny it now
-I won't back down
-
-[bridge]
-No, I won't back down, won't turn around
-Until you're back where you belong
-I'll cross the oceans wide, stand by your side
-Together we are strong
-
-[outro]
-Every road you take, I'll be one step behind
-Every dream you chase, love's the tie that binds
-You can't fight this feeling now
-I won't back down

requirements.txt

@@ -10,3 +10,4 @@ tensorboard
 descript-audiotools>=0.7.2
 descript-audio-codec
 scipy==1.10.1
+accelerate>=0.26.0