Spaces:
Build error
Build error
import gradio as gr | |
import subprocess | |
import os | |
import shutil | |
import tempfile | |
import spaces | |
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList | |
import torch | |
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(): | |
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() | |
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.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 = "./xcodec_mini_infer" | |
) | |
# Add xcodec_mini_infer and descriptaudiocodec to sys path | |
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 os | |
import sys | |
import torch | |
import numpy as np | |
import json | |
import re | |
import uuid | |
import gradio as gr | |
from tqdm import tqdm | |
from omegaconf import OmegaConf | |
import torchaudio | |
from torchaudio.transforms import Resample | |
import soundfile as sf | |
from einops import rearrange | |
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList | |
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 | |
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 mmtokenizer import _MMSentencePieceTokenizer | |
# Load models once at startup | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
# Load language model | |
print("Loading language 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() | |
# Load tokenizers and codec tools | |
print("Loading tokenizers...") | |
mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model") | |
codectool = CodecManipulator("xcodec", 0, 1) | |
# Load codec models | |
print("Loading codec models...") | |
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() | |
# Load vocoders | |
print("Loading vocoders...") | |
vocal_decoder, inst_decoder = build_codec_model( | |
'./xcodec_mini_infer/decoders/config.yaml', | |
'./xcodec_mini_infer/decoders/decoder_131000.pth', | |
'./xcodec_mini_infer/decoders/decoder_151000.pth' | |
) | |
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 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] | |
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() | |
wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit) | |
torchaudio.save(path, wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16) | |
def run_inference(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=2000): | |
try: | |
# Create temporary output directory | |
output_dir = tempfile.mkdtemp() | |
stage1_output_dir = os.path.join(output_dir, "stage1") | |
os.makedirs(stage1_output_dir, exist_ok=True) | |
# Process inputs | |
structured_lyrics = split_lyrics(lyrics_txt_content) | |
full_lyrics = "\n".join(structured_lyrics) | |
prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genre_txt_content}\n{full_lyrics}"] + structured_lyrics | |
# Generation parameters | |
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]') | |
run_n_segments = min(num_segments + 1, len(structured_lyrics)) | |
# Generate tokens | |
raw_output = None | |
for i in tqdm(range(1, run_n_segments)): | |
section_text = prompt_texts[i].replace('[start_of_segment]', '').replace('[end_of_segment]', '') | |
guidance_scale = 1.5 if i <= 1 else 1.2 | |
prompt_ids = 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 = prompt_ids if i == 1 else torch.cat([raw_output, prompt_ids], dim=1) | |
if input_ids.shape[-1] > 16384 - max_new_tokens - 1: | |
input_ids = input_ids[:, -(16384 - max_new_tokens - 1):] | |
with torch.no_grad(): | |
output_seq = model.generate( | |
input_ids=input_ids, | |
max_new_tokens=max_new_tokens, | |
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, | |
) | |
raw_output = output_seq if i == 1 else torch.cat([raw_output, output_seq[:, input_ids.shape[-1]:]], dim=1) | |
# Process generated tokens | |
ids = raw_output[0].cpu().numpy() | |
soa_idx = np.where(ids == mmtokenizer.soa)[0] | |
eoa_idx = np.where(ids == mmtokenizer.eoa)[0] | |
vocals, instrumentals = [], [] | |
for i in range(len(soa_idx)): | |
codec_ids = ids[soa_idx[i]+1:eoa_idx[i]] | |
codec_ids = codec_ids[:2 * (len(codec_ids) // 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])) | |
# Generate audio | |
vocals = np.concatenate(vocals, axis=1) | |
instrumentals = np.concatenate(instrumentals, axis=1) | |
with torch.no_grad(): | |
vocal_audio = codec_model.decode(torch.tensor(vocals.astype(np.int16)).long().unsqueeze(0).permute(1, 0, 2).to(device)) | |
inst_audio = codec_model.decode(torch.tensor(instrumentals.astype(np.int16)).long().unsqueeze(0).permute(1, 0, 2).to(device)) | |
# Mix and save audio | |
final_audio = (vocal_audio.cpu().squeeze() + inst_audio.cpu().squeeze()) / 2 | |
output_path = os.path.join(output_dir, "final_output.wav") | |
save_audio(final_audio.unsqueeze(0), output_path, 16000) | |
return output_path | |
except Exception as e: | |
print(f"Error during inference: {str(e)}") | |
raise gr.Error(f"Generation failed: {str(e)}") | |
# Gradio UI | |
with gr.Blocks() as demo: | |
gr.Markdown("# YuE Music Generator") | |
with gr.Row(): | |
with gr.Column(): | |
genre_txt = gr.Textbox(label="Genre Tags", placeholder="e.g., female vocal, jazz, piano") | |
lyrics_txt = gr.Textbox(label="Lyrics", lines=10, placeholder="Enter lyrics with sections like [verse], [chorus]") | |
num_segments = gr.Slider(1, 10, value=2, label="Number of Segments") | |
max_tokens = gr.Slider(500, 3000, value=2000, label="Max Tokens") | |
btn = gr.Button("Generate Music") | |
with gr.Column(): | |
audio_out = gr.Audio(label="Generated Music") | |
examples = 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"""] | |
], | |
inputs=[genre_txt, lyrics_txt], | |
outputs=audio_out | |
) | |
btn.click( | |
fn=run_inference, | |
inputs=[genre_txt, lyrics_txt, num_segments, max_tokens], | |
outputs=audio_out | |
) | |
if __name__ == "__main__": | |
demo.launch() |