Update app.py

app.py

@@ -56,25 +56,36 @@ from omegaconf import OmegaConf
 import torchaudio
 from torchaudio.transforms import Resample
 import soundfile as sf
+
 from tqdm import tqdm
 from einops import rearrange
 from codecmanipulator import CodecManipulator
 from mmtokenizer import _MMSentencePieceTokenizer
 from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
+import glob
+import time
+import copy
+from collections import Counter
 from models.soundstream_hubert_new import SoundStream
 from vocoder import build_codec_model, process_audio
 from post_process_audio import replace_low_freq_with_energy_matched
 
-# Initialize device
 device = "cuda:0"
 
-# Load models once and reuse
-print("Loading models...")
 model = AutoModelForCausalLM.from_pretrained(
     "m-a-p/YuE-s1-7B-anneal-en-cot",
     torch_dtype=torch.float16,
-    attn_implementation="flash_attention_2",
-).to(device).eval()
+    attn_implementation="flash_attention_2",  # To enable flashattn, you have to install flash-attn
+).to(device)
+# assistant_model = AutoModelForCausalLM.from_pretrained(
+#     "m-a-p/YuE-s2-1B-general",
+#     torch_dtype=torch.float16,
+#     attn_implementation="flash_attention_2",  # To enable flashattn, you have to install flash-attn
+# ).to(device)
+# assistant_model = torch.compile(assistant_model)
+# model = torch.compile(model)
+# assistant_model.eval()
+model.eval()
 
 basic_model_config = './xcodec_mini_infer/final_ckpt/config.yaml'
 resume_path = './xcodec_mini_infer/final_ckpt/ckpt_00360000.pth'
@@ -82,142 +93,308 @@ config_path = './xcodec_mini_infer/decoders/config.yaml'
 vocal_decoder_path = './xcodec_mini_infer/decoders/decoder_131000.pth'
 inst_decoder_path = './xcodec_mini_infer/decoders/decoder_151000.pth'
 
-# Load codec model
+mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
+
+codectool = CodecManipulator("xcodec", 0, 1)
 model_config = OmegaConf.load(basic_model_config)
+# Load codec model
 codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
-codec_model.load_state_dict(torch.load(resume_path, map_location='cpu')['codec_model'])
+parameter_dict = torch.load(resume_path, map_location='cpu')
+codec_model.load_state_dict(parameter_dict['codec_model'])
+# codec_model = torch.compile(codec_model)
 codec_model.eval()
 
 # Preload and compile vocoders
 vocal_decoder, inst_decoder = build_codec_model(config_path, vocal_decoder_path, inst_decoder_path)
-vocal_decoder.to(device).eval()
-inst_decoder.to(device).eval()
-
-# Tokenizer and codec tool
-mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-codectool = CodecManipulator("xcodec", 0, 1)
-
-def generate_music(genre_txt, lyrics_txt, max_new_tokens=5, run_n_segments=2, use_audio_prompt=False, audio_prompt_path="", prompt_start_time=0.0, prompt_end_time=30.0, rescale=False):
+vocal_decoder.to(device)
+inst_decoder.to(device)
+# vocal_decoder = torch.compile(vocal_decoder)
+# inst_decoder = torch.compile(inst_decoder)
+vocal_decoder.eval()
+inst_decoder.eval()
+
+
+def generate_music(
+        max_new_tokens=5,
+        run_n_segments=2,
+        genre_txt=None,
+        lyrics_txt=None,
+        use_audio_prompt=False,
+        audio_prompt_path="",
+        prompt_start_time=0.0,
+        prompt_end_time=30.0,
+        cuda_idx=0,
+        rescale=False,
+):
     if use_audio_prompt and not audio_prompt_path:
