app.py

import gradio as gr
import subprocess
import os
import shutil
import tempfile
import spaces
import torch
import sys
import uuid
import re

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
)

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"
)

# 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)

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'))


# don't change above code

import argparse
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
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

device = "cuda:0"

stage2_model = "m-a-p/YuE-s2-1B-general"
model_stage2 = AutoModelForCausalLM.from_pretrained(
    stage2_model,
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2"
    ).to(device)
model_stage2.eval()

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()

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'

mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")

codectool = CodecManipulator("xcodec", 0, 1)
codectool_stage2 = CodecManipulator("xcodec", 0, 8)
model_config = OmegaConf.load(basic_model_config)
# Load codec model
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.eval()

# Preload and compile vocoders
vocal_decoder, inst_decoder = build_codec_model(config_path, vocal_decoder_path, inst_decoder_path)
vocal_decoder.to(device)
inst_decoder.to(device)
vocal_decoder.eval()
inst_decoder.eval()


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)
    # 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


def stage2_generate(model, prompt, batch_size=1): # set batch_size=1 for gradio demo
    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 * 300
            idx_end = (i + 1) * 300
            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
    for frames_idx in range(codec_ids.shape[1]):
        cb0 = codec_ids[:, frames_idx:frames_idx+1]
        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,
            )

        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

    # 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=1): # set batch_size=1 for gradio demo
    stage2_result = []
    for i in tqdm(range(len(stage1_output_set))):
        output_filename = os.path.join(stage2_output_dir, os.path.basename(stage1_output_set[i]))

        if os.path.exists(output_filename):
            print(f'{output_filename} stage2 has done.')
            continue

        # Load the prompt
        prompt = np.load(stage1_output_set[i]).astype(np.int32)

        # Only accept 6s segments
        output_duration = prompt.shape[-1] // 50 // 6 * 6
        num_batch = output_duration // 6

        if num_batch <= batch_size:
            # If num_batch is less than or equal to batch_size, we can infer the entire prompt at once
            output = stage2_generate(model, prompt[:, :output_duration*50], batch_size=num_batch)
        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)

            for seg in range(num_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
                )
                segments.append(segment)

            # Concatenate all the segments
            output = np.concatenate(segments, axis=0)

        # Process the ending part of the prompt
        if output_duration*50 != prompt.shape[-1]:
            ending = stage2_generate(model, prompt[:, output_duration*50:], batch_size=1)
            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


