import gradio as gr
import subprocess
import os
import spaces
import sys
import shutil
import tempfile
import uuid
import re
import time
import copy
from collections import Counter
from tqdm import tqdm
from einops import rearrange
import numpy as np
import json

import torch
import torchaudio
from torchaudio.transforms import Resample
import soundfile as sf

# --- Install flash-attn (if needed) ---
print("Installing flash-attn...")
subprocess.run(
    "pip install flash-attn --no-build-isolation",
    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
    shell=True
)

# --- Download and set up stage1 files ---
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 required 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"))

# --- Additional imports (vocoder & post processing) ---
from omegaconf import OmegaConf
from codecmanipulator import CodecManipulator
from mmtokenizer import _MMSentencePieceTokenizer
from transformers import AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
from models.soundstream_hubert_new import SoundStream

# Import vocoder functions (ensure these modules exist)
from vocoder import build_codec_model, process_audio
from post_process_audio import replace_low_freq_with_energy_matched

# ----------------------- Global Configuration -----------------------
# Stage1 and Stage2 model identifiers (change if needed)
STAGE1_MODEL = "m-a-p/YuE-s1-7B-anneal-en-cot"
STAGE2_MODEL = "m-a-p/YuE-s2-1B-general"
# Vocoder model files (paths in the xcodec snapshot)
BASIC_MODEL_CONFIG = os.path.join(folder_path, "final_ckpt/config.yaml")
RESUME_PATH = os.path.join(folder_path, "final_ckpt/ckpt_00360000.pth")
VOCAL_DECODER_PATH = os.path.join(folder_path, "decoders/decoder_131000.pth")
INST_DECODER_PATH = os.path.join(folder_path, "decoders/decoder_151000.pth")
VOCODER_CONFIG_PATH = os.path.join(folder_path, "decoders/config.yaml")

# Misc settings
MAX_NEW_TOKENS = 15         # Duration slider (in seconds, scaled internally)
RUN_N_SEGMENTS = 2          # Number of segments to generate
STAGE2_BATCH_SIZE = 4       # Batch size for stage2 inference

# You may change these defaults via Gradio input (see below)

# ----------------------- Device Setup -----------------------
device = "cuda:0" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# ----------------------- Load Stage1 Models and Tokenizer -----------------------
print("Loading Stage 1 model and tokenizer...")
model = AutoModelForCausalLM.from_pretrained(
    STAGE1_MODEL,
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
).to(device)
model.eval()

model_stage2 = AutoModelForCausalLM.from_pretrained(
    STAGE2_MODEL,
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
).to(device)
model_stage2.eval()

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

# Two separate codec manipulators: one for Stage1 and one for Stage2 (with a higher number of quantizers)
codectool = CodecManipulator("xcodec", 0, 1)
codectool_stage2 = CodecManipulator("xcodec", 0, 8)

