modified: app.py

app.py

@@ -1,29 +1,13 @@
 import gradio as gr
 import subprocess
-import os
+import os 
 import shutil
 import tempfile
 import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
 import torch
-from huggingface_hub import snapshot_download
-import sys
-import uuid
-import numpy as np
-import json
-from omegaconf import OmegaConf
-import torchaudio
-from torchaudio.transforms import Resample
-import soundfile as sf
-from tqdm import tqdm
-from einops import rearrange
-import time
-from codecmanipulator import CodecManipulator
-from mmtokenizer import _MMSentencePieceTokenizer
-import re
-
 
-is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ.get('SPACE_ID', '') else False
+is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False
 
 # Install required package
 def install_flash_attn():
@@ -43,10 +27,14 @@ def install_flash_attn():
 # Install flash-attn
 install_flash_attn()
 
-# Download xcodec_mini_infer
+from huggingface_hub import snapshot_download 
+
+# Create xcodec_mini_infer folder
 folder_path = './xcodec_mini_infer'
+
+# Create the folder if it doesn't exist
 if not os.path.exists(folder_path):
-    os.makedirs(folder_path, exist_ok=True)
+    os.mkdir(folder_path)
     print(f"Folder created at: {folder_path}")
 else:
     print(f"Folder already exists at: {folder_path}")
@@ -56,131 +44,208 @@ snapshot_download(
     local_dir = "./xcodec_mini_infer"
 )
 
-# Add to path
-sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
-sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
-
-# Load Model (do this ONCE)
-print("Loading Models...")
-device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
-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()
-
-mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-
-print("Models Loaded!")
-
-
+# Change to the "inference" directory
+inference_dir = "."
+try:
+    os.chdir(inference_dir)
+    print(f"Changed working directory to: {os.getcwd()}")
+except FileNotFoundError:
+    print(f"Directory not found: {inference_dir}")
+    exit(1)
 
 def empty_output_folder(output_dir):
-    for file in os.listdir(output_dir):
+    # List all files in the output directory
+    files = os.listdir(output_dir)
+    
+    # Iterate over the files and remove them
+    for file in files:
         file_path = os.path.join(output_dir, file)
         try:
             if os.path.isdir(file_path):
+                # If it's a directory, remove it recursively
                 shutil.rmtree(file_path)
             else:
+                # If it's a file, delete it
                 os.remove(file_path)
         except Exception as e:
             print(f"Error deleting file {file_path}: {e}")
 
+# Function to create a temporary file with string content
 def create_temp_file(content, prefix, suffix=".txt"):
     temp_file = tempfile.NamedTemporaryFile(delete=False, mode="w", prefix=prefix, suffix=suffix)
-    content = content.strip() + "\n\n"
+    # Ensure content ends with newline and normalize line endings
+    content = content.strip() + "\n\n"  # Add extra newline at end
     content = content.replace("\r\n", "\n").replace("\r", "\n")
     temp_file.write(content)
     temp_file.close()
+    
+    # Debug: Print file contents
+    print(f"\nContent written to {prefix}{suffix}:")
+    print(content)
+    print("---")
+    
     return temp_file.name
 
-
 def get_last_mp3_file(output_dir):
-    mp3_files = [file for file in os.listdir(output_dir) if file.endswith('.mp3')]
+    # List all files in the output directory
+    files = os.listdir(output_dir)
+    
+    # Filter only .mp3 files
+    mp3_files = [file for file in files if file.endswith('.mp3')]
+    
     if not mp3_files:
         print("No .mp3 files found in the output folder.")
         return None
+    
+    # Get the full path for the mp3 files
     mp3_files_with_path = [os.path.join(output_dir, file) for file in mp3_files]
+    
+    # Sort the files based on the modification time (most recent first)
     mp3_files_with_path.sort(key=lambda x: os.path.getmtime(x), reverse=True)
+    
+    # Return the most recent .mp3 file
     return mp3_files_with_path[0]
 
-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)
-    audio = torch.mean(audio, dim=0, keepdim=True)
-    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
-
-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)
+device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
 
+model = AutoModelForCausalLM.from_pretrained(
+    "m-a-p/YuE-s1-7B-anneal-en-cot", 
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2", # To enable flashattn, you have to install flash-attn
+    )
+model.to(device)
+model.eval()
+
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
+import argparse
+import torch
+import numpy as np
+import json
+from omegaconf import OmegaConf
+import torchaudio
+from torchaudio.transforms import Resample
+import soundfile as sf
+
+import uuid
+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
+import re
 
-@spaces.GPU(duration=120)
 def generate_music(
-    genre_txt=None,
-    lyrics_txt=None,
+    stage1_model="m-a-p/YuE-s1-7B-anneal-en-cot",
     max_new_tokens=3000,
     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,
     output_dir="./output",
     keep_intermediate=False,
+    disable_offload_model=False,
     cuda_idx=0,
+    basic_model_config='./xcodec_mini_infer/final_ckpt/config.yaml',
+    resume_path='./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth',
+    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',
     rescale=False,
 ):
