import gradio as gr import subprocess import os import shutil import tempfile import spaces import torch import sys import uuid import re import numpy as np import json import time import copy from collections import Counter # Install flash-attn and set environment variable to skip cuda build print("Installing flash-attn...") subprocess.run( "pip install flash-attn --no-build-isolation", env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"}, shell=True ) # Download snapshot from huggingface_hub from huggingface_hub import snapshot_download folder_path = './xcodec_mini_infer' 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 ) # Change working directory to current folder 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) # Append necessary module paths base_path = os.path.dirname(os.path.abspath(__file__)) sys.path.append(os.path.join(base_path, 'xcodec_mini_infer')) sys.path.append(os.path.join(base_path, 'xcodec_mini_infer', 'descriptaudiocodec')) # Other imports 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 from models.soundstream_hubert_new import SoundStream # Device setup device = "cuda:0" # Load and (optionally) compile the LM model model = AutoModelForCausalLM.from_pretrained( "m-a-p/YuE-s1-7B-anneal-en-cot", torch_dtype=torch.float16, attn_implementation="flash_attention_2", ).to(device) model.eval() try: # torch.compile is available in PyTorch 2.0+ model = torch.compile(model) except Exception as e: print("torch.compile not used for model:", e) # File paths for codec model checkpoint basic_model_config = os.path.join(folder_path, 'final_ckpt/config.yaml') resume_path = os.path.join(folder_path, 'final_ckpt/ckpt_00360000.pth') # Initialize tokenizer and codec manipulator mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model") codectool = CodecManipulator("xcodec", 0, 1) # Load codec model config and initialize codec model model_config = OmegaConf.load(basic_model_config) # Dynamically create the model from its name in the config. codec_class = eval(model_config.generator.name) codec_model = codec_class(**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.eval() try: codec_model = torch.compile(codec_model) except Exception as e: print("torch.compile not used for codec_model:", e) # Pre-compile the regex pattern for splitting lyrics LYRICS_PATTERN = re.compile(r"\[(\w+)\](.*?)\n(?=\[|\Z)", re.DOTALL) # ------------------ GPU decorated generation function ------------------ # @spaces.GPU(duration=120) 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 'use_audio_prompt' is enabled!") max_new_tokens = max_new_tokens * 100 # scaling factor with tempfile.TemporaryDirectory() as output_dir: stage1_output_dir = os.path.join(output_dir, "stage1") os.makedirs(stage1_output_dir, exist_ok=True) # -- In-place logits processor that blocks token ranges -- class BlockTokenRangeProcessor(LogitsProcessor): def __init__(self, start_id, end_id): # Pre-create a tensor for indices if possible 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 # -- Audio processing utility -- def load_audio_mono(filepath, sampling_rate=16000): audio, sr = torchaudio.load(filepath) audio = audio.mean(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 # -- Lyrics splitting using precompiled regex -- def split_lyrics(lyrics: str): segments = LYRICS_PATTERN.findall(lyrics) # Return segments with formatting (strip extra whitespace) return [f"[{tag}]\n{text.strip()}\n\n" for tag, text in segments] # Prepare prompt texts genres = genre_txt.strip() if genre_txt else "" lyrics_segments = split_lyrics(lyrics_txt + "\n") full_lyrics = "\n".join(lyrics_segments) # The first prompt is a global instruction; the rest are segments. prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"] prompt_texts += lyrics_segments random_id = uuid.uuid4() raw_output = None # Decoding config parameters top_p = 0.93 temperature = 1.0 repetition_penalty = 1.2 # Pre-tokenize static tokens start_of_segment = mmtokenizer.tokenize('[start_of_segment]') end_of_segment = mmtokenizer.tokenize('[end_of_segment]') soa_token = mmtokenizer.soa # start-of-audio token id eoa_token = mmtokenizer.eoa # end-of-audio token id # Pre-tokenize the global prompt (first element) global_prompt_ids = mmtokenizer.tokenize(prompt_texts[0]) run_n_segments = min(run_n_segments + 1, len(prompt_texts)) # Loop over segments. (Note: Each segment is processed sequentially.) for i, p in enumerate(tqdm(prompt_texts[:run_n_segments], desc="Generating segments")): # Remove any spurious tokens in the text section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '') guidance_scale = 1.5 if i <= 1 else 1.2 if i == 0: # Skip generation on the instruction segment. continue # Build prompt IDs differently depending on whether audio prompt is enabled. if i == 1: if use_audio_prompt: audio_prompt = load_audio_mono(audio_prompt_path) audio_prompt = audio_prompt.unsqueeze(0) with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.float16): raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=0.5) # Process raw codes (transpose and convert to numpy) raw_codes = raw_codes.transpose(0, 1).cpu().numpy().astype(np.int16) code_ids = codectool.npy2ids(raw_codes[0]) # Slice using prompt start/end time (assuming 50 tokens per second) audio_prompt_codec = code_ids[int(prompt_start_time * 50): int(prompt_end_time * 50)] audio_prompt_codec_ids = [soa_token] + codectool.sep_ids + audio_prompt_codec + [eoa_token] sentence_ids = mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]") head_id = global_prompt_ids + sentence_ids else: head_id = global_prompt_ids prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [soa_token] + codectool.sep_ids else: prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [soa_token] + codectool.sep_ids prompt_ids_tensor = torch.as_tensor(prompt_ids, device=device).unsqueeze(0) if raw_output is not None: # Concatenate previous outputs with the new prompt input_ids = torch.cat([raw_output, prompt_ids_tensor], dim=1) else: input_ids = prompt_ids_tensor # Enforce maximum context window by slicing if needed 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.cuda.amp.autocast(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=eoa_token, pad_token_id=eoa_token, logits_processor=LogitsProcessorList([ BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016) ]), guidance_scale=guidance_scale, use_cache=True ) # Ensure the output ends with an end-of-audio token if output_seq[0, -1].item() != eoa_token: tensor_eoa = torch.as_tensor([[eoa_token]], device=device) output_seq = torch.cat((output_seq, tensor_eoa), dim=1) # For subsequent segments, append only the newly generated tokens. if raw_output is not None: new_tokens = output_seq[:, input_ids.shape[-1]:] raw_output = torch.cat([raw_output, prompt_ids_tensor, new_tokens], dim=1) else: raw_output = output_seq # Save raw output codec tokens to temporary files and check token pairs. ids = raw_output[0].cpu().numpy() soa_idx = np.where(ids == soa_token)[0] eoa_idx = np.where(ids == eoa_token)[0] 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_list = [] instrumentals_list = [] # If using an audio prompt, skip the first pair (it may be reference) start_idx = 1 if use_audio_prompt else 0 for i in range(start_idx, len(soa_idx)): codec_ids = ids[soa_idx[i] + 1: eoa_idx[i]] if codec_ids[0] == 32016: codec_ids = codec_ids[1:] # Force even length and reshape into 2 channels. codec_ids = codec_ids[:2 * (len(codec_ids) // 2)] codec_ids = np.array(codec_ids) reshaped = rearrange(codec_ids, "(n b) -> b n", b=2) vocals_list.append(codectool.ids2npy(reshaped[0])) instrumentals_list.append(codectool.ids2npy(reshaped[1])) vocals = np.concatenate(vocals_list, axis=1) instrumentals = np.concatenate(instrumentals_list, axis=1) # Save the numpy arrays to temporary files vocal_save_path = os.path.join(stage1_output_dir, f"vocal_{str(random_id).replace('.', '@')}.npy") inst_save_path = os.path.join(stage1_output_dir, f"instrumental_{str(random_id).replace('.', '@')}.npy") np.save(vocal_save_path, vocals) np.save(inst_save_path, instrumentals) stage1_output_set = [vocal_save_path, inst_save_path] print("Converting to Audio...") # Utility function for saving audio with in-place clipping def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False): os.makedirs(os.path.dirname(path), exist_ok=True) limit = 0.99 max_val = wav.abs().max().item() if rescale and max_val > 0: wav = wav * (limit / max_val) else: wav = wav.clamp(-limit, limit) torchaudio.save(path, wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16) # Reconstruct tracks by decoding codec tokens 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_path in stage1_output_set: codec_result = np.load(npy_path) with torch.inference_mode(): # Adjust shape: (1, T, C) expected by the decoder input_tensor = torch.as_tensor(codec_result.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device) decoded_waveform = codec_model.decode(input_tensor) decoded_waveform = decoded_waveform.cpu().squeeze(0) save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy_path))[0] + ".mp3") tracks.append(save_path) save_audio(decoded_waveform, save_path, sample_rate=16000) # Mix vocal and instrumental tracks (using torch to avoid extra I/O if possible) for inst_path in tracks: try: if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) and 'instrumental' in inst_path: vocal_path = inst_path.replace('instrumental', 'vocal') if not os.path.exists(vocal_path): continue # Read using soundfile vocal_stem, sr = sf.read(vocal_path) instrumental_stem, _ = sf.read(inst_path) mix_stem = (vocal_stem + instrumental_stem) / 1.0 mix_path = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace('instrumental', 'mixed')) # Write the mix to disk (if needed) or return in memory # Here we return three tuples: (sr, mix), (sr, vocal), (sr, instrumental) return (sr, (mix_stem * 32767).astype(np.int16)), (sr, (vocal_stem * 32767).astype(np.int16)), (sr, (instrumental_stem * 32767).astype(np.int16)) except Exception as e: print("Mixing error:", e) return None, None, None # ------------------ Inference function and Gradio UI ------------------ # def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=15): try: mixed_audio_data, vocal_audio_data, instrumental_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 mixed_audio_data, vocal_audio_data, instrumental_audio_data except Exception as e: gr.Warning("An Error Occurred: " + str(e)) return None, None, None finally: print("Temporary files deleted.") # Build Gradio UI 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(): num_segments = gr.Number(label="Number of Segments", value=2, interactive=True) max_new_tokens = gr.Slider(label="Duration of song", minimum=1, maximum=30, step=1, value=15, interactive=True) submit_btn = gr.Button("Submit") music_out = gr.Audio(label="Mixed Audio Result") with gr.Accordion(label="Vocal and Instrumental Result", open=False): vocal_out = gr.Audio(label="Vocal Audio") instrumental_out = gr.Audio(label="Instrumental Audio") gr.Examples( examples=[ [ "Bass Metalcore Thrash Metal Furious bright vocal male Angry aggressive vocal Guitar", """[verse] Step back cause I'll ignite Won't quit without a fight No escape, gear up, it's a fierce fight Brace up, raise your hands up and light Fear the might. Step back cause I'll ignite Won't back down without a fight It keeps going and going, the heat is on. [chorus] Hot flame. Hot flame. Still here, still holding aim I don't care if I'm bright or dim: nah. I've made it clear, I'll make it again All I want is my crew and my gain. I'm feeling wild, got a bit of rebel style. Locked inside my mind, hot flame. """ ], [ "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'mma keep it real 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, vocal_out, instrumental_out], cache_examples=True, cache_mode="eager", fn=infer ) submit_btn.click( fn=infer, inputs=[genre_txt, lyrics_txt, num_segments, max_new_tokens], outputs=[music_out, vocal_out, instrumental_out] ) gr.Markdown("## Call for Contributions\nIf you find this space interesting please feel free to contribute.") demo.queue().launch(show_error=True)