# Load codec (xcodec) model for Stage1 & Stage2 decoding
model_config = OmegaConf.load(BASIC_MODEL_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()

# Precompile regex for splitting lyrics
LYRICS_PATTERN = re.compile(r"\[(\w+)\](.*?)\n(?=\[|\Z)", re.DOTALL)

# ----------------------- Utility Functions -----------------------
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

def split_lyrics(lyrics: str):
    segments = LYRICS_PATTERN.findall(lyrics)
    return [f"[{tag}]\n{text.strip()}\n\n" for tag, text in segments]

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

# ----------------------- Stage2 Functions -----------------------
def stage2_generate(model_stage2, prompt, batch_size=16):
    """
    Given a prompt (a numpy array of raw codec ids), upsample using the Stage2 model.
    """
    # Unflatten prompt: assume prompt shape (1, T) and then reformat.
    print(f"stage2_generate: received prompt with shape: {prompt.shape}")
    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)

    # Build new prompt tokens for Stage2:
    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_concat = np.concatenate(codec_list, axis=0)
        prompt_ids = np.concatenate(
            [
                np.tile([mmtokenizer.soa, mmtokenizer.stage_1], (batch_size, 1)),
                codec_ids_concat,
                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(),
                np.array([mmtokenizer.stage_2]),
            ]
        ).astype(np.int32)
        prompt_ids = prompt_ids[np.newaxis, ...]

    codec_ids_tensor = torch.as_tensor(codec_ids).to(device)
    prompt_ids_tensor = torch.as_tensor(prompt_ids).to(device)
    len_prompt = prompt_ids_tensor.shape[-1]

    block_list = LogitsProcessorList([
        BlockTokenRangeProcessor(0, 46358),
        BlockTokenRangeProcessor(53526, mmtokenizer.vocab_size)
    ])

    # Teacher forcing generate loop: generate tokens in fixed 7-token steps per frame.
    for frames_idx in range(codec_ids_tensor.shape[1]):
        cb0 = codec_ids_tensor[:, frames_idx:frames_idx+1]
        prompt_ids_tensor = torch.cat([prompt_ids_tensor, cb0], dim=1)
        with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.float16):
            stage2_output = model_stage2.generate(
                input_ids=prompt_ids_tensor,
                min_new_tokens=7,
                max_new_tokens=7,
                eos_token_id=mmtokenizer.eoa,
                pad_token_id=mmtokenizer.eoa,
                logits_processor=block_list,
                use_cache=True
            )
        # Ensure exactly 7 new tokens were added.
        assert stage2_output.shape[1] - prompt_ids_tensor.shape[1] == 7, (
            f"output new tokens={stage2_output.shape[1]-prompt_ids_tensor.shape[1]}"
        )
        prompt_ids_tensor = stage2_output

    # Return new tokens (excluding prompt)
    if batch_size > 1:
        output = prompt_ids_tensor.cpu().numpy()[:, len_prompt:]
        # If desired, reshape/split per batch element
        output_list = [output[i] for i in range(batch_size)]
        output = np.concatenate(output_list, axis=0)
    else:
        output = prompt_ids_tensor[0].cpu().numpy()[len_prompt:]
    return output

def stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_size=4):
    stage2_result = []
    for path in tqdm(stage1_output_set, desc="Stage2 Inference"):
        output_filename = os.path.join(stage2_output_dir, os.path.basename(path))
        if os.path.exists(output_filename):
            print(f"{output_filename} already processed.")
            stage2_result.append(output_filename)
            continue

        prompt = np.load(path).astype(np.int32)
        # Ensure prompt is 2D.
        if prompt.ndim == 1:
            prompt = prompt[np.newaxis, :]
        print(f"Loaded prompt from {path} with shape: {prompt.shape}")

        # Compute total duration in seconds (assuming 50 tokens per second)
        total_duration_sec = prompt.shape[-1] // 50
        if total_duration_sec < 6:
            # Not enough tokens for a full 6-sec segment; use the entire prompt.
            output_duration = total_duration_sec
            print(f"Prompt too short for 6-sec segmentation. Using full duration: {output_duration} seconds.")
        else:
            output_duration = (total_duration_sec // 6) * 6

        # If after the above, output_duration is still zero, raise an error.
        if output_duration == 0:
            raise ValueError(f"Output duration computed as 0 for {path}. Prompt length: {prompt.shape[-1]} tokens")
        
        num_batch = output_duration // 6

        # Process prompt in batches
        if num_batch <= batch_size:
            output = stage2_generate(model_stage2, prompt[:, :output_duration*50], batch_size=num_batch)
        else:
            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
                end_idx = min((seg + 1) * batch_size * 300, output_duration * 50)
                current_batch = batch_size if (seg != num_segments - 1 or num_batch % batch_size == 0) else num_batch % batch_size
                segment_prompt = prompt[:, start_idx:end_idx]
                if segment_prompt.shape[-1] == 0:
                    print(f"Warning: empty segment detected for seg {seg}, start {start_idx}, end {end_idx}. Skipping this segment.")
                    continue
                segment = stage2_generate(model_stage2, segment_prompt, batch_size=current_batch)
                segments.append(segment)
            if len(segments) == 0:
                raise ValueError(f"No valid segments produced for {path}.")
            output = np.concatenate(segments, axis=0)

        # Process any remaining tokens if prompt length not fully used.
        if output_duration * 50 != prompt.shape[-1]:
            ending = stage2_generate(model_stage2, prompt[:, output_duration * 50:], batch_size=1)
            output = np.concatenate([output, ending], axis=0)

        # Convert Stage2 output tokens back to numpy using Stage2’s codec manipulator.
        output = codectool_stage2.ids2npy(output)

        # Fix any invalid codes (if needed)
        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_common = sorted(counter.items(), key=lambda x: x[1], reverse=True)[0][0]
                    fixed_output[i, j] = most_common

        np.save(output_filename, fixed_output)
        stage2_result.append(output_filename)
    return stage2_result

# ----------------------- Main Generation Function (Stage1 + Stage2) -----------------------
@spaces.GPU(duration=175)
def generate_music(
    genre_txt="",
    lyrics_txt="",
    max_new_tokens=2,
    run_n_segments=1,
    use_audio_prompt=False,
    audio_prompt_path="",
    prompt_start_time=0.0,
    prompt_end_time=30.0,
    rescale=False,
):
    # Scale max_new_tokens (e.g. seconds * 50 tokens per second)
    max_new_tokens_scaled = max_new_tokens * 50

    # Use a temporary directory to store intermediate stage outputs.
    with tempfile.TemporaryDirectory() as tmp_dir:
        stage1_output_dir = os.path.join(tmp_dir, "stage1")
        stage2_output_dir = os.path.join(tmp_dir, "stage2")
        os.makedirs(stage1_output_dir, exist_ok=True)
        os.makedirs(stage2_output_dir, exist_ok=True)

        # ---------------- Stage 1: Text-to-Music Generation ----------------
        genres = genre_txt.strip()
        lyrics_segments = split_lyrics(lyrics_txt + "\n")
        full_lyrics = "\n".join(lyrics_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
        top_p = 0.93
        temperature = 1.0
        repetition_penalty = 1.2

        # Pre-tokenize special tokens
        start_of_segment = mmtokenizer.tokenize("[start_of_segment]")
        end_of_segment = mmtokenizer.tokenize("[end_of_segment]")
        soa_token = mmtokenizer.soa
        eoa_token = mmtokenizer.eoa

        global_prompt_ids = mmtokenizer.tokenize(prompt_texts[0])
        run_n = min(run_n_segments + 1, len(prompt_texts))
        for i, p in enumerate(tqdm(prompt_texts[:run_n], desc="Stage1 Generation")):
            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 = 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)
                    raw_codes = raw_codes.transpose(0, 1).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 = [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:
                input_ids = torch.cat([raw_output, prompt_ids_tensor], dim=1)
            else:
                input_ids = prompt_ids_tensor

            max_context = 16384 - max_new_tokens_scaled - 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_scaled,
                    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,
                )
                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)
            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 Stage1 outputs (vocal & instrumental) as npy files.
        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: {len(soa_idx)} vs {len(eoa_idx)}")
        vocals_list = []
        instrumentals_list = []
        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 * (len(codec_ids) // 2)]
            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)
        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]

        # (Optional) Offload Stage1 model from GPU to free memory.
        model.cpu()
        torch.cuda.empty_cache()

        # ---------------- Stage 2: Refinement/Upsampling ----------------
        print("Stage 2 inference...")
        
        stage2_result = stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_size=STAGE2_BATCH_SIZE)
        print("Stage 2 inference completed.")

        # ---------------- Reconstruct Audio from Stage2 Tokens ----------------
        recons_output_dir = os.path.join(tmp_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)
            with torch.inference_mode():
                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))[0] + ".mp3")
            tracks.append(save_path)
            save_audio(decoded_waveform, save_path, 16000, rescale)
        # Mix vocal and instrumental tracks:
        mix_audio = None
        vocal_audio = None
        instrumental_audio = None
        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
                    vocal_data, sr = sf.read(vocal_path)
                    instrumental_data, _ = sf.read(inst_path)
                    mix_data = (vocal_data + instrumental_data) / 1.0
                    recons_mix = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace("instrumental", "mixed"))
                    sf.write(recons_mix, mix_data, sr)
                    mix_audio = (sr, (mix_data * 32767).astype(np.int16))
                    vocal_audio = (sr, (vocal_data * 32767).astype(np.int16))
                    instrumental_audio = (sr, (instrumental_data * 32767).astype(np.int16))
            except Exception as e:
                print("Mixing error:", e)
                return None, None, None

        # ---------------- Vocoder Upsampling and Post Processing ----------------
        print("Vocoder upsampling...")
        vocal_decoder, inst_decoder = build_codec_model(VOCODER_CONFIG_PATH, VOCAL_DECODER_PATH, INST_DECODER_PATH)
        vocoder_output_dir = os.path.join(tmp_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_stems_dir, exist_ok=True)
        os.makedirs(vocoder_mix_dir, exist_ok=True)
        # Process each track with the vocoder (here we process vocal and instrumental separately)
        if vocal_audio is not None and instrumental_audio is not None:
            vocal_output = process_audio(
                stage2_result[0],
                os.path.join(vocoder_stems_dir, "vocal.mp3"),
                rescale,
                None,
                vocal_decoder,
                codec_model,
            )
            instrumental_output = process_audio(
                stage2_result[1],
                os.path.join(vocoder_stems_dir, "instrumental.mp3"),
                rescale,
                None,
                inst_decoder,
                codec_model,
            )
            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, rescale)
                print(f"Created vocoder mix: {vocoder_mix}")
            except RuntimeError as e:
                print(e)
                print("Mixing vocoder outputs failed!")
        else:
            print("Missing vocal/instrumental outputs for vocoder stage.")

        # Post-process: Replace low frequency of Stage1 reconstruction with energy-matched vocoder mix.
        final_mix_path = os.path.join(tmp_dir, "final_mix.mp3")
        try:
            replace_low_freq_with_energy_matched(
                a_file=recons_mix,      # Stage1 mix at 16kHz
                b_file=vocoder_mix,      # Vocoder mix at 48kHz
                c_file=final_mix_path,
                cutoff_freq=5500.0
            )
        except Exception as e:
            print("Post processing error:", e)
            final_mix_path = recons_mix  # Fall back to Stage1 mix

        # Return final outputs as tuples: (sample_rate, np.int16 audio)
        final_audio, vocal_audio, instrumental_audio = None, None, None
        try:
            final_audio_data, sr = sf.read(final_mix_path)
            final_audio = (sr, (final_audio_data * 32767).astype(np.int16))
        except Exception as e:
            print("Final mix read error:", e)
        return final_audio, vocal_audio, instrumental_audio