-        raise FileNotFoundError("Please provide an audio prompt filepath when enabling 'use_audio_prompt'!")
-
-    max_new_tokens *= 100
-    top_p = 0.93
-    temperature = 1.0
-    repetition_penalty = 1.2
-
-    # Split lyrics into segments
-    def split_lyrics(lyrics):
-        pattern = r"\[(\w+)\](.*?)\n(?=\[|\Z)"
-        segments = re.findall(pattern, lyrics, re.DOTALL)
-        return [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
-
-    lyrics = split_lyrics(lyrics_txt + "\n")
-    full_lyrics = "\n".join(lyrics)
-    prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genre_txt.strip()}\n{full_lyrics}"] + lyrics
-
-    raw_output = None
-    stage1_output_set = []
+        raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
+    cuda_idx = cuda_idx
+    max_new_tokens = max_new_tokens * 100
+    
+    with tempfile.TemporaryDirectory() as output_dir:
+        stage1_output_dir = os.path.join(output_dir, f"stage1")
+        os.makedirs(stage1_output_dir, exist_ok=True)
+
+        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
+
+        def load_audio_mono(filepath, sampling_rate=16000):
+            audio, sr = torchaudio.load(filepath)
+            # Convert to mono
+            audio = torch.mean(audio, dim=0, keepdim=True)
+            # Resample if needed
+            if sr != sampling_rate:
+                resampler = Resample(orig_freq=sr, new_freq=sampling_rate)
+                audio = resampler(audio)
+            return audio
+
+        def split_lyrics(lyrics: str):
+            pattern = r"\[(\w+)\](.*?)\n(?=\[|\Z)"
+            segments = re.findall(pattern, lyrics, re.DOTALL)
+            structured_lyrics = [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
+            return structured_lyrics
+
+        # Call the function and print the result
+        stage1_output_set = []
+
+        genres = genre_txt.strip()
+        lyrics = split_lyrics(lyrics_txt + "\n")
+        # intruction
+        full_lyrics = "\n".join(lyrics)
+        prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"]
+        prompt_texts += lyrics
+
+        random_id = uuid.uuid4()
+        output_seq = None
+        # Here is suggested decoding config
+        top_p = 0.93
+        temperature = 1.0
+        repetition_penalty = 1.2
+        # special tokens
+        start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
+        end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
+
+        raw_output = None
+
+        # Format text prompt
+        run_n_segments = min(run_n_segments + 1, len(lyrics))
+
+        print(list(enumerate(tqdm(prompt_texts[:run_n_segments]))))
+
+        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
+            if i == 0:
+                continue
+            if i == 1:
+                if use_audio_prompt:
+                    audio_prompt = load_audio_mono(audio_prompt_path)
+                    audio_prompt.unsqueeze_(0)
+                    with torch.no_grad():
+                        raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=0.5)
+                    raw_codes = raw_codes.transpose(0, 1)
+                    raw_codes = raw_codes.cpu().numpy().astype(np.int16)
+                    # Format audio prompt
+                    code_ids = codectool.npy2ids(raw_codes[0])
+                    audio_prompt_codec = code_ids[int(prompt_start_time * 50): int(prompt_end_time * 50)]  # 50 is tps of xcodec
+                    audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [
+                        mmtokenizer.eoa]
+                    sentence_ids = mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize(
+                        "[end_of_reference]")
+                    head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
+                else:
+                    head_id = mmtokenizer.tokenize(prompt_texts[0])
+                prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+            else:
+                prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+
+            prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device)
+            input_ids = torch.cat([raw_output, prompt_ids], dim=1) if i > 1 else prompt_ids
+            # Use window slicing in case output sequence exceeds the context of model
+            max_context = 16384 - max_new_tokens - 1
+            if input_ids.shape[-1] > max_context:
+                print(
+                    f'Section {i}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
+                input_ids = input_ids[:, -(max_context):]
+            with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.float16):
+                output_seq = model.generate(
+                    input_ids=input_ids,
+                    max_new_tokens=max_new_tokens,
+                    min_new_tokens=100,
+                    do_sample=True,
+                    top_p=top_p,
+                    temperature=temperature,
+                    repetition_penalty=repetition_penalty,
+                    eos_token_id=mmtokenizer.eoa,
+                    pad_token_id=mmtokenizer.eoa,
+                    logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
+                    guidance_scale=guidance_scale,
+                    use_cache=True,
+                    top_k=50,
+                    num_beams=1
+                )
+                if output_seq[0][-1].item() != mmtokenizer.eoa:
+                    tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
+                    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)
+            else:
+                raw_output = output_seq
+            print(len(raw_output))
+
+        # save raw output and check sanity
+        ids = raw_output[0].cpu().numpy()
+        soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
+        eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
+        if len(soa_idx) != len(eoa_idx):
+            raise ValueError(f'invalid pairs of soa and eoa, Num of soa: {len(soa_idx)}, Num of eoa: {len(eoa_idx)}')
+
+        vocals = []
+        instrumentals = []
+        range_begin = 1 if use_audio_prompt else 0
+        for i in range(range_begin, len(soa_idx)):
+            codec_ids = ids[soa_idx[i] + 1:eoa_idx[i]]
+            if codec_ids[0] == 32016:
+                codec_ids = codec_ids[1:]
+            codec_ids = codec_ids[:2 * (codec_ids.shape[0] // 2)]
+            vocals_ids = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0])
+            vocals.append(vocals_ids)
+            instrumentals_ids = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[1])
+            instrumentals.append(instrumentals_ids)
+        vocals = np.concatenate(vocals, axis=1)
+        instrumentals = np.concatenate(instrumentals, axis=1)
+        
+        vocal_save_path = os.path.join(stage1_output_dir, f"vocal_{random_id}".replace('.', '@') + '.npy')
+        inst_save_path = os.path.join(stage1_output_dir, f"instrumental_{random_id}".replace('.', '@') + '.npy')
+        np.save(vocal_save_path, vocals)
+        np.save(inst_save_path, instrumentals)
+        stage1_output_set.append(vocal_save_path)
+        stage1_output_set.append(inst_save_path)
+        
 
