modified: app.py

app.py

@@ -7,7 +7,7 @@ import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
 import torch
 
-is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False 
+is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False
 
 # Install required package
 def install_flash_attn():
@@ -49,65 +49,83 @@ import sys
 sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
 sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
 
-
-import os
-import sys
-import torch
+import argparse
 import numpy as np
 import json
-import re
-import uuid
-import gradio as gr
-from tqdm import tqdm
 from omegaconf import OmegaConf
 import torchaudio
 from torchaudio.transforms import Resample
 import soundfile as sf
+
+import uuid
+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
+import re
 
-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'))
-from codecmanipulator import CodecManipulator
-from mmtokenizer import _MMSentencePieceTokenizer
-
-# Load models once at startup
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load language model
-print("Loading language model...")
-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()
+# --- Arguments and Model Loading from infer.py ---
+parser = argparse.ArgumentParser()
+# Model Configuration:
+parser.add_argument("--stage1_model", type=str, default="m-a-p/YuE-s1-7B-anneal-en-cot", help="The model checkpoint path or identifier for the Stage 1 model.")
+parser.add_argument("--max_new_tokens", type=int, default=3000, help="The maximum number of new tokens to generate in one pass during text generation.")
+parser.add_argument("--run_n_segments", type=int, default=2, help="The number of segments to process during the generation.")
+# Prompt
+parser.add_argument("--genre_txt", type=str, default="", help="The file path to a text file containing genre tags that describe the musical style or characteristics (e.g., instrumental, genre, mood, vocal timbre, vocal gender). This is used as part of the generation prompt.") # Modified: removed required=True and using default=""
+parser.add_argument("--lyrics_txt", type=str, default="", help="The file path to a text file containing the lyrics for the music generation. These lyrics will be processed and split into structured segments to guide the generation process.") # Modified: removed required=True and using default=""
+parser.add_argument("--use_audio_prompt", action="store_true", help="If set, the model will use an audio file as a prompt during generation. The audio file should be specified using --audio_prompt_path.")
+parser.add_argument("--audio_prompt_path", type=str, default="", help="The file path to an audio file to use as a reference prompt when --use_audio_prompt is enabled.")
+parser.add_argument("--prompt_start_time", type=float, default=0.0, help="The start time in seconds to extract the audio prompt from the given audio file.")
+parser.add_argument("--prompt_end_time", type=float, default=30.0, help="The end time in seconds to extract the audio prompt from the given audio file.")
+# Output
+parser.add_argument("--output_dir", type=str, default="./output", help="The directory where generated outputs will be saved.")
+parser.add_argument("--keep_intermediate", action="store_true", help="If set, intermediate outputs will be saved during processing.")
+parser.add_argument("--disable_offload_model", action="store_true", help="If set, the model will not be offloaded from the GPU to CPU after Stage 1 inference.")
+parser.add_argument("--cuda_idx", type=int, default=0)
+# Config for xcodec and upsampler
+parser.add_argument('--basic_model_config', default='./xcodec_mini_infer/final_ckpt/config.yaml', help='YAML files for xcodec configurations.')
+parser.add_argument('--resume_path', default='./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', help='Path to the xcodec checkpoint.')
+parser.add_argument('--config_path', type=str, default='./xcodec_mini_infer/decoders/config.yaml', help='Path to Vocos config file.')
+parser.add_argument('--vocal_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_131000.pth', help='Path to Vocos decoder weights.')
+parser.add_argument('--inst_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_151000.pth', help='Path to Vocos decoder weights.')
+parser.add_argument('-r', '--rescale', action='store_true', help='Rescale output to avoid clipping.')
+
+
+args = parser.parse_args([]) # Modified: Pass empty list to parse_args to avoid command line parsing in Gradio
+
+if args.use_audio_prompt and not args.audio_prompt_path:
+    raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
+model_name = args.stage1_model # Modified: Renamed 'model' to 'model_name' to avoid shadowing the loaded model later
+cuda_idx = args.cuda_idx
+max_new_tokens_config = args.max_new_tokens # Modified: Renamed 'max_new_tokens' to 'max_new_tokens_config' to avoid shadowing the Gradio input
+stage1_output_dir = os.path.join(args.output_dir, f"stage1")
+os.makedirs(stage1_output_dir, exist_ok=True)
+
+# load tokenizer and model
+device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
+
+# Now you can use `device` to move your tensors or models to the GPU (if available)
+print(f"Using device: {device}")
 
