import os
import sys
from mmgp import offload
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 re
import copy
from tqdm import tqdm
from collections import Counter
import argparse
import numpy as np
import torch
import torchaudio
import time
from datetime import datetime
from torchaudio.transforms import Resample
import soundfile as sf
from einops import rearrange
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
from omegaconf import OmegaConf
from codecmanipulator import CodecManipulator
from mmtokenizer import _MMSentencePieceTokenizer
from xcodec_mini_infer.models.soundstream_hubert_new import SoundStream
from xcodec_mini_infer.vocoder import build_codec_model, process_audio
from xcodec_mini_infer.post_process_audio import replace_low_freq_with_energy_matched
import gradio as gr
parser = argparse.ArgumentParser()
# Model Configuration:
parser.add_argument("--max_new_tokens", type=int, default=3000,
help="The maximum number of new tokens to generate in one pass during text generation.")
parser.add_argument("--run_n_segments", type=int, default=2,
help="The number of segments to process during the generation.")
# Prompt
parser.add_argument("--genre_txt", type=str, default="prompt_examples/genrerock.txt",
help="The file path to a text file containing genre tags that describe the musical style or characteristics (e.g., instrumental, genre, mood, vocal timbre, vocal gender). This is used as part of the generation prompt.")
parser.add_argument("--lyrics_txt", type=str, default="prompt_examples/lastxmas.txt",
help="The file path to a text file containing the lyrics for the music generation. These lyrics will be processed and split into structured segments to guide the generation process.")
parser.add_argument("--use_audio_prompt", action="store_true",
help="If set, the model will use an audio file as a prompt during generation. The audio file should be specified using --audio_prompt_path.")
parser.add_argument("--audio_prompt_path", type=str, default="",
help="The file path to an audio file to use as a reference prompt when --use_audio_prompt is enabled.")
parser.add_argument("--prompt_start_time", type=float, default=0.0,
help="The start time in seconds to extract the audio prompt from the given audio file.")
parser.add_argument("--prompt_end_time", type=float, default=30.0,
help="The end time in seconds to extract the audio prompt from the given audio file.")
parser.add_argument("--use_dual_tracks_prompt", action="store_true",
help="If set, the model will use dual tracks as a prompt during generation. The vocal and instrumental files should be specified using --vocal_track_prompt_path and --instrumental_track_prompt_path.")
parser.add_argument("--vocal_track_prompt_path", type=str, default="",
help="The file path to a vocal track file to use as a reference prompt when --use_dual_tracks_prompt is enabled.")
parser.add_argument("--instrumental_track_prompt_path", type=str, default="",
help="The file path to an instrumental track file to use as a reference prompt when --use_dual_tracks_prompt is enabled.")
# Output
parser.add_argument("--output_dir", type=str, default="./output",
help="The directory where generated outputs will be saved.")
parser.add_argument("--keep_intermediate", action="store_true",
help="If set, intermediate outputs will be saved during processing.")
parser.add_argument("--disable_offload_model", action="store_true",
help="If set, the model will not be offloaded from the GPU to CPU after Stage 1 inference.")
parser.add_argument("--cuda_idx", type=int, default=0)
# Config for xcodec and upsampler
parser.add_argument('--basic_model_config', default='./xcodec_mini_infer/final_ckpt/config.yaml',
help='YAML files for xcodec configurations.')
parser.add_argument('--resume_path', default='./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth',
help='Path to the xcodec checkpoint.')
parser.add_argument('--config_path', type=str, default='./xcodec_mini_infer/decoders/config.yaml',
help='Path to Vocos config file.')
parser.add_argument('--vocal_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_131000.pth',
help='Path to Vocos decoder weights.')
parser.add_argument('--inst_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_151000.pth',
help='Path to Vocos decoder weights.')
parser.add_argument('-r', '--rescale', action='store_true', help='Rescale output to avoid clipping.')
parser.add_argument("--profile", type=int, default=3)
parser.add_argument("--verbose", type=int, default=1)
parser.add_argument("--compile", action="store_true")
parser.add_argument("--sdpa", action="store_true")
parser.add_argument("--icl", action="store_true")
parser.add_argument("--turbo-stage2", action="store_true")
# Gradio server
parser.add_argument("--server_name", type=str, default="localhost",
help="The server name for the wWbUI. By default it exposes the service to all network interfaces. Set to localhost, if you want to restrict access to the local machine.")
parser.add_argument("--server_port", type=int, default=7860, help="The port number for the WebUI.")
args = parser.parse_args()
# set up arguments
profile = args.profile
compile = args.compile
sdpa = args.sdpa
use_icl = args.icl
if use_icl:
args.stage1_model = "m-a-p/YuE-s1-7B-anneal-en-icl"
else:
args.stage1_model = "m-a-p/YuE-s1-7B-anneal-en-cot"
args.stage2_model = "m-a-p/YuE-s2-1B-general"
args.stage2_batch_size = [20, 20, 20, 4, 3, 2][profile]
if sdpa:
attn_implementation = "sdpa"
else:
attn_implementation = "flash_attention_2"
if args.use_audio_prompt and not args.audio_prompt_path:
raise FileNotFoundError(
"Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
if args.use_dual_tracks_prompt and not args.vocal_track_prompt_path and not args.instrumental_track_prompt_path:
raise FileNotFoundError(
"Please offer dual tracks prompt filepath using '--vocal_track_prompt_path' and '--inst_decoder_path', when you enable '--use_dual_tracks_prompt'!")
stage1_model = args.stage1_model
stage2_model = args.stage2_model
cuda_idx = args.cuda_idx
max_new_tokens = args.max_new_tokens
stage1_output_dir = os.path.join(args.output_dir, f"stage1")
stage2_output_dir = stage1_output_dir.replace('stage1', 'stage2')
os.makedirs(stage1_output_dir, exist_ok=True)
os.makedirs(stage2_output_dir, exist_ok=True)
# load tokenizer and model
device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
model = AutoModelForCausalLM.from_pretrained(
stage1_model,
torch_dtype=torch.bfloat16,
attn_implementation=attn_implementation, # To enable flashattn, you have to install flash-attn
)
# to device, if gpu is available
model.to(device)
model.eval()
model_stage2 = AutoModelForCausalLM.from_pretrained(
stage2_model,
torch_dtype=torch.float16,
attn_implementation=attn_implementation,
)
model_stage2.to(device)
model_stage2.eval()
# remove test on arguments for method 'model.generate' in case transformers patch not applied
def nop(nada):
pass
model._validate_model_kwargs = nop
model_stage2._validate_model_kwargs = nop
pipe = {"transformer": model, "stage2": model_stage2}
quantizeTransformer = profile == 3 or profile == 4 or profile == 5
codectool = CodecManipulator("xcodec", 0, 1)
codectool_stage2 = CodecManipulator("xcodec", 0, 8)
model_config = OmegaConf.load(args.basic_model_config)
codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
parameter_dict = torch.load(args.resume_path, map_location="cpu", weights_only=False)
codec_model.load_state_dict(parameter_dict['codec_model'])
codec_model.to('cpu')
codec_model.eval()
kwargs = {}
if profile == 5:
kwargs["budgets"] = {"transformer": 500, "*": 3000}
kwargs["pinnedMemory"] = True
elif profile == 4:
kwargs["budgets"] = {"transformer": 3000, "*": 5000}
elif profile == 2:
kwargs["budgets"] = 5000
offload.profile(pipe, profile_no=profile, compile=compile, quantizeTransformer=quantizeTransformer,
verboseLevel=args.verbose if args.verbose is not None else 1, **kwargs) # pinnedMemory=False,
class BlockTokenRangeProcessor(LogitsProcessor):
def __init__(self, start_id, end_id):
self.blocked_token_ids = list(range(start_id, end_id))
self.start_id = start_id
self.end_id = end_id
def __call__(self, input_ids, scores):
# scores[:, self.blocked_token_ids] = -float("inf")
scores[:, self.start_id: self.end_id] = -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 encode_audio(codec_model, audio_prompt, device, target_bw=0.5):
if len(audio_prompt.shape) < 3:
audio_prompt.unsqueeze_(0)
with torch.no_grad():
raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=target_bw)
raw_codes = raw_codes.transpose(0, 1)
raw_codes = raw_codes.cpu().numpy().astype(np.int16)
return raw_codes
def split_lyrics(lyrics):
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 get_song_id(seed, genres, top_p, temperature, repetition_penalty, max_new_tokens):
timestamp = datetime.now().strftime("%Y%m%d-%H%M-%S.%f")[:-3]
genres = re.sub(r'[^a-zA-Z0-9_-]', '_', genres.replace(' ', '-'))
genres = re.sub(r'_+', '_', genres).strip('_')
genres = genres[:180]
song_id = f"{timestamp}_{genres}_seed{seed}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}"
return song_id[:240]
def stage1_inference(genres, lyrics_input, run_n_segments, max_new_tokens, seed, state=None, callback=None):
# Tips:
# genre tags support instrumental，genre，mood，vocal timbr and vocal gender
# all kinds of tags are needed
genres = genres.strip()
lyrics = split_lyrics(lyrics_input)
# instruction
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
# 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]')
# Format text prompt
run_n_segments = min(run_n_segments, len(lyrics))
for i, p in enumerate(tqdm(prompt_texts[1:run_n_segments + 1]), 1):
# print(f"---Stage 1: Generating Sequence {i} out of {run_n_segments}")
state["stage"] = f"Stage 1: Generating Sequence {i} out of {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 == 1:
if args.use_dual_tracks_prompt or args.use_audio_prompt:
if args.use_dual_tracks_prompt:
vocals_ids = load_audio_mono(args.vocal_track_prompt_path)
instrumental_ids = load_audio_mono(args.instrumental_track_prompt_path)
vocals_ids = encode_audio(codec_model, vocals_ids, device, target_bw=0.5)
instrumental_ids = encode_audio(codec_model, instrumental_ids, device, target_bw=0.5)
vocals_ids = codectool.npy2ids(vocals_ids[0])
instrumental_ids = codectool.npy2ids(instrumental_ids[0])
min_size = min(len(vocals_ids), len(instrumental_ids))
vocals_ids = vocals_ids[0: min_size]
instrumental_ids = instrumental_ids[0: min_size]
ids_segment_interleaved = rearrange([np.array(vocals_ids), np.array(instrumental_ids)],
'b n -> (n b)')
audio_prompt_codec = ids_segment_interleaved[
int(args.prompt_start_time * 50 * 2): int(args.prompt_end_time * 50 * 2)]
audio_prompt_codec = audio_prompt_codec.tolist()
elif args.use_audio_prompt:
audio_prompt = load_audio_mono(args.audio_prompt_path)
raw_codes = encode_audio(codec_model, audio_prompt, device, target_bw=0.5)
# Format audio prompt
code_ids = codectool.npy2ids(raw_codes[0])
audio_prompt_codec = code_ids[int(args.prompt_start_time * 50): int(
args.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
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
# 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.')
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, 32017)]),
guidance_scale=guidance_scale,
callback=callback,
)
torch.cuda.empty_cache()
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
# 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()
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 args.use_audio_prompt or args.use_dual_tracks_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)
song_id = get_song_id(seed, genres, top_p, temperature, repetition_penalty, max_new_tokens)
vocal_save_path = os.path.join(stage1_output_dir, f"{song_id}_vtrack.npy")
inst_save_path = os.path.join(stage1_output_dir, f"{song_id}_itrack.npy")
np.save(vocal_save_path, vocals)
np.save(inst_save_path, instrumentals)
stage1_output_set = []
stage1_output_set.append(vocal_save_path)
stage1_output_set.append(inst_save_path)
return stage1_output_set
def stage2_generate(model, prompt, batch_size=16, segment_duration=6, state=None, callback=None):
codec_ids = codectool.unflatten(prompt, n_quantizer=1)
codec_ids = codectool.offset_tok_ids(
codec_ids,
global_offset=codectool.global_offset,
codebook_size=codectool.codebook_size,
num_codebooks=codectool.num_codebooks,
).astype(np.int32)
# Prepare prompt_ids based on batch size or single input
if batch_size > 1:
codec_list = []
for i in range(batch_size):
idx_begin = i * segment_duration * 50
idx_end = (i + 1) * segment_duration * 50
codec_list.append(codec_ids[:, idx_begin:idx_end])
codec_ids = np.concatenate(codec_list, axis=0)
prompt_ids = np.concatenate(
[
np.tile([mmtokenizer.soa, mmtokenizer.stage_1], (batch_size, 1)),
codec_ids,
np.tile([mmtokenizer.stage_2], (batch_size, 1)),
],
axis=1
)
else:
prompt_ids = np.concatenate([
np.array([mmtokenizer.soa, mmtokenizer.stage_1]),
codec_ids.flatten(), # Flatten the 2D array to 1D
np.array([mmtokenizer.stage_2])
]).astype(np.int32)
prompt_ids = prompt_ids[np.newaxis, ...]
codec_ids = torch.as_tensor(codec_ids).to(device)
prompt_ids = torch.as_tensor(prompt_ids).to(device)
len_prompt = prompt_ids.shape[-1]
block_list = LogitsProcessorList(
[BlockTokenRangeProcessor(0, 46358), BlockTokenRangeProcessor(53526, mmtokenizer.vocab_size)])
# Teacher forcing generate loop
max_tokens = codec_ids.shape[1] * 8
i = 0
real_max_length = codec_ids.shape[1] * 8 + prompt_ids.shape[1]
session_cache = {"real_max_length": real_max_length}
codec_ids.shape[1]
for frames_idx in range(codec_ids.shape[1]):
if i % 96 == 0:
# print(f"Tokens: {i} out of {max_tokens}")
callback(i, real_max_length)
cb0 = codec_ids[:, frames_idx:frames_idx + 1]
# print(f"insert cb0: {cb0}")
prompt_ids = torch.cat([prompt_ids, cb0], dim=1)
input_ids = prompt_ids
with torch.no_grad():
stage2_output = model.generate(input_ids=input_ids,
min_new_tokens=7,
max_new_tokens=7,
eos_token_id=mmtokenizer.eoa,
pad_token_id=mmtokenizer.eoa,
logits_processor=block_list,
session_cache=session_cache,
)
assert stage2_output.shape[1] - prompt_ids.shape[
1] == 7, f"output new tokens={stage2_output.shape[1] - prompt_ids.shape[1]}"
prompt_ids = stage2_output
i += 8
del session_cache
torch.cuda.empty_cache()
# Return output based on batch size
if batch_size > 1:
output = prompt_ids.cpu().numpy()[:, len_prompt:]
output_list = [output[i] for i in range(batch_size)]
output = np.concatenate(output_list, axis=0)
else:
output = prompt_ids[0].cpu().numpy()[len_prompt:]
return output
def stage2_inference(model, stage1_output_set, stage2_output_dir, batch_size=4, segment_duration=6, state=None,
callback=None):
stage2_result = []
for i in tqdm(range(len(stage1_output_set))):
if i == 0:
# print("---Stage 2.1: Sampling Vocal track")
prefix = "Stage 2.1: Sampling Vocal track"
else:
# print("---Stage 2.2: Sampling Instrumental track")
prefix = "Stage 2.2: Sampling Instrumental track"
output_filename = os.path.join(stage2_output_dir, os.path.basename(stage1_output_set[i]))
if os.path.exists(output_filename) and False:
print(f'{output_filename} stage2 has done.')
stage2_result.append(output_filename)
continue
# Load the prompt
prompt = np.load(stage1_output_set[i]).astype(np.int32)
segment_length = 3
# Only accept 6s segments ( = segment_duration )
output_duration = prompt.shape[-1] // 50 // segment_duration * segment_duration
num_batch = output_duration // segment_duration
any_trail = output_duration * 50 != prompt.shape[-1]
if num_batch <= batch_size:
# If num_batch is less than or equal to batch_size, we can infer the entire prompt at once
# print("Only one segment to process for this track")
max_segments = 2 if any_trail else 1
if max_segments == 1:
state["stage"] = prefix
else:
state["stage"] = prefix + f", segment 1 out of {max_segments}"
output = stage2_generate(model, prompt[:, :output_duration * 50], batch_size=num_batch,
segment_duration=segment_duration, state=state, callback=callback)
else:
# If num_batch is greater than batch_size, process in chunks of batch_size
segments = []
num_segments = (num_batch // batch_size) + (1 if num_batch % batch_size != 0 else 0)
max_segments = num_segments + 1 if any_trail else num_segments
for seg in range(num_segments):
# print(f"Segment {seg+1} out of {max_segments}")
state["stage"] = prefix + f", segment {seg + 1} out of {max_segments}"
start_idx = seg * batch_size * 300
# Ensure the end_idx does not exceed the available length
end_idx = min((seg + 1) * batch_size * 300, output_duration * 50) # Adjust the last segment
current_batch_size = batch_size if seg != num_segments - 1 or num_batch % batch_size == 0 else num_batch % batch_size
segment = stage2_generate(
model,
prompt[:, start_idx:end_idx],
batch_size=current_batch_size,
segment_duration=segment_duration,
state=state,
callback=callback
)
segments.append(segment)
# Concatenate all the segments
output = np.concatenate(segments, axis=0)
# Process the ending part of the prompt
if any_trail:
# print(f"Segment {max_segments} / {max_segments}")
state["stage"] = prefix + f", segment {max_segments} out of {max_segments}"
ending = stage2_generate(model, prompt[:, output_duration * 50:], batch_size=1,
segment_duration=segment_duration, state=state, callback=callback)
output = np.concatenate([output, ending], axis=0)
output = codectool_stage2.ids2npy(output)
# Fix invalid codes (a dirty solution, which may harm the quality of audio)
# We are trying to find better one
fixed_output = copy.deepcopy(output)
for i, line in enumerate(output):
for j, element in enumerate(line):
if element < 0 or element > 1023:
counter = Counter(line)
most_frequant = sorted(counter.items(), key=lambda x: x[1], reverse=True)[0][0]
fixed_output[i, j] = most_frequant
# save output
np.save(output_filename, fixed_output)
stage2_result.append(output_filename)
return stage2_result
def build_callback(state, progress, status):
def callback(tokens_processed, max_tokens):
prefix = state["prefix"]
status = prefix + state["stage"]
tokens_processed += 1
if state.get("abort", False):
status_msg = status + " - Aborting"
raise Exception("abort")
# pipe._interrupt = True
# elif step_idx == num_inference_steps:
# status_msg = status + " - VAE Decoding"
else:
status_msg = status # + " - Denoising"
progress(tokens_processed / max_tokens, desc=status_msg, unit=" %")
return callback
def abort_generation(state):
if "in_progress" in state:
state["abort"] = True
return gr.Button(interactive=False)
else:
return gr.Button(interactive=True)
def refresh_gallery(state):
file_list = state.get("file_list", None)
if len(file_list) > 0:
return file_list[0], file_list
else:
return None, file_list
def finalize_gallery(state):
if "in_progress" in state:
del state["in_progress"]
time.sleep(0.2)
return gr.Button(interactive=True)
def generate_song(genres_input, lyrics_input, run_n_segments, seed, max_new_tokens, vocal_track_prompt,
instrumental_track_prompt, prompt_start_time, prompt_end_time, repeat_generation, state,
progress=gr.Progress()):
args.use_audio_prompt = False
args.use_dual_tracks_prompt = False
# Call the function and print the result
if "abort" in state:
del state["abort"]
state["in_progress"] = True
state["selected"] = 0
file_list = state.get("file_list", [])
if len(file_list) == 0:
state["file_list"] = file_list
if use_icl:
if prompt_start_time > prompt_end_time:
raise gr.Error(f"'Start time' should be less than 'End Time'")
if (prompt_end_time - prompt_start_time) > 30:
raise gr.Error(f"The duration for the audio prompt should not exceed 30s")
if vocal_track_prompt == None:
raise gr.Error(f"You must provide at least a Vocal audio prompt")
args.prompt_start_time = prompt_start_time
args.prompt_end_time = prompt_end_time
if instrumental_track_prompt == None:
args.use_audio_prompt = True
args.audio_prompt_path = vocal_track_prompt
else:
args.use_dual_tracks_prompt = True
args.vocal_track_prompt_path = vocal_track_prompt
args.instrumental_track_prompt_path = instrumental_track_prompt
segment_duration = 3 if args.turbo_stage2 else 6
import random
if seed <= 0:
seed = random.randint(0, 999999999)
genres_input = genres_input.replace("\r", "").split("\n")
song_no = 0
total_songs = repeat_generation * len(genres_input)
start_time = time.time()
for genres_no, genres in enumerate(genres_input):
for gen_no in range(repeat_generation):
song_no += 1
prefix = ""
status = f"Song {song_no}/{total_songs}"
if len(genres_input) > 1:
prefix += f"Genres {genres_no + 1}/{len(genres_input)} > "
if repeat_generation > 1:
prefix += f"Generation {gen_no + 1}/{repeat_generation} > "
state["prefix"] = prefix
# return "output/cot_inspiring-female-uplifting-pop-airy-vocal-electronic-bright-vocal-vocal_tp0@93_T1@0_rp1@2_maxtk3000_mixed_e0a99c45-7f63-41c9-826f-9bde7417db4c.mp3"
torch.cuda.manual_seed(seed)
random.seed(seed)
callback = build_callback(state, progress, status)
# if True:
try:
stage1_output_set = stage1_inference(genres, lyrics_input, run_n_segments, max_new_tokens, seed, state,
callback)
# random_id ="5b4b4613-1cc2-4d84-af7a-243f853f168b"
# stage1_output_set = [ "output/stage1/inspiring-female-uplifting-pop-airy-vocal-electronic-bright-vocal_tp0@93_T1@0_rp1@2_maxtk3000_5b4b4613-1cc2-4d84-af7a-243f853f168b_vtrack.npy",
# "output/stage1/inspiring-female-uplifting-pop-airy-vocal-electronic-bright-vocal_tp0@93_T1@0_rp1@2_maxtk3000_5b4b4613-1cc2-4d84-af7a-243f853f168b_itrack.npy"]
stage2_result = stage2_inference(model_stage2, stage1_output_set, stage2_output_dir,
batch_size=args.stage2_batch_size, segment_duration=segment_duration,
state=state, callback=callback)
except Exception as e:
s = str(e)
if "abort" in s:
stage2_result = None
else:
raise
if stage2_result == None:
end_time = time.time()
yield f"Song Generation Aborted. Total Generation Time: {end_time - start_time:.1f}s"
return
print(stage2_result)
print('Stage 2 DONE.\n')
# 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(args.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 stage2_result:
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, device="cpu"))
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 '_itrack' in inst_path:
# find pair
vocal_path = inst_path.replace('_itrack', '_vtrack')
if not os.path.exists(vocal_path):
continue
# mix
recons_mix = os.path.join(recons_mix_dir,
os.path.basename(inst_path).replace('_itrack', '_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)
# vocoder to upsample audios
vocal_decoder, inst_decoder = build_codec_model(args.config_path, args.vocal_decoder_path,
args.inst_decoder_path)
vocoder_output_dir = os.path.join(args.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)
for npy in stage2_result:
if '_itrack' in npy:
# Process instrumental
instrumental_output = process_audio(
npy,
os.path.join(vocoder_stems_dir, 'itrack.mp3'),
args.rescale,
args,
inst_decoder,
codec_model
)
else:
# Process vocal
vocal_output = process_audio(
npy,
os.path.join(vocoder_stems_dir, 'vtrack.mp3'),
args.rescale,
args,
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, args.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
output_file = os.path.join(args.output_dir, os.path.basename(recons_mix))
replace_low_freq_with_energy_matched(
a_file=recons_mix, # 16kHz
b_file=vocoder_mix, # 48kHz
c_file=output_file,
cutoff_freq=5500.0
)
file_list.insert(0, output_file)
if song_no < total_songs:
yield status
else:
end_time = time.time()
yield f"Total Generation Time: {end_time - start_time:.1f}s"
seed += 1
# return output_file
def create_demo():
with gr.Blocks() as demo:
gr.Markdown("YuEGP v3
")
gr.Markdown(
"YuE is a groundbreaking series of open-source foundation models designed for music generation, specifically for transforming lyrics into full songs (lyrics2song).
")
gr.Markdown(
"GPU Poor version by DeepBeepMeep (Updates / Original). Switch to profile 1 for fast generation (requires a 16 GB VRAM GPU), 1 min of song will take only 4 minutes
")
if use_icl:
gr.Markdown(
"With In Context Learning Mode in addition to the lyrics and genres info, you can provide audio prompts to describe your expectations. You can generate a song with either:
")
gr.Markdown("- a single mixed (song/instruments) Audio track prompt
")
gr.Markdown("- a Vocal track and an Instrumental track prompt
")
gr.Markdown(
"Given some Lyrics and sample audio songs, you can try different Genres Prompt by separating each prompt by a carriage return.")
else:
gr.Markdown(
"Given some Lyrics, you can try different Genres Prompt by separating each prompt by a carriage return.")
with gr.Row():
with gr.Column():
with open(os.path.join("prompt_examples", "lyrics.txt")) as f:
lyrics_file = f.read()
# lyrics_file.replace("\n", "\n\r")
genres_input = gr.Text(label="Genres Prompt (one Genres Prompt per line for multiple generations)",
value="inspiring female uplifting pop airy vocal electronic bright vocal",
lines=3)
lyrics_input = gr.Text(label="Lyrics", lines=20, value=lyrics_file)
repeat_generation = gr.Slider(1, 25.0, value=1.0, step=1,
label="Number of Generated Songs per Genres Prompt")
with gr.Column():
state = gr.State({})
number_sequences = gr.Slider(1, 10, value=2, step=1,
label="Number of Sequences (paragraphs in Lyrics, the higher this number, the higher the VRAM consumption)")
max_new_tokens = gr.Slider(300, 6000, value=3000, step=300,
label="Number of tokens per sequence (1000 tokens = 10s, the higher this number, the higher the VRAM consumption) ")
seed = gr.Slider(0, 999999999, value=123, step=1, label="Seed (0 for random)")
with gr.Row():
with gr.Column():
gen_status = gr.Text(label="Status", interactive=False)
generate_btn = gr.Button("Generate")
abort_btn = gr.Button("Abort")
output = gr.Audio(label="Last Generated Song")
files_history = gr.Files(label="History of Generated Songs (From most Recent to Oldest)",
type='filepath', height=150)
abort_btn.click(abort_generation, state, abort_btn)
gen_status.change(refresh_gallery, inputs=[state], outputs=[output, files_history])
with gr.Row(visible=use_icl): # use_icl
with gr.Column():
vocal_track_prompt = gr.Audio(label="Audio track prompt / Vocal track prompt", type='filepath')
with gr.Column():
instrumental_track_prompt = gr.Audio(
label="Intrumental track prompt (optional if Vocal track prompt set)", type='filepath')
with gr.Row(visible=use_icl):
with gr.Column():
prompt_start_time = gr.Slider(0.0, 300.0, value=0.0, step=0.5, label="Audio Prompt Start time")
prompt_end_time = gr.Slider(0.0, 300.0, value=30.0, step=0.5, label="Audio Prompt End time")
abort_btn.click(abort_generation, state, abort_btn)
generate_btn.click(
fn=generate_song,
inputs=[
genres_input,
lyrics_input,
number_sequences,
seed,
max_new_tokens,
vocal_track_prompt,
instrumental_track_prompt,
prompt_start_time,
prompt_end_time,
repeat_generation,
state
],
outputs=[gen_status] # ,state
).then(
finalize_gallery,
[state],
[abort_btn]
)
return demo
if __name__ == "__main__":
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
demo = create_demo()
demo.launch(
server_name=args.server_name,
server_port=args.server_port,
allowed_paths=[args.output_dir])