@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 offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
    cuda_idx = cuda_idx
    max_new_tokens = max_new_tokens * 100

    with tempfile.TemporaryDirectory() as output_dir:
        stage1_output_dir = os.path.join(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)

        stage1_output_set = []

        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))

        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)
                    # Format audio prompt
                    code_ids = codectool.npy2ids(raw_codes[0])
                    audio_prompt_codec = code_ids[int(prompt_start_time * 50): int(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.inference_mode(), torch.autocast(device_type='cuda', 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=mmtokenizer.eoa,
                    pad_token_id=mmtokenizer.eoa,
                    logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
                    guidance_scale=guidance_scale,
                    use_cache=True
                )
                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))

        # 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 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"vocal_{random_id}".replace('.', '@') + '.npy')
        inst_save_path = os.path.join(stage1_output_dir, f"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("Stage 2 inference...")
        stage2_result = stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_size=1) # set batch_size=1 for gradio demo
        print('Stage 2 DONE.\n')

        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 from stage 1
        recons_output_dir = os.path.join(output_dir, "recons_stage1") # changed folder name to recons_stage1
        recons_mix_dir = os.path.join(recons_output_dir, 'mix')
        os.makedirs(recons_mix_dir, exist_ok=True)
        tracks_stage1 = [] # changed variable name to tracks_stage1
        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] + "_stage1.mp3") # changed filename to include _stage1
            tracks_stage1.append(save_path) # changed variable name to tracks_stage1
            save_audio(decodec_rlt, save_path, 16000)

        # reconstruct tracks from stage 2 and vocoder
        recons_output_dir = os.path.join(output_dir, "recons_stage2_vocoder") # changed folder name to recons_stage2_vocoder
        recons_mix_dir = os.path.join(recons_output_dir, 'mix')
        os.makedirs(recons_mix_dir, exist_ok=True)
        tracks_stage2_vocoder = [] # changed variable name to tracks_stage2_vocoder
        vocoder_stems_dir = os.path.join(recons_output_dir, 'stems') # vocoder output stems in recons_stage2_vocoder
        os.makedirs(vocoder_stems_dir, exist_ok=True)

        vocal_output = None # initialize for mix error handling
        instrumental_output = None # initialize for mix error handling

        for npy in stage2_result:
            if 'instrumental' in npy:
                # Process instrumental
                instrumental_output = process_audio(
                    npy,
                    os.path.join(vocoder_stems_dir, 'instrumental.mp3'), # vocoder output to vocoder_stems_dir
                    rescale,
                    None, # Removed args, use default vocoder args
                    inst_decoder,
                    codec_model
                )
            else:
                # Process vocal
                vocal_output = process_audio(
                    npy,
                    os.path.join(vocoder_stems_dir, 'vocal.mp3'), # vocoder output to vocoder_stems_dir
                    rescale,
                    None, # Removed args, use default vocoder args
                    vocal_decoder,
                    codec_model
                )

        # mix tracks from vocoder output
        try:
            mix_output = instrumental_output + vocal_output
            vocoder_mix = os.path.join(recons_mix_dir, 'mixed_stage2_vocoder.mp3') # mixed output in recons_stage2_vocoder, changed filename
            save_audio(mix_output, vocoder_mix, 44100, rescale)
            print(f"Created mix: {vocoder_mix}")
            tracks_stage2_vocoder.append(vocoder_mix) # add mixed vocoder output path
        except RuntimeError as e:
            print(e)
            vocoder_mix = None # set to None if mix failed
            print(f"mix {vocoder_mix} failed! inst: {instrumental_output.shape if instrumental_output is not None else 'None'}, vocal: {vocal_output.shape if vocal_output is not None else 'None'}")


        # mix tracks from stage 1
        mixed_stage1_path = None
        vocal_stage1_path = None
        instrumental_stage1_path = None
        for inst_path in tracks_stage1: # changed variable name to tracks_stage1
            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_stage1', 'mixed_stage1')) # changed mixed filename
                    vocal_stem, sr = sf.read(vocal_path)
                    instrumental_stem, _ = sf.read(inst_path)
                    mix_stem = (vocal_stem + instrumental_stem) / 1

                    sf.write(recons_mix, mix_stem, sr)
                    mixed_stage1_path = recons_mix # store mixed stage 1 path
                    vocal_stage1_path = vocal_path # store vocal stage 1 path
                    instrumental_stage1_path = inst_path # store instrumental stage 1 path

            except Exception as e:
                print(e)


        # Post process - skip post process for gradio to simplify.
        # recons_mix_final_path = os.path.join(output_dir, os.path.basename(mixed_stage1_path).replace('_stage1', '_final')) # final output path
        # replace_low_freq_with_energy_matched(
        #     a_file=mixed_stage1_path,     # 16kHz
        #     b_file=vocoder_mix,     # 48kHz
        #     c_file=recons_mix_final_path,
        #     cutoff_freq=5500.0
        # )


        if vocoder_mix is not None: # return vocoder mix if successful
            mixed_audio_data, sr_vocoder_mix = sf.read(vocoder_mix)
            vocal_audio_data = None # stage 2 vocoder stems are not mixed and returned in this demo, set to None
            instrumental_audio_data = None # stage 2 vocoder stems are not mixed and returned in this demo, set to None
            return (sr_vocoder_mix, (mixed_audio_data * 32767).astype(np.int16)), vocal_audio_data, instrumental_audio_data
        elif mixed_stage1_path is not None: # if vocoder failed, return stage 1 mix
            mixed_audio_data_stage1, sr_stage1_mix = sf.read(mixed_stage1_path)
            vocal_audio_data_stage1, sr_vocal_stage1 = sf.read(vocal_stage1_path)
            instrumental_audio_data_stage1, sr_inst_stage1 = sf.read(instrumental_stage1_path)
            return (sr_stage1_mix, (mixed_audio_data_stage1 * 32767).astype(np.int16)), (sr_vocal_stage1, (vocal_audio_data_stage1 * 32767).astype(np.int16)), (sr_inst_stage1, (instrumental_audio_data_stage1 * 32767).astype(np.int16))
        else: # if both failed, return None
             return None, None, None


def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=15):
    # Execute the command
    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 Occured: " + str(e))
        return None, None, None
    finally:
        print("Temporary files deleted.")


# Gradio
with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation")
        gr.HTML("""
        <div style="display:flex;column-gap:4px;">
            <a href="https://github.com/multimodal-art-projection/YuE">
                <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
            </a>
            <a href="https://map-yue.github.io">
                <img src='https://img.shields.io/badge/Project-Page-green'>
            </a>
            <a href="https://huggingface.co/spaces/innova-ai/YuE-music-generator-demo?duplicate=true">
                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
            </a>
        </div>
        """)
        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 (Stage 2 + Vocoder)")
                with gr.Accordion(label="Stage 1 Vocal and Instrumental Result", open=False):
                    vocal_out = gr.Audio(label="Vocal Audio (Stage 1)")
                    instrumental_out = gr.Audio(label="Instrumental Audio (Stage 1)")

        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'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, 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)