-# Load tokenizers and codec tools
-print("Loading tokenizers...")
 mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-codectool = CodecManipulator("xcodec", 0, 1)
 
-# Load codec models
-print("Loading codec models...")
-model_config = OmegaConf.load('./xcodec_mini_infer/final_ckpt/config.yaml')
+codectool = CodecManipulator("xcodec", 0, 1)
+model_config = OmegaConf.load(args.basic_model_config)
 codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
-parameter_dict = torch.load('./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', map_location='cpu')
+parameter_dict = torch.load(args.resume_path, map_location='cpu')
 codec_model.load_state_dict(parameter_dict['codec_model'])
 codec_model.to(device)
 codec_model.eval()
 
-# Load vocoders
-print("Loading vocoders...")
-vocal_decoder, inst_decoder = build_codec_model(
-    './xcodec_mini_infer/decoders/config.yaml',
-    './xcodec_mini_infer/decoders/decoder_131000.pth',
-    './xcodec_mini_infer/decoders/decoder_151000.pth'
-)
-
 class BlockTokenRangeProcessor(LogitsProcessor):
     def __init__(self, start_id, end_id):
         self.blocked_token_ids = list(range(start_id, end_id))
@@ -116,118 +134,325 @@ class BlockTokenRangeProcessor(LogitsProcessor):
         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):
     pattern = r"\[(\w+)\](.*?)\n(?=\[|\Z)"
     segments = re.findall(pattern, lyrics, re.DOTALL)
