modified: app.py

app.py

@@ -1,151 +1,175 @@
 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 uuid
+import time
+import copy
+from collections import Counter
+import re
+import numpy as np
+import torchaudio
+import soundfile as sf
+from torchaudio.transforms import Resample
+from einops import rearrange
+from tqdm import tqdm
+from omegaconf import OmegaConf
 
-is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False
+# --- Constants and Environment Setup ---
+IS_SHARED_UI = "innova-ai/YuE-music-generator-demo" in os.environ.get('SPACE_ID', '')
+OUTPUT_DIR = "./output"
+XCODEC_FOLDER = "./xcodec_mini_infer"
+MM_TOKENIZER_PATH = "./mm_tokenizer_v0.2_hf/tokenizer.model"
+STAGE1_MODEL_NAME = "m-a-p/YuE-s1-7B-anneal-en-cot"
 
-# Install required package
+# --- Utility Functions ---
 def install_flash_attn():
+    """Installs flash-attn using pip."""
     try:
         print("Installing flash-attn...")
-        # Install flash attention
         subprocess.run(
             "pip install flash-attn --no-build-isolation",
             env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
             shell=True,
+            check=True  # Raise an exception if the command fails
         )
         print("flash-attn installed successfully!")
     except subprocess.CalledProcessError as e:
         print(f"Failed to install flash-attn: {e}")
         exit(1)
 
-# Install flash-attn
-install_flash_attn()
-
-from huggingface_hub import snapshot_download 
-
-# Create xcodec_mini_infer folder
-folder_path = './xcodec_mini_infer'
+def download_xcodec_model(folder_path):
+  """Downloads xcodec model from huggingface hub."""
+  if not os.path.exists(folder_path):
+      os.makedirs(folder_path, exist_ok=True)
+      print(f"Folder created at: {folder_path}")
+  else:
+      print(f"Folder already exists at: {folder_path}")
 
-# Create the folder if it doesn't exist
-if not os.path.exists(folder_path):
-    os.mkdir(folder_path)
-    print(f"Folder created at: {folder_path}")
-else:
-    print(f"Folder already exists at: {folder_path}")
+  snapshot_download(
+      repo_id = "m-a-p/xcodec_mini_infer",
+      local_dir = folder_path
+  )
+  print(f"Downloaded xcodec model to {folder_path}")
 
-snapshot_download(
-    repo_id = "m-a-p/xcodec_mini_infer",
-    local_dir = "./xcodec_mini_infer"
-)
 
-# 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 change_working_directory(directory):
+    """Changes the working directory."""
+    try:
+        os.chdir(directory)
+        print(f"Changed working directory to: {os.getcwd()}")
+    except FileNotFoundError:
+        print(f"Directory not found: {directory}")
+        exit(1)
 
 def empty_output_folder(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:
+    """Clears the output directory."""
+    if not os.path.exists(output_dir):
+        return
+    for file in os.listdir(output_dir):
         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)
-    # Ensure content ends with newline and normalize line endings
-    content = content.strip() + "\n\n"  # Add extra newline at end
+    """Creates a temporary file with given content."""
+    content = content.strip() + "\n\n"
     content = content.replace("\r\n", "\n").replace("\r", "\n")
-    temp_file.write(content)
-    temp_file.close()
-    
-    # Debug: Print file contents
+    with tempfile.NamedTemporaryFile(delete=False, mode="w", prefix=prefix, suffix=suffix) as temp_file:
+        temp_file.write(content)
+        temp_file_name = temp_file.name
     print(f"\nContent written to {prefix}{suffix}:")
     print(content)
     print("---")
-    
-    return temp_file.name
+    return temp_file_name
 
 def get_last_mp3_file(output_dir):
-    # 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')]
-    
+    """Returns the path to the most recently modified .mp3 file in the directory, or None if none exists."""
+    mp3_files = [os.path.join(output_dir, file) for file in os.listdir(output_dir) 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]
-
-device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
+    return max(mp3_files, key=os.path.getmtime)
+
+def load_audio_mono(filepath, sampling_rate=16000):
+    """Loads an audio file and converts it to mono at the desired sample rate."""
+    audio, sr = torchaudio.load(filepath)
+    audio = torch.mean(audio, dim=0, keepdim=True)  # Convert to mono
+    if sr != sampling_rate:
+        resampler = Resample(orig_freq=sr, new_freq=sampling_rate)
+        audio = resampler(audio)
+    return audio
+
+def split_lyrics(lyrics: str):
+    """Splits lyrics into segments based on the [section] tags."""
+    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]
+
+def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
+    """Saves a torch audio tensor to a file."""
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    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)
+
+# --- Model Initialization ---
+def initialize_models(device):
+    """Initializes and loads all required models."""
+    print(f"Using device: {device}")
+    # Load Stage 1 Model
+    stage1_model = AutoModelForCausalLM.from_pretrained(
+        STAGE1_MODEL_NAME,
+        torch_dtype=torch.float16,
+        attn_implementation="flash_attention_2",
+    ).to(device).eval()
+
+    # Load Tokenizer
+    mmtokenizer = _MMSentencePieceTokenizer(MM_TOKENIZER_PATH)
+
+     # Load Codec Model
+    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'))
+    from codecmanipulator import CodecManipulator
+    from models.soundstream_hubert_new import SoundStream
 