+    if use_audio_prompt and not audio_prompt_path:
+        raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
+    
+    model = stage1_model
+    cuda_idx = cuda_idx
+    max_new_tokens = max_new_tokens
+    stage1_output_dir = os.path.join(output_dir, f"stage1")
+    os.makedirs(stage1_output_dir, exist_ok=True)
+
+    # load tokenizer and model
+    device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
+
+    # Now you can use `device` to move your tensors or models to the GPU (if available)
+    print(f"Using device: {device}")
+
+    mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
+
     codectool = CodecManipulator("xcodec", 0, 1)
-    model_config = OmegaConf.load('./xcodec_mini_infer/final_ckpt/config.yaml')
+    model_config = OmegaConf.load(basic_model_config)
     codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
-    parameter_dict = torch.load('./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', map_location='cpu')
+    parameter_dict = torch.load(resume_path, map_location='cpu')
     codec_model.load_state_dict(parameter_dict['codec_model'])
     codec_model.to(device)
     codec_model.eval()
 
-    if use_audio_prompt and not audio_prompt_path:
-        raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
-
-    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 = []
+    # Tips:
+    # genre tags support instrumental，genre，mood，vocal timbr and vocal gender
+    # # all kinds of tags are needed
+    # with open(genre_txt) as f:
+    #     genres = f.read().strip()
+    # with open(lyrics_txt) as f:
+    #     lyrics = split_lyrics(f.read())
     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])):
@@ -196,8 +261,9 @@ def generate_music(
                     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)] 
+                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
@@ -209,6 +275,7 @@ def generate_music(
 
         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.')
@@ -236,6 +303,7 @@ def generate_music(
             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()
@@ -256,22 +324,36 @@ def generate_music(
         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"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_vocal_{random_id}".replace('.', '@')+'.npy')
     inst_save_path = os.path.join(stage1_output_dir, f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_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)
 
 
+    # offload model
+    if not disable_offload_model:
+        model.cpu()
+        del model
+        torch.cuda.empty_cache()
+
     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=[]
@@ -300,9 +382,83 @@ def generate_music(
                 sf.write(recons_mix, mix_stem, sr)
         except Exception as e:
             print(e)
+    
+
+    # vocoder to upsample audios
+    vocal_decoder, inst_decoder = build_codec_model(config_path, vocal_decoder_path, inst_decoder_path)
+    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
+    replace_low_freq_with_energy_matched(
+        a_file=recons_mix,     # 16kHz
+        b_file=vocoder_mix,     # 48kHz
+        c_file=os.path.join(output_dir, os.path.basename(recons_mix)),
+        cutoff_freq=5500.0
+    )
+    print("All process Done")
     return recons_mix
 
+
+@spaces.GPU(duration=120)
+def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=200):
+
+    # Ensure the output folder exists
+    output_dir = "./output"
+    os.makedirs(output_dir, exist_ok=True)
+    print(f"Output folder ensured at: {output_dir}")
+
+    empty_output_folder(output_dir)
+    
+    # Execute the command
+    try:
+        music = generate_music(stage1_model=model, genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments, output_dir=output_dir, cuda_idx=0, max_new_tokens=max_new_tokens)
+        
+        return music
+    except subprocess.CalledProcessError as e:
+        print(f"Error occurred: {e}")
+        return None
+    finally:
+        # Clean up temporary files
+        print("Temporary files deleted.")
+
 # Gradio 
+
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation")
@@ -327,7 +483,7 @@ with gr.Blocks() as demo:
             with gr.Column():
                 if is_shared_ui:
                     num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
-                    max_new_tokens = gr.Slider(label="Max New Tokens", info="100 tokens equals 1 second long music", minimum=100, maximum="3000", step=100, value=500, interactive=True)
+                    max_new_tokens = gr.Slider(label="Max New Tokens", info="100 tokens equals 1 second long music", minimum=100, maximum="3000", step=100, value=500, interactive=True) # increase it after testing
                 else:
                     num_segments = gr.Number(label="Number of Song Segments", value=2, interactive=True)
                     max_new_tokens = gr.Slider(label="Max New Tokens", minimum=500, maximum="24000", step=500, value=3000, interactive=True)
@@ -375,12 +531,12 @@ Living out my dreams with this mic and a deal
             outputs = [music_out],
             cache_examples = False,
             # cache_mode="lazy",
-            fn=generate_music
+            fn=infer
         )
     
     submit_btn.click(
-        fn = generate_music,
+        fn = infer, 
         inputs = [genre_txt, lyrics_txt, num_segments, max_new_tokens],
         outputs = [music_out]
     )
-demo.queue().launch(show_api=False, show_error=True)
+demo.queue().launch(show_api=False, show_error=True)