Update app.py

app.py

@@ -8,23 +8,21 @@ import torch
 import sys
 import uuid
 import re
-import numpy as np
-import json
-import time
-import copy
-from collections import Counter
 
-# Install flash-attn and set environment variable to skip cuda build
 print("Installing flash-attn...")
+# Install flash attention
 subprocess.run(
     "pip install flash-attn --no-build-isolation",
     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
     shell=True
 )
 
-# Download snapshot from huggingface_hub
 from huggingface_hub import snapshot_download
+
+# 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}")
@@ -33,10 +31,10 @@ else:
 
 snapshot_download(
     repo_id="m-a-p/xcodec_mini_infer",
-    local_dir=folder_path
+    local_dir="./xcodec_mini_infer"
 )
 
-# Change working directory to current folder
+# Change to the "inference" directory
 inference_dir = "."
 try:
     os.chdir(inference_dir)
@@ -45,179 +43,179 @@ except FileNotFoundError:
     print(f"Directory not found: {inference_dir}")
     exit(1)
 
-# Append necessary module paths
-base_path = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(os.path.join(base_path, 'xcodec_mini_infer'))
-sys.path.append(os.path.join(base_path, 'xcodec_mini_infer', 'descriptaudiocodec'))
+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
 
-# Other imports
+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 # removed vocoder
+#from post_process_audio import replace_low_freq_with_energy_matched # removed post process
 
-# Device setup
 device = "cuda:0"
 
-# Load and (optionally) compile the LM model
 model = AutoModelForCausalLM.from_pretrained(
     "m-a-p/YuE-s1-7B-anneal-en-cot",
     torch_dtype=torch.float16,
     attn_implementation="flash_attention_2",
+    # low_cpu_mem_usage=True,
 ).to(device)
 model.eval()
-try:
-    # torch.compile is available in PyTorch 2.0+
-    model = torch.compile(model)
-except Exception as e:
-    print("torch.compile not used for model:", e)
 
-# File paths for codec model checkpoint
-basic_model_config = os.path.join(folder_path, 'final_ckpt/config.yaml')
-resume_path = os.path.join(folder_path, 'final_ckpt/ckpt_00360000.pth')
+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' # removed vocoder
+#vocal_decoder_path = './xcodec_mini_infer/decoders/decoder_131000.pth' # removed vocoder
+#inst_decoder_path = './xcodec_mini_infer/decoders/decoder_151000.pth' # removed vocoder
 
-# Initialize tokenizer and codec manipulator
 mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-codectool = CodecManipulator("xcodec", 0, 1)
 
-# Load codec model config and initialize codec model
+codectool = CodecManipulator("xcodec", 0, 1)
 model_config = OmegaConf.load(basic_model_config)
-# Dynamically create the model from its name in the config.
-codec_class = eval(model_config.generator.name)
-codec_model = codec_class(**model_config.generator.config).to(device)
+# 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 = torch.compile(codec_model)
 codec_model.eval()
-try:
-    codec_model = torch.compile(codec_model)
-except Exception as e:
-    print("torch.compile not used for codec_model:", e)
 
-# Pre-compile the regex pattern for splitting lyrics
-LYRICS_PATTERN = re.compile(r"\[(\w+)\](.*?)\n(?=\[|\Z)", re.DOTALL)
+# Preload and compile vocoders # removed vocoder
+#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 = torch.compile(vocal_decoder)
+#inst_decoder = torch.compile(inst_decoder)
+#vocal_decoder.eval()
+#inst_decoder.eval()
 
-# ------------------ GPU decorated generation function ------------------ #
-@spaces.GPU(duration=175)
+
+@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,
+        max_new_tokens=5,
+        run_n_segments=2,
+        genre_txt=None,
+        lyrics_txt=None,
+        use_audio_prompt=False,
+        audio_prompt_path="",
+        prompt_start_time=0.0,
+        prompt_end_time=30.0,
+        cuda_idx=0,
+        rescale=False,
 ):
     if use_audio_prompt and not audio_prompt_path:
-        raise FileNotFoundError("Please provide an audio prompt filepath when 'use_audio_prompt' is enabled!")
-    max_new_tokens = max_new_tokens * 50  # scaling factor
+        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, "stage1")
+        stage1_output_dir = os.path.join(output_dir, f"stage1")
         os.makedirs(stage1_output_dir, exist_ok=True)
 
-        # -- In-place logits processor that blocks token ranges --
         class BlockTokenRangeProcessor(LogitsProcessor):
             def __init__(self, start_id, end_id):