-    class BlockTokenRangeProcessor(LogitsProcessor):
-        def __init__(self, start_id, end_id):
-            self.blocked_token_ids = list(range(start_id, end_id))
+        print("Converting to Audio...")
+
+        # convert audio tokens to audio
+        def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
+            folder_path = os.path.dirname(path)
+            if not os.path.exists(folder_path):
+                os.makedirs(folder_path)
+            limit = 0.99
+            max_val = wav.abs().max()
+            wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
+            torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)
+
+        # reconstruct tracks
+        recons_output_dir = os.path.join(output_dir, "recons")
+        recons_mix_dir = os.path.join(recons_output_dir, 'mix')
+        os.makedirs(recons_mix_dir, exist_ok=True)
+        tracks = []
+        for npy in stage1_output_set:
+            codec_result = np.load(npy)
+            decodec_rlt = []
+            with torch.no_grad():
+                decoded_waveform = codec_model.decode(
+                    torch.as_tensor(codec_result.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(
+                        device))
+            decoded_waveform = decoded_waveform.cpu().squeeze(0)
+            decodec_rlt.append(torch.as_tensor(decoded_waveform))
+            decodec_rlt = torch.cat(decodec_rlt, dim=-1)
+            save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy))[0] + ".mp3")
+            tracks.append(save_path)
+            save_audio(decodec_rlt, save_path, 16000)
+        # mix tracks
+        for inst_path in tracks:
+            try:
+                if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) \
+                        and 'instrumental' in inst_path:
+                    # find pair
+                    vocal_path = inst_path.replace('instrumental', 'vocal')
+                    if not os.path.exists(vocal_path):
+                        continue
+                    # mix
+                    recons_mix = os.path.join(recons_mix_dir,
+                                              os.path.basename(inst_path).replace('instrumental', 'mixed'))
+                    vocal_stem, sr = sf.read(inst_path)
+                    instrumental_stem, _ = sf.read(vocal_path)
+                    mix_stem = (vocal_stem + instrumental_stem) / 1
+                    sf.write(recons_mix, mix_stem, sr)
+            except Exception as e:
+                print(e)
+
+        # vocoder to upsample audios
+        vocoder_output_dir = os.path.join(output_dir, 'vocoder')
+        vocoder_stems_dir = os.path.join(vocoder_output_dir, 'stems')
+        vocoder_mix_dir = os.path.join(vocoder_output_dir, 'mix')
+        os.makedirs(vocoder_mix_dir, exist_ok=True)
+        os.makedirs(vocoder_stems_dir, exist_ok=True)
+        instrumental_output = None
+        vocal_output = None
+        for npy in stage1_output_set:
+            if 'instrumental' in npy:
+                # Process instrumental
+                instrumental_output = process_audio(
+                    npy,
+                    os.path.join(vocoder_stems_dir, 'instrumental.mp3'),
+                    rescale,
+                    argparse.Namespace(**locals()),  # Convert local variables to argparse.Namespace
+                    inst_decoder,
+                    codec_model
+                )
+            else:
+                # Process vocal
+                vocal_output = process_audio(
+                    npy,
+                    os.path.join(vocoder_stems_dir, 'vocal.mp3'),
+                    rescale,
+                    argparse.Namespace(**locals()),  # Convert local variables to argparse.Namespace
+                    vocal_decoder,
+                    codec_model
+                )
+        # mix tracks
+        try:
+            mix_output = instrumental_output + vocal_output
+            vocoder_mix = os.path.join(vocoder_mix_dir, os.path.basename(recons_mix))
+            save_audio(mix_output, vocoder_mix, 44100, rescale)
+            print(f"Created mix: {vocoder_mix}")
+        except RuntimeError as e:
+            print(e)
+            print(f"mix {vocoder_mix} failed! inst: {instrumental_output.shape}, vocal: {vocal_output.shape}")
+
+        # Post process
+        final_output_path = os.path.join(output_dir, os.path.basename(recons_mix))
+        replace_low_freq_with_energy_matched(
+            a_file=recons_mix,  # 16kHz
+            b_file=vocoder_mix,  # 48kHz
+            c_file=final_output_path,
+            cutoff_freq=5500.0
+        )
+        print("All process Done")
+        
+        # Load the final audio file and return the numpy array
+        final_audio, sr = torchaudio.load(final_output_path)
+        return (sr, final_audio.squeeze().numpy())
 