-    return [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
-
-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()
-    wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
-    torchaudio.save(path, wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)
-
-@spaces.GPU(duration=150)
-def run_inference(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=2000):
-    try:
-        # Create temporary output directory
-        output_dir = tempfile.mkdtemp()
-        stage1_output_dir = os.path.join(output_dir, "stage1")
-        os.makedirs(stage1_output_dir, exist_ok=True)
-
-        # Process inputs
-        structured_lyrics = split_lyrics(lyrics_txt_content)
-        full_lyrics = "\n".join(structured_lyrics)
-        prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genre_txt_content}\n{full_lyrics}"] + structured_lyrics
-
-        # Generation parameters
-        top_p = 0.93
-        temperature = 1.0
-        repetition_penalty = 1.2
-        start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
-        end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
-        run_n_segments = min(num_segments + 1, len(structured_lyrics))
-
-        # Generate tokens
-        raw_output = None
-        for i in tqdm(range(1, run_n_segments)):
-            section_text = prompt_texts[i].replace('[start_of_segment]', '').replace('[end_of_segment]', '')
-            guidance_scale = 1.5 if i <= 1 else 1.2
-            prompt_ids = 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 = prompt_ids if i == 1 else torch.cat([raw_output, prompt_ids], dim=1)
-            if input_ids.shape[-1] > 16384 - max_new_tokens - 1:
-                input_ids = input_ids[:, -(16384 - max_new_tokens - 1):]
-
-            with torch.no_grad():
-                output_seq = model.generate(
-                    input_ids=input_ids,
-                    max_new_tokens=max_new_tokens,
-                    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,
+    structured_lyrics = [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
+    return structured_lyrics
+
+def generate_music(genres, lyrics_content, num_segments_run, max_new_tokens_run): # Modified: Function to encapsulate generation logic
+    stage1_output_set_local = [] # Modified: Local variable to store output paths
+
+    lyrics = split_lyrics(lyrics_content)
+    print(len(lyrics))
+    # 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(num_segments_run+1, len(lyrics)) # Modified: Use passed num_segments_run
+
+    print(list(enumerate(tqdm(prompt_texts[:run_n_segments]))))
+
+    global model # Modified: Declare model as global to use the loaded model in Gradio scope
+
+    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 args.use_audio_prompt:
+                audio_prompt = load_audio_mono(args.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(args.prompt_start_time *50): int(args.prompt_end_time *50)] # 50 is tps of xcodec
+                audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
+                sentence_ids = mmtokenizer.tokenize("[start_of_reference]") +  audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
+                head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
+            else:
+                head_id = mmtokenizer.tokenize(prompt_texts[0])
+            prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+        else:
+            prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+
+        prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device)
+        input_ids = torch.cat([raw_output, prompt_ids], dim=1) if i > 1 else prompt_ids
+        # Use window slicing in case output sequence exceeds the context of model
+        max_context = 16384-max_new_tokens_config-1 # Modified: Use max_new_tokens_config
+        if input_ids.shape[-1] > max_context:
+            print(f'Section {i}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
+            input_ids = input_ids[:, -(max_context):]
+        with torch.no_grad():
+            output_seq = model.generate(
+                input_ids=input_ids,
+                max_new_tokens=max_new_tokens_run, # Modified: Use max_new_tokens_run
+                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,
                 )
-
-            raw_output = output_seq if i == 1 else torch.cat([raw_output, output_seq[:, input_ids.shape[-1]:]], dim=1)
-
-        # Process generated tokens
-        ids = raw_output[0].cpu().numpy()
-        soa_idx = np.where(ids == mmtokenizer.soa)[0]
-        eoa_idx = np.where(ids == mmtokenizer.eoa)[0]
-        vocals, instrumentals = [], []
-
-        for i in range(len(soa_idx)):
-            codec_ids = ids[soa_idx[i]+1:eoa_idx[i]]
-            codec_ids = codec_ids[:2 * (len(codec_ids) // 2)]
-            vocals.append(codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0]))
-            instrumentals.append(codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[1]))
-
-        # Generate audio
-        vocals = np.concatenate(vocals, axis=1)
-        instrumentals = np.concatenate(instrumentals, axis=1)
-        
+            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 args.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"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens_run}_vocal_{random_id}".replace('.', '@')+'.npy') # Modified: Use max_new_tokens_run in filename
+    inst_save_path = os.path.join(stage1_output_dir, f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens_run}_instrumental_{random_id}".replace('.', '@')+'.npy') # Modified: Use max_new_tokens_run in filename
+    np.save(vocal_save_path, vocals)
+    np.save(inst_save_path, instrumentals)
+    stage1_output_set_local.append(vocal_save_path)
+    stage1_output_set_local.append(inst_save_path)
+
+
+    # offload model - Removed offloading for gradio integration to keep model loaded
+    # if not args.disable_offload_model:
+    #     model.cpu()
+    #     del model
+    #     torch.cuda.empty_cache()
+
+    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
+    recons_output_dir = os.path.join(args.output_dir, "recons")
+    recons_mix_dir = os.path.join(recons_output_dir, 'mix')
+    os.makedirs(recons_mix_dir, exist_ok=True)
+    tracks = []
+    for npy in stage1_output_set_local: # Modified: Use stage1_output_set_local
+        codec_result = np.load(npy)
+        decodec_rlt=[]
         with torch.no_grad():