-                # Pre-create a tensor for indices if possible
                 self.blocked_token_ids = list(range(start_id, end_id))
+
             def __call__(self, input_ids, scores):
                 scores[:, self.blocked_token_ids] = -float("inf")
                 return scores
 
-        # -- Audio processing utility --
         def load_audio_mono(filepath, sampling_rate=16000):
             audio, sr = torchaudio.load(filepath)
-            audio = audio.mean(dim=0, keepdim=True)  # convert to mono
+            # 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
 
-        # -- Lyrics splitting using precompiled regex --
         def split_lyrics(lyrics: str):
-            segments = LYRICS_PATTERN.findall(lyrics)
-            # Return segments with formatting (strip extra whitespace)
-            return [f"[{tag}]\n{text.strip()}\n\n" for tag, text in segments]
-
-        # Prepare prompt texts
-        genres = genre_txt.strip() if genre_txt else ""
-        lyrics_segments = split_lyrics(lyrics_txt + "\n")
-        full_lyrics = "\n".join(lyrics_segments)
-        # The first prompt is a global instruction; the rest are segments.
+            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
+
+        # Call the function and print the result
+        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_segments
+        prompt_texts += lyrics
 
         random_id = uuid.uuid4()
-        raw_output = None
-
-        # Decoding config parameters
+        output_seq = None
+        # Here is suggested decoding config
         top_p = 0.93
         temperature = 1.0
         repetition_penalty = 1.2
-
-        # Pre-tokenize static tokens
+        # special tokens
         start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
         end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
-        soa_token = mmtokenizer.soa  # start-of-audio token id
-        eoa_token = mmtokenizer.eoa  # end-of-audio token id
-
-        # Pre-tokenize the global prompt (first element)
-        global_prompt_ids = mmtokenizer.tokenize(prompt_texts[0])
-        run_n_segments = min(run_n_segments + 1, len(prompt_texts))
-        
-        # Loop over segments. (Note: Each segment is processed sequentially.)
-        for i, p in enumerate(tqdm(prompt_texts[:run_n_segments], desc="Generating segments")):
-            # Remove any spurious tokens in the text
+
+        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:
-                # Skip generation on the instruction segment.
                 continue
-
-            # Build prompt IDs differently depending on whether audio prompt is enabled.
             if i == 1:
                 if use_audio_prompt:
                     audio_prompt = load_audio_mono(audio_prompt_path)
-                    audio_prompt = audio_prompt.unsqueeze(0)
-                    with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.float16):
+                    audio_prompt.unsqueeze_(0)
+                    with torch.no_grad():
                         raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=0.5)
-                    # Process raw codes (transpose and convert to numpy)
-                    raw_codes = raw_codes.transpose(0, 1).cpu().numpy().astype(np.int16)
+                    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])
-                    # Slice using prompt start/end time (assuming 50 tokens per second)
-                    audio_prompt_codec = code_ids[int(prompt_start_time * 50): int(prompt_end_time * 50)]
-                    audio_prompt_codec_ids = [soa_token] + codectool.sep_ids + audio_prompt_codec + [eoa_token]
-                    sentence_ids = mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
-                    head_id = global_prompt_ids + sentence_ids
+                    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 = global_prompt_ids
-                prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [soa_token] + codectool.sep_ids
-            else:
-                prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [soa_token] + codectool.sep_ids
-
-            prompt_ids_tensor = torch.as_tensor(prompt_ids, device=device).unsqueeze(0)
-            if raw_output is not None:
-                # Concatenate previous outputs with the new prompt
-                input_ids = torch.cat([raw_output, prompt_ids_tensor], dim=1)
+                    head_id = mmtokenizer.tokenize(prompt_texts[0])
+                prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
             else:
-                input_ids = prompt_ids_tensor
+                prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
 
-            # Enforce maximum context window by slicing if needed
+            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:
-                input_ids = input_ids[:, -max_context:]
-
-            with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.float16):
+                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,
@@ -226,149 +224,140 @@ def generate_music(
                     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)
-                    ]),
+                    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
                 )
-                # Ensure the output ends with an end-of-audio token
-                if output_seq[0, -1].item() != eoa_token:
-                    tensor_eoa = torch.as_tensor([[eoa_token]], device=device)
+                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)
-            # For subsequent segments, append only the newly generated tokens.
-            if raw_output is not None:
-                new_tokens = output_seq[:, input_ids.shape[-1]:]
-                raw_output = torch.cat([raw_output, prompt_ids_tensor, new_tokens], dim=1)
+            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 codec tokens to temporary files and check token pairs.
+        # save raw output and check sanity
         ids = raw_output[0].cpu().numpy()
