Update app.py

app.py

@@ -1,26 +1,20 @@
 import gradio as gr
 import subprocess
 import os
-import spaces
-import sys
 import shutil
 import tempfile
+import spaces
+import torch
+import sys
 import uuid
 import re
+import numpy as np
+import json
 import time
 import copy
 from collections import Counter
-from tqdm import tqdm
-from einops import rearrange
-import numpy as np
-import json
-import argparse
-import torch
-import torchaudio
-from torchaudio.transforms import Resample
-import soundfile as sf
 
-# --- Install flash-attn (if needed) ---
+# Install flash-attn and set environment variable to skip cuda build
 print("Installing flash-attn...")
 subprocess.run(
     "pip install flash-attn --no-build-isolation",
@@ -28,9 +22,9 @@ subprocess.run(
     shell=True
 )
 
-# --- Download and set up stage1 files ---
+# Download snapshot from huggingface_hub
 from huggingface_hub import snapshot_download
-folder_path = "./xcodec_mini_infer"
+folder_path = './xcodec_mini_infer'
 if not os.path.exists(folder_path):
     os.mkdir(folder_path)
     print(f"Folder created at: {folder_path}")
@@ -51,319 +45,156 @@ except FileNotFoundError:
     print(f"Directory not found: {inference_dir}")
     exit(1)
 
-# --- Append required module paths ---
+# 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(base_path, 'xcodec_mini_infer'))
+sys.path.append(os.path.join(base_path, 'xcodec_mini_infer', 'descriptaudiocodec'))
 
-# --- Additional imports (vocoder & post processing) ---
+# 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 AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
+import glob
 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
-
-default_args = argparse.Namespace(cuda_idx=0)
-
-# You may change these defaults via Gradio input (see below)
+# Device setup
+device = "cuda:0"
 
-# ----------------------- 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...")
+# Load and (optionally) compile the LM model
 model = AutoModelForCausalLM.from_pretrained(
-    STAGE1_MODEL,
+    "m-a-p/YuE-s1-7B-anneal-en-cot",
     torch_dtype=torch.float16,
     attn_implementation="flash_attention_2",
 ).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)
 
-model_stage2 = AutoModelForCausalLM.from_pretrained(
-    STAGE2_MODEL,
-    torch_dtype=torch.float16,
-    attn_implementation="flash_attention_2",
-).to(device)
-model_stage2.eval()
+# 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")
-
-# 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)
+# Load codec model config and initialize codec model
+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)
-parameter_dict = torch.load(RESUME_PATH, map_location="cpu")
-codec_model.load_state_dict(parameter_dict["codec_model"])
+parameter_dict = torch.load(resume_path, map_location='cpu')
+codec_model.load_state_dict(parameter_dict['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)
 
-# Precompile regex for splitting lyrics
+# Pre-compile the regex pattern 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)
+# ------------------ GPU decorated generation function ------------------ #
+@spaces.GPU(duration=120)
 def generate_music(
-    genre_txt="",
-    lyrics_txt="",
-    max_new_tokens=2,
-    run_n_segments=1,
+    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,
 ):
-    # Scale max_new_tokens (e.g. seconds * 50 tokens per second)
-    max_new_tokens_scaled = max_new_tokens * 50
+    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
 
-    # 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")
+    with tempfile.TemporaryDirectory() as output_dir:
+        stage1_output_dir = os.path.join(output_dir, "stage1")
         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()
+        # -- 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
+            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.
         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
+        # Decoding config parameters
         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
+        # Pre-tokenize static 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 = 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]", "")
+        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 = 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)
+                    # 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])
+                    # 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]")
@@ -376,17 +207,20 @@ def generate_music(
 
             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:
                 input_ids = prompt_ids_tensor
 
-            max_context = 16384 - max_new_tokens_scaled - 1
+            # Enforce maximum context window by slicing if needed
+            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):
                 output_seq = model.generate(
                     input_ids=input_ids,
-                    max_new_tokens=max_new_tokens_scaled,
+                    max_new_tokens=max_new_tokens,
                     min_new_tokens=100,
                     do_sample=True,
                     top_p=top_p,
@@ -399,177 +233,147 @@ def generate_music(
                         BlockTokenRangeProcessor(32016, 32016)
                     ]),
                     guidance_scale=guidance_scale,
-                    use_cache=True,
+                    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)
                     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)
             else:
                 raw_output = output_seq
 