-            vocal_audio = codec_model.decode(torch.tensor(vocals.astype(np.int16)).long().unsqueeze(0).permute(1, 0, 2).to(device))
-            inst_audio = codec_model.decode(torch.tensor(instrumentals.astype(np.int16)).long().unsqueeze(0).permute(1, 0, 2).to(device))
+            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] + ".mp3")
+        tracks.append(save_path)
+        save_audio(decodec_rlt, save_path, 16000)
+    # mix tracks
+    for inst_path in tracks:
+        try:
+            if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) \
+                and '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', 'mixed'))
+                vocal_stem, sr = sf.read(inst_path)
+                instrumental_stem, _ = sf.read(vocal_path)
+                mix_stem = (vocal_stem + instrumental_stem) / 1
+                sf.write(recons_mix, mix_stem, sr)
+        except Exception as e:
+            print(e)
+
+    # vocoder to upsample audios
+    vocal_decoder, inst_decoder = build_codec_model(args.config_path, args.vocal_decoder_path, args.inst_decoder_path)
+    vocoder_output_dir = os.path.join(args.output_dir, 'vocoder')
+    vocoder_stems_dir = os.path.join(vocoder_output_dir, 'stems')
+    vocoder_mix_dir = os.path.join(vocoder_output_dir, 'mix')
+    os.makedirs(vocoder_mix_dir, exist_ok=True)
+    os.makedirs(vocoder_stems_dir, exist_ok=True)
+
+    instrumental_output = None # Initialize outside try block
+    vocal_output = None # Initialize outside try block
+    recons_mix_path = "" # Initialize outside try block
+
+
+    for npy in stage1_output_set_local: # Modified: Use stage1_output_set_local
+        if 'instrumental' in npy:
+            # Process instrumental
+            instrumental_output = process_audio(
+                npy,
+                os.path.join(vocoder_stems_dir, 'instrumental.mp3'),
+                args.rescale,
+                args,
+                inst_decoder,
+                codec_model
+            )
+        else:
+            # Process vocal
+            vocal_output = process_audio(
+                npy,
+                os.path.join(vocoder_stems_dir, 'vocal.mp3'),
+                args.rescale,
+                args,
+                vocal_decoder,
+                codec_model
+            )
+    # mix tracks
+    try:
+        mix_output = instrumental_output + vocal_output
+        recons_mix_path_temp = os.path.join(recons_mix_dir, os.path.basename(recons_mix)) # Use recons_mix from previous step
+        save_audio(mix_output, recons_mix_path_temp, 44100, args.rescale)
+        print(f"Created mix: {recons_mix_path_temp}")
+        recons_mix_path = recons_mix_path_temp # Assign to outer scope variable
+    except RuntimeError as e:
+        print(e)
+        print(f"mix {recons_mix_path} failed! inst: {instrumental_output.shape}, vocal: {vocal_output.shape}")
+
+    # Post process
+    final_output_path = os.path.join(args.output_dir, os.path.basename(recons_mix_path)) # Use recons_mix_path from previous step
+    replace_low_freq_with_energy_matched(
+        a_file=recons_mix_path,     # 16kHz # Use recons_mix_path
+        b_file=recons_mix_path_temp,     # 48kHz # Use recons_mix_path_temp
+        c_file=final_output_path,
+        cutoff_freq=5500.0
+    )
+    print("All process Done")
+    return final_output_path # Modified: Return the final output audio path
 
-        # Mix and save audio
-        final_audio = (vocal_audio.cpu().squeeze() + inst_audio.cpu().squeeze()) / 2
-        output_path = os.path.join(output_dir, "final_output.wav")
-        save_audio(final_audio.unsqueeze(0), output_path, 16000)
 
