Update app.py

app.py

@@ -8,21 +8,23 @@ 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}")
@@ -31,10 +33,10 @@ else:
 
 snapshot_download(
     repo_id="m-a-p/xcodec_mini_infer",
-    local_dir="./xcodec_mini_infer"
+    local_dir=folder_path
 )
 
-# Change to the "inference" directory
+# Change working directory to current folder
 inference_dir = "."
 try:
     os.chdir(inference_dir)
@@ -43,179 +45,178 @@ 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
+# 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'))
 
-import argparse
-import numpy as np
-import json
+# Other imports
 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)
 
-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
+# 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')
 
+# 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
 model_config = OmegaConf.load(basic_model_config)
-# Load codec model
-codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
+# 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)
 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)
 
-# 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()
-
+# Pre-compile the regex pattern for splitting lyrics
+LYRICS_PATTERN = re.compile(r"\[(\w+)\](.*?)\n(?=\[|\Z)", re.DOTALL)
 
+# ------------------ GPU decorated generation function ------------------ #
 @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 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
+        raise FileNotFoundError("Please provide an audio prompt filepath when 'use_audio_prompt' is enabled!")
+    max_new_tokens = max_new_tokens * 100  # scaling factor
 
     with tempfile.TemporaryDirectory() as output_dir:
-        stage1_output_dir = os.path.join(output_dir, f"stage1")
+        stage1_output_dir = os.path.join(output_dir, "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)
-            # Convert to mono
-            audio = torch.mean(audio, dim=0, keepdim=True)
-            # Resample if needed
+            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
 
+        # -- Lyrics splitting using precompiled regex --
         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
-
-        # 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)
+            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.
         prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"]
-        prompt_texts += lyrics
+        prompt_texts += lyrics_segments
 
         random_id = uuid.uuid4()
-        output_seq = None
-        # Here is suggested decoding config
+        raw_output = None
+
+        # Decoding config parameters
         top_p = 0.93
         temperature = 1.0
         repetition_penalty = 1.2
-        # special tokens
+
+        # Pre-tokenize static 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])):
+        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
             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.unsqueeze_(0)
-                    with torch.no_grad():
+                    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)
-                    raw_codes = raw_codes.cpu().numpy().astype(np.int16)
-                    # Format audio prompt
+                    # Process raw codes (transpose and convert to numpy)
+                    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)]  # 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
+                    # 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
                 else:
-                    head_id = mmtokenizer.tokenize(prompt_texts[0])
-                prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+                    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)
             else:
-                prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+                input_ids = prompt_ids_tensor
 
-            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
+            # Enforce maximum context window by slicing if needed
             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):
+                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,
@@ -224,140 +225,149 @@ def generate_music(
                     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)]),
+                    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() != mmtokenizer.eoa:
-                    tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
+                # 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)
                     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)
+            # 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)
             else:
                 raw_output = output_seq
-            print(len(raw_output))
 
-        # save raw output and check sanity
+        # Save raw output codec tokens to temporary files and check token pairs.
         ids = raw_output[0].cpu().numpy()
-        soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
-        eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
+        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, 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]]
+            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]]
             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')
+            # 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")
         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)
+        stage1_output_set = [vocal_save_path, inst_save_path]
 
         print("Converting to Audio...")
 
-        # convert audio tokens to audio
+        # Utility function for saving audio with in-place clipping
         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)
+            os.makedirs(os.path.dirname(path), exist_ok=True)
             limit = 0.99
-            max_val = wav.abs().max()
-            wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
-            torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)
+            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)
 
-        # reconstruct tracks
+        # Reconstruct tracks by decoding codec tokens
         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 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))
+        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)
             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] + ".mp3")
+            save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy_path))[0] + ".mp3")
             tracks.append(save_path)
-            save_audio(decodec_rlt, save_path, 16000)
-        # mix tracks
+            save_audio(decoded_waveform, save_path, sample_rate=16000)
+
+        # Mix vocal and instrumental tracks (using torch to avoid extra I/O if possible)
         for inst_path in tracks:
             try:
-                if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) \
-                        and 'instrumental' in inst_path:
-                    # find pair
+                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
-                    # 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
+                    # 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)
                     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(e)
+                print("Mixing error:", e)
                 return None, None, None
 
-
+# ------------------ Inference function and Gradio UI ------------------ #
 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 Occured: " + str(e))
+        gr.Warning("An Error Occurred: " + str(e))
         return None, None, None
     finally:
         print("Temporary files deleted.")
 
-
-# Gradio
+# Build Gradio UI
 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")
@@ -399,7 +409,7 @@ 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
+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
@@ -419,4 +429,5 @@ Living out my dreams with this mic and a deal
         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)
+    
+demo.queue().launch(show_error=True)