import gradio as gr import subprocess 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 # Install required package def install_flash_attn(): 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, ) 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() # Download xcodec_mini_infer folder_path = './xcodec_mini_infer' 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}") snapshot_download( repo_id = "m-a-p/xcodec_mini_infer", 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") codectool = CodecManipulator("xcodec", 0, 1) model_config = OmegaConf.load('./xcodec_mini_infer/final_ckpt/config.yaml') 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') codec_model.load_state_dict(parameter_dict['codec_model']) codec_model.to(device) codec_model.eval() print("Models Loaded!") def empty_output_folder(output_dir): for file in os.listdir(output_dir): file_path = os.path.join(output_dir, file) try: if os.path.isdir(file_path): shutil.rmtree(file_path) else: os.remove(file_path) except Exception as e: print(f"Error deleting file {file_path}: {e}") 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" content = content.replace("\r\n", "\n").replace("\r", "\n") temp_file.write(content) temp_file.close() 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')] if not mp3_files: print("No .mp3 files found in the output folder.") return None mp3_files_with_path = [os.path.join(output_dir, file) for file in mp3_files] mp3_files_with_path.sort(key=lambda x: os.path.getmtime(x), reverse=True) 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) @spaces.GPU(duration=120) def generate_music( genre_txt=None, lyrics_txt=None, max_new_tokens=3000, run_n_segments=2, use_audio_prompt=False, audio_prompt_path="", prompt_start_time=0.0, prompt_end_time=30.0, output_dir="./output", keep_intermediate=False, cuda_idx=0, 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'!") stage1_output_dir = os.path.join(output_dir, f"stage1") os.makedirs(stage1_output_dir, exist_ok=True) stage1_output_set = [] genres = genre_txt.strip() 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] {genres}\n{full_lyrics}"] prompt_texts += lyrics random_id = uuid.uuid4() output_seq = None top_p = 0.93 temperature = 1.0 repetition_penalty = 1.2 start_of_segment = mmtokenizer.tokenize('[start_of_segment]') end_of_segment = mmtokenizer.tokenize('[end_of_segment]') raw_output = None 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) 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_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 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.no_grad(): 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, ) 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)) 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"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) print("Converting to Audio...") 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) return recons_mix # Gradio with gr.Blocks() as demo: with gr.Column(): gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation") gr.HTML("""
Duplicate this Space
""") with gr.Row(): with gr.Column(): genre_txt = gr.Textbox(label="Genre") lyrics_txt = gr.Textbox(label="Lyrics") 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) 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) submit_btn = gr.Button("Submit") music_out = gr.Audio(label="Audio Result") gr.Examples( examples = [ [ "female blues airy vocal bright vocal piano sad romantic guitar jazz", """[verse] In the quiet of the evening, shadows start to fall Whispers of the night wind echo through the hall Lost within the silence, I hear your gentle voice Guiding me back homeward, making my heart rejoice [chorus] Don't let this moment fade, hold me close tonight With you here beside me, everything's alright Can't imagine life alone, don't want to let you go Stay with me forever, let our love just flow """ ], [ "rap piano street tough piercing vocal hip-hop synthesizer clear vocal male", """[verse] Woke up in the morning, sun is shining bright Chasing all my dreams, gotta get my mind right City lights are fading, but my vision's clear Got my team beside me, no room for fear Walking through the streets, beats inside my head Every step I take, closer to the bread People passing by, they don't understand Building up my future with my own two hands [chorus] This is my life, and I'm aiming for the top Never gonna quit, no, I'm never gonna stop Through the highs and lows, I'mma keep it real 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=generate_music ) submit_btn.click( fn = generate_music, inputs = [genre_txt, lyrics_txt, num_segments, max_new_tokens], outputs = [music_out] ) demo.queue().launch(show_api=False, show_error=True)