-        return output_path
+# Gradio UI
+model = AutoModelForCausalLM.from_pretrained( # Load model here for Gradio scope
+    "m-a-p/YuE-s1-7B-anneal-en-cot",
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2", # To enable flashattn, you have to install flash-attn
+    ).to(device).eval() # Modified: Load model globally for Gradio to access
+
+def empty_output_folder(output_dir):
+    # List all files in the output directory
+    files = os.listdir(output_dir)
+
+    # Iterate over the files and remove them
+    for file in files:
+        file_path = os.path.join(output_dir, file)
+        try:
+            if os.path.isdir(file_path):
+                # If it's a directory, remove it recursively
+                shutil.rmtree(file_path)
+            else:
+                # If it's a file, delete it
+                os.remove(file_path)
+        except Exception as e:
+            print(f"Error deleting file {file_path}: {e}")
+
+@spaces.GPU(duration=120)
+def infer_gradio(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=200): # Modified: Renamed infer to infer_gradio to avoid conflict
+
+    # Ensure the output folder exists
+    output_dir = "./output"
+    os.makedirs(output_dir, exist_ok=True)
+    print(f"Output folder ensured at: {output_dir}")
+
+    empty_output_folder(output_dir)
+
+    # Call the generation function directly
+    output_audio_path = generate_music(genre_txt_content, lyrics_txt_content, int(num_segments), int(max_new_tokens)) # Modified: Call generate_music and pass num_segments and max_new_tokens as int
+
+    if output_audio_path and os.path.exists(output_audio_path):
+        print("Generated audio file:", output_audio_path)
+        return output_audio_path
+    else:
+        print("No audio file generated or path is invalid.")
+        return None
 
-    except Exception as e:
-        print(f"Error during inference: {str(e)}")
-        raise gr.Error(f"Generation failed: {str(e)}")
 
-# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# YuE Music Generator")
-    with gr.Row():
-        with gr.Column():
-            genre_txt = gr.Textbox(label="Genre Tags", placeholder="e.g., female vocal, jazz, piano")
-            lyrics_txt = gr.Textbox(label="Lyrics", lines=10, placeholder="Enter lyrics with sections like [verse], [chorus]")
-            num_segments = gr.Slider(1, 10, value=2, label="Number of Segments")
-            max_tokens = gr.Slider(500, 3000, value=2000, label="Max Tokens")
-            btn = gr.Button("Generate Music")
-        with gr.Column():
-            audio_out = gr.Audio(label="Generated Music")
-
-    examples = gr.Examples(
-        examples=[
-            ["female blues airy vocal bright vocal piano sad romantic guitar jazz",
-             """[verse]
+    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():
+                if is_shared_ui:
+                    num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
+                    max_new_tokens = gr.Slider(label="Max New Tokens", minimum=500, maximum="3000", step=500, value=500, interactive=True) # increase it after testing
+                else:
+                    num_segments = gr.Number(label="Number of Song Segments", value=2, interactive=True)
+                    max_new_tokens = gr.Slider(label="Max New Tokens", minimum=500, maximum="24000", step=500, value=3000, interactive=True)
+                submit_btn = gr.Button("Submit")
+                music_out = gr.Audio(label="Audio Result")
+
+        gr.Examples(
+            examples = [
+                [
+                    "female blues airy vocal bright vocal piano sad romantic guitar jazz",
+                    """[verse]
 In the quiet of the evening, shadows start to fall
 Whispers of the night wind echo through the hall
 Lost within the silence, I hear your gentle voice
@@ -237,23 +462,39 @@ Guiding me back homeward, making my heart rejoice
 Don't let this moment fade, hold me close tonight
 With you here beside me, everything's alright
 Can't imagine life alone, don't want to let you go
-Stay with me forever, let our love just flow"""],
-            ["rap piano street tough piercing vocal hip-hop synthesizer clear vocal male",
-             """[verse]
+Stay with me forever, let our love just flow
+                    """
+                ],
+                [
+                    "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"""]
-        ],
-        inputs=[genre_txt, lyrics_txt],
-        outputs=audio_out
-    )
+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
 
-    btn.click(
-        fn=run_inference,
-        inputs=[genre_txt, lyrics_txt, num_segments, max_tokens],
-        outputs=audio_out
-    )
+[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],
+            cache_examples = False,
+            # cache_mode="lazy",
+            fn=infer_gradio # Modified: Use infer_gradio
+        )
 
-if __name__ == "__main__":
-    demo.launch()
+    submit_btn.click(
+        fn = infer_gradio, # Modified: Use infer_gradio
+        inputs = [genre_txt, lyrics_txt, num_segments, max_new_tokens],
+        outputs = [music_out]
+    )
+demo.queue().launch(show_api=False, show_error=True)