-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()
+    codectool = CodecManipulator("xcodec", 0, 1)
+    basic_model_config=os.path.join(XCODEC_FOLDER, "final_ckpt", "config.yaml")
+    resume_path=os.path.join(XCODEC_FOLDER, "final_ckpt", "ckpt_00360000.pth")
+    model_config = OmegaConf.load(basic_model_config)
+    codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
+    parameter_dict = torch.load(resume_path, map_location='cpu')
+    codec_model.load_state_dict(parameter_dict['codec_model'])
+    codec_model.to(device).eval()
+    
+    return stage1_model, mmtokenizer, codectool, codec_model
 
-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
+# --- Logits Processor ---
+class BlockTokenRangeProcessor(LogitsProcessor):
+    def __init__(self, start_id, end_id):
+        self.blocked_token_ids = list(range(start_id, end_id))
 
-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
+    def __call__(self, input_ids, scores):
+        scores[:, self.blocked_token_ids] = -float("inf")
+        return scores
 
+# --- Music Generation Core Function ---
 def generate_music(
-    stage1_model="m-a-p/YuE-s1-7B-anneal-en-cot",
+    stage1_model,
+    mmtokenizer,
+    codectool,
+    codec_model,
     max_new_tokens=3000,
     run_n_segments=2,
     genre_txt=None,
@@ -154,42 +178,22 @@ def generate_music(
     audio_prompt_path="",
     prompt_start_time=0.0,
     prompt_end_time=30.0,
-    output_dir="./output",
+    output_dir=OUTPUT_DIR,
     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(basic_model_config)
-    codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
-    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()
-
+    # Load Model Parameters for decoding
     class BlockTokenRangeProcessor(LogitsProcessor):
         def __init__(self, start_id, end_id):
             self.blocked_token_ids = list(range(start_id, end_id))
@@ -198,56 +202,24 @@ def generate_music(
             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())
+    # Split lyrics
     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))
+    stage1_output_set = []
 
     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
@@ -281,7 +253,7 @@ def generate_music(
             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.no_grad():
-            output_seq = model.generate(
+            output_seq = stage1_model.generate(
                 input_ids=input_ids, 
                 max_new_tokens=max_new_tokens, 
                 min_new_tokens=100, 
@@ -295,7 +267,7 @@ def generate_music(
                 guidance_scale=guidance_scale,
                 )
             if output_seq[0][-1].item() != mmtokenizer.eoa:
-                tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
+                tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(stage1_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)
@@ -331,24 +303,15 @@ def generate_music(
     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
+        stage1_model.cpu()
+        del stage1_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')
@@ -384,80 +347,37 @@ def generate_music(
             print(e)
     return recons_mix
 
-    # 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")
-
-
+# --- Gradio Interface ---
 @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)
+    """Main function that runs model and returns output audio."""
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    print(f"Output folder ensured at: {OUTPUT_DIR}")
+    empty_output_folder(OUTPUT_DIR)
+    
+    device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
+    stage1_model, mmtokenizer, codectool, codec_model = initialize_models(device)
     
-    # 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
+         music = generate_music(
+            stage1_model=stage1_model, 
+            mmtokenizer=mmtokenizer,
+            codectool=codectool,
+            codec_model=codec_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")
@@ -480,9 +400,9 @@ with gr.Blocks() as demo:
                 lyrics_txt = gr.Textbox(label="Lyrics")
                 
             with gr.Column():
-                if is_shared_ui:
+                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) # increase it after testing
+                    max_new_tokens = gr.Slider(label="Max New Tokens", info="100 tokens equals 1 second of music", minimum=100, maximum="3000", step=100, value=500, interactive=True) 
                 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)
@@ -529,7 +449,6 @@ Living out my dreams with this mic and a deal
              inputs = [genre_txt, lyrics_txt],
             outputs = [music_out],
             cache_examples = False,
-            # cache_mode="lazy",
             fn=infer
         )
     
@@ -538,4 +457,14 @@ 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_api=False, show_error=True) 
+
+# --- Initialization and Execution ---
+if __name__ == "__main__":
+    # Install Flash Attention
+    install_flash_attn()
+    # Download xcodec mini infer
+    download_xcodec_model(XCODEC_FOLDER)
+    # Change to inference working directory
+    change_working_directory(".")
+
+    demo.queue().launch(show_api=False, show_error=True)