-        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
-
-        if i == 0:
-            continue
-
-        if i == 1 and use_audio_prompt:
-            audio_prompt = load_audio_mono(audio_prompt_path)
-            audio_prompt = audio_prompt.unsqueeze(0).to(device)
-            raw_codes = codec_model.encode(audio_prompt, target_bw=0.5).transpose(0, 1).cpu().numpy().astype(np.int16)
-            audio_prompt_codec = codectool.npy2ids(raw_codes[0])[int(prompt_start_time * 50): int(prompt_end_time * 50)]
-            audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
-            sentence_ids = mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
-            head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
-        else:
-            head_id = mmtokenizer.tokenize(prompt_texts[0])
-
-        prompt_ids = head_id + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
-        prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device)
-
-        input_ids = torch.cat([raw_output, prompt_ids], dim=1) if i > 1 else prompt_ids
-
-        max_context = 16384 - max_new_tokens - 1
-        if input_ids.shape[-1] > max_context:
-            input_ids = input_ids[:, -(max_context):]
-
-        with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.float16):
-            output_seq = model.generate(
-                input_ids=input_ids,
-                max_new_tokens=max_new_tokens,
-                min_new_tokens=100,
-                do_sample=True,
-                top_p=top_p,
-                temperature=temperature,
-                repetition_penalty=repetition_penalty,
-                eos_token_id=mmtokenizer.eoa,
-                pad_token_id=mmtokenizer.eoa,
-                logits_processor=LogitsProcessorList([
-                    BlockTokenRangeProcessor(0, 32002),
-                    BlockTokenRangeProcessor(32016, 32016)
-                ]),
-                guidance_scale=guidance_scale,
-                use_cache=True,
-                top_k=50,
-                num_beams=1
-            )
-
-        if output_seq[0][-1].item() != mmtokenizer.eoa:
-            tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(device)
-            output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
-
-        raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1) if i > 1 else output_seq
-
-    # Process and save outputs
-    ids = raw_output[0].cpu().numpy()
-    soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
-    eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
-
-    vocals, instrumentals = [], []
-    for i in range(len(soa_idx)):
-        codec_ids = ids[soa_idx[i] + 1:eoa_idx[i]]
-        if codec_ids[0] == 32016:
-            codec_ids = codec_ids[1:]
-        codec_ids = codec_ids[:2 * (codec_ids.shape[0] // 2)]
-        vocals.append(codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0]))
-        instrumentals.append(codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[1]))
-
-    vocals = np.concatenate(vocals, axis=1)
-    instrumentals = np.concatenate(instrumentals, axis=1)
-
-    # Decode and mix audio
-    decoded_vocals = codec_model.decode(torch.as_tensor(vocals.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device)).cpu().squeeze(0)
-    decoded_instrumentals = codec_model.decode(torch.as_tensor(instrumentals.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device)).cpu().squeeze(0)
-
-    mixed_audio = (decoded_vocals + decoded_instrumentals) / 2
-    mixed_audio_np = mixed_audio.detach().numpy()  # Convert to NumPy array
-    mixed_audio_int16 = (mixed_audio_np * 32767).astype(np.int16)  # Convert to int16
-    
-    # Return the sample rate and the converted audio data
-    return (16000, mixed_audio_int16)
 
 @spaces.GPU(duration=120)
 def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=10):
+    # Execute the command
     try:
-        return generate_music(genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments, max_new_tokens=max_new_tokens)
+        audio_data = generate_music(genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments,
+                               cuda_idx=0, max_new_tokens=max_new_tokens)
+        return audio_data
     except Exception as e:
-        gr.Warning("An Error Occurred: " + str(e))
+        gr.Warning("An Error Occured: " + str(e))
         return None
+    finally:
+        print("Temporary files deleted.")
 
-# Gradio Interface
+
+# Gradio
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation")
@@ -316,5 +493,4 @@ Living out my dreams with this mic and a deal
         inputs=[genre_txt, lyrics_txt, num_segments, max_new_tokens],
         outputs=[music_out]
     )
-
 demo.queue().launch(show_error=True)