-        soa_idx = np.where(ids == soa_token)[0]
-        eoa_idx = np.where(ids == eoa_token)[0]
+        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_list = []
-        instrumentals_list = []
-        # If using an audio prompt, skip the first pair (it may be reference)
-        start_idx = 1 if use_audio_prompt else 0
-        for i in range(start_idx, len(soa_idx)):
-            codec_ids = ids[soa_idx[i] + 1: eoa_idx[i]]
+            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:]
-            # Force even length and reshape into 2 channels.
-            codec_ids = codec_ids[:2 * (len(codec_ids) // 2)]
-            codec_ids = np.array(codec_ids)
-            reshaped = rearrange(codec_ids, "(n b) -> b n", b=2)
-            vocals_list.append(codectool.ids2npy(reshaped[0]))
-            instrumentals_list.append(codectool.ids2npy(reshaped[1]))
-        vocals = np.concatenate(vocals_list, axis=1)
-        instrumentals = np.concatenate(instrumentals_list, axis=1)
-
-        # Save the numpy arrays to temporary files
-        vocal_save_path = os.path.join(stage1_output_dir, f"vocal_{str(random_id).replace('.', '@')}.npy")
-        inst_save_path = os.path.join(stage1_output_dir, f"instrumental_{str(random_id).replace('.', '@')}.npy")
+            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 = [vocal_save_path, inst_save_path]
+        stage1_output_set.append(vocal_save_path)
+        stage1_output_set.append(inst_save_path)
 
         print("Converting to Audio...")
 
-        # Utility function for saving audio with in-place clipping
+        # convert audio tokens to audio
         def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
-            os.makedirs(os.path.dirname(path), exist_ok=True)
+            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().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)
+            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 by decoding codec tokens
+        # reconstruct tracks
         recons_output_dir = os.path.join(output_dir, "recons")
-        recons_mix_dir = os.path.join(recons_output_dir, "mix")
+        recons_mix_dir = os.path.join(recons_output_dir, 'mix')
         os.makedirs(recons_mix_dir, exist_ok=True)
         tracks = []
-        for npy_path in stage1_output_set:
-            codec_result = np.load(npy_path)
-            with torch.inference_mode():
-                # Adjust shape: (1, T, C) expected by the decoder
-                input_tensor = torch.as_tensor(codec_result.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device)
-                decoded_waveform = codec_model.decode(input_tensor)
+        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)
-            save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy_path))[0] + ".mp3")
+            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] + ".mp3")
             tracks.append(save_path)
-            save_audio(decoded_waveform, save_path, sample_rate=16000)
-
-        # Mix vocal and instrumental tracks (using torch to avoid extra I/O if possible)
+            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 'instrumental' in inst_path:
+                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
-                    # Read using soundfile
-                    vocal_stem, sr = sf.read(vocal_path)
-                    instrumental_stem, _ = sf.read(inst_path)
-                    mix_stem = (vocal_stem + instrumental_stem) / 1.0
-                    mix_path = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace('instrumental', 'mixed'))
-                    # Write the mix to disk (if needed) or return in memory
-                    # Here we return three tuples: (sr, mix), (sr, vocal), (sr, instrumental)
+                    # mix
+                    recons_mix = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace('instrumental', 'mixed'))
+                    vocal_stem, sr = sf.read(inst_path)
+                    instrumental_stem, _ = sf.read(vocal_path)
+                    mix_stem = (vocal_stem + instrumental_stem) / 1
                     return (sr, (mix_stem * 32767).astype(np.int16)), (sr, (vocal_stem * 32767).astype(np.int16)), (sr, (instrumental_stem * 32767).astype(np.int16))
             except Exception as e:
-                print("Mixing error:", e)
+                print(e)
                 return None, None, None
 
-# ------------------ Inference function and Gradio UI ------------------ #
-def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=23):
+
+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
-        )
+        mixed_audio_data, vocal_audio_data, instrumental_audio_data = generate_music(genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments,
+                               cuda_idx=0, max_new_tokens=max_new_tokens)
         return mixed_audio_data, vocal_audio_data, instrumental_audio_data
     except Exception as e:
-        gr.Warning("An Error Occurred: " + str(e))
+        gr.Warning("An Error Occured: " + str(e))
         return None, None, None
     finally:
         print("Temporary files deleted.")
 
-# Build Gradio UI
+
+# 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>
-            """
-        )
+        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")
                 with gr.Accordion(label="Vocal and Instrumental Result", open=False):
                     vocal_out = gr.Audio(label="Vocal Audio")