# ----------------------- Gradio Interface -----------------------
with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# YuE: Full-Song Generation (Stage1 + Stage2)")
        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>
            </div>
            """
        )
        with gr.Row():
            with gr.Column():
                genre_txt = gr.Textbox(label="Genre", placeholder="e.g. Bass Metalcore Thrash Metal Furious bright vocal male")
                lyrics_txt = gr.Textbox(label="Lyrics", placeholder="Paste lyrics with segments such as [verse], [chorus], etc.")
            with gr.Column():
                num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
                max_new_tokens = gr.Slider(label="Duration of song (sec)", minimum=1, maximum=30, step=1, value=15, interactive=True)
                use_audio_prompt = gr.Checkbox(label="Use Audio Prompt", value=False)
                audio_prompt_path = gr.Textbox(label="Audio Prompt Filepath (if used)", placeholder="Path to audio file")
                submit_btn = gr.Button("Submit")
                music_out = gr.Audio(label="Mixed Audio Result")
                with gr.Accordion(label="Vocal and Instrumental Results", 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=generate_music
        )
    submit_btn.click(
        fn=generate_music,
        inputs=[genre_txt, lyrics_txt, max_new_tokens, num_segments, use_audio_prompt, audio_prompt_path],
        outputs=[music_out, vocal_out, instrumental_out]
    )
    gr.Markdown("## Contributions Welcome\nFeel free to contribute improvements or fixes.")

demo.queue().launch(show_error=True)