-        # Save Stage1 outputs (vocal & instrumental) as npy files.
+        # Save raw output codec tokens to temporary files and check token pairs.
         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)}")
+            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 = []
-        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 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:]
+            # 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 = [vocal_save_path, inst_save_path]
 
-        # (Optional) Offload Stage1 model from GPU to free memory.
-        model.cpu()
-        torch.cuda.empty_cache()
+        print("Converting to Audio...")
 
-        # ---------------- 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.")
+        # Utility function for saving audio with in-place clipping
+        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)
 
-        # ---------------- Reconstruct Audio from Stage2 Tokens ----------------
-        recons_output_dir = os.path.join(tmp_dir, "recons")
+        # 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")
         os.makedirs(recons_mix_dir, exist_ok=True)
         tracks = []
-        for npy in stage2_result:
-            codec_result = np.load(npy)
+        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)
-            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(decoded_waveform, save_path, 16000, rescale)
-        # Mix vocal and instrumental tracks:
-        mix_audio = None
-        vocal_audio = None
-        instrumental_audio = None
+            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:
-                    vocal_path = inst_path.replace("instrumental", "vocal")
+                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))
+                    # 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("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,
-                default_args,
-                vocal_decoder,
-                codec_model,
-            )
-            instrumental_output = process_audio(
-                stage2_result[1],
-                os.path.join(vocoder_stems_dir, "instrumental.mp3"),
-                rescale,
-                default_args,
-                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)
-            vocoder_audio_data, sr = sf.read(vocoder_mix)
-            recons_audio, sr = sf.read(recons_mix)
-            final_audio = (sr, (final_audio_data * 32767).astype(np.int16))
-            vocal_audio = (sr, (vocoder_audio_data*32767).astype(np.int16))
-            instrumental_audio = (sr, (recons_audio*32767).astype(np.int16))
-        except Exception as e:
-            print("Final mix read error:", e)
-        return final_audio, vocal_audio, instrumental_audio
-
-# ----------------------- Gradio Interface -----------------------
+# ------------------ Inference function and Gradio UI ------------------ #
+def infer(genre_txt_content, lyrics_txt_content, num_segments=1, max_new_tokens=25):
+    try:
+        mixed_audio_data, vocal_audio_data, instrumental_audio_data = generate_music(
+            genre_txt=genre_txt_content,
+            lyrics_txt=lyrics_txt_content,
+            run_n_segments=num_segments,
+            cuda_idx=0,
+            max_new_tokens=max_new_tokens
+        )
+        return mixed_audio_data, vocal_audio_data, instrumental_audio_data
+    except Exception as e:
+        gr.Warning("An Error Occurred: " + str(e))
+        return None, None, None
+    finally:
+        print("Temporary files deleted.")
+
+# Build Gradio UI
 with gr.Blocks() as demo:
     with gr.Column():
-        gr.Markdown("# YuE: Full-Song Generation (Stage1 + Stage2)")
+        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>
+            <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", 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.")
+                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 (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")
+                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 Results", open=False):
+                with gr.Accordion(label="Vocal and Instrumental Result", open=False):
                     vocal_out = gr.Audio(label="Vocal Audio")
                     instrumental_out = gr.Audio(label="Instrumental Audio")
+
         gr.Examples(
             examples=[
                 [
@@ -617,13 +421,14 @@ Living out my dreams with this mic and a deal
             outputs=[music_out, vocal_out, instrumental_out],
             cache_examples=True,
             cache_mode="eager",
-            fn=generate_music
+            fn=infer
         )
+
     submit_btn.click(
-        fn=generate_music,
-        inputs=[genre_txt, lyrics_txt, max_new_tokens, num_segments, use_audio_prompt, audio_prompt_path],
+        fn=infer,
+        inputs=[genre_txt, lyrics_txt, num_segments, max_new_tokens],
         outputs=[music_out, vocal_out, instrumental_out]
     )
-    gr.Markdown("## Contributions Welcome\nFeel free to contribute improvements or fixes.")
-
-demo.queue().launch(show_error=True)
+    gr.Markdown("## Call for Contributions\nIf you find this space interesting please feel free to contribute.")
+    
+demo.queue().launch(show_error=True)