Update app.py

app.py

@@ -1,20 +1,26 @@
 import gradio as gr
 import subprocess
 import os
-import shutil
-import tempfile
 import spaces
-import torch
 import sys
+import shutil
+import tempfile
 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 torch
+import torchaudio
+from torchaudio.transforms import Resample
+import soundfile as sf
 
-# Install flash-attn and set environment variable to skip cuda build
+# --- Install flash-attn (if needed) ---
 print("Installing flash-attn...")
 subprocess.run(
     "pip install flash-attn --no-build-isolation",
@@ -22,9 +28,9 @@ subprocess.run(
     shell=True
 )
 
-# Download snapshot from huggingface_hub
+# --- Download and set up stage1 files ---
 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}")
@@ -45,156 +51,287 @@ except FileNotFoundError:
     print(f"Directory not found: {inference_dir}")
     exit(1)
 
-# Append necessary module paths
+# --- Append required module paths ---
 base_path = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(os.path.join(base_path, 'xcodec_mini_infer'))
-sys.path.append(os.path.join(base_path, 'xcodec_mini_infer', 'descriptaudiocodec'))
+sys.path.append(os.path.join(base_path, "xcodec_mini_infer"))
+sys.path.append(os.path.join(base_path, "xcodec_mini_infer", "descriptaudiocodec"))
 
-# Other imports
+# --- Additional imports (vocoder & post processing) ---
 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
+from transformers import AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
 from models.soundstream_hubert_new import SoundStream
 
-# Device setup
-device = "cuda:0"
-
-# Load and (optionally) compile the LM model
+# Import vocoder functions (ensure these modules exist)
+from vocoder import build_codec_model, process_audio
+from post_process_audio import replace_low_freq_with_energy_matched
+
+# ----------------------- Global Configuration -----------------------
+# Stage1 and Stage2 model identifiers (change if needed)
+STAGE1_MODEL = "m-a-p/YuE-s1-7B-anneal-en-cot"
+STAGE2_MODEL = "m-a-p/YuE-s2-1B-general"
+# Vocoder model files (paths in the xcodec snapshot)
+BASIC_MODEL_CONFIG = os.path.join(folder_path, "final_ckpt/config.yaml")
+RESUME_PATH = os.path.join(folder_path, "final_ckpt/ckpt_00360000.pth")
+VOCAL_DECODER_PATH = os.path.join(folder_path, "decoders/decoder_131000.pth")
+INST_DECODER_PATH = os.path.join(folder_path, "decoders/decoder_151000.pth")
+VOCODER_CONFIG_PATH = os.path.join(folder_path, "decoders/config.yaml")
+
+# Misc settings
+MAX_NEW_TOKENS = 15         # Duration slider (in seconds, scaled internally)
+RUN_N_SEGMENTS = 2          # Number of segments to generate
+STAGE2_BATCH_SIZE = 4       # Batch size for stage2 inference
+
+# You may change these defaults via Gradio input (see below)
+
+# ----------------------- Device Setup -----------------------
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# ----------------------- Load Stage1 Models and Tokenizer -----------------------
+print("Loading Stage 1 model and tokenizer...")
 model = AutoModelForCausalLM.from_pretrained(
-    "m-a-p/YuE-s1-7B-anneal-en-cot",
+    STAGE1_MODEL,
     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)
-
-# 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')
+    print("torch.compile skipped for Stage1 model:", e)
 
-# 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 model config and initialize codec model
-model_config = OmegaConf.load(basic_model_config)
-# Dynamically create the model from its name in the config.
+# Load codec (xcodec) model for Stage1 & Stage2 decoding
+model_config = OmegaConf.load(BASIC_MODEL_CONFIG)
 codec_class = eval(model_config.generator.name)
 codec_model = codec_class(**model_config.generator.config).to(device)
-parameter_dict = torch.load(resume_path, map_location='cpu')
-codec_model.load_state_dict(parameter_dict['codec_model'])
+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)
+    print("torch.compile skipped for codec_model:", e)
 
-# Pre-compile the regex pattern for splitting lyrics
+# Precompile regex for splitting lyrics
 LYRICS_PATTERN = re.compile(r"\[(\w+)\](.*?)\n(?=\[|\Z)", re.DOTALL)
 
-# ------------------ GPU decorated generation function ------------------ #
+# ----------------------- 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.
+    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.no_grad():
+            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,
+            )
+        # 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)
+        # Only process multiples of 6 seconds; here 50 tokens per second.
+        output_duration = (prompt.shape[-1] // 50) // 6 * 6
+        num_batch = output_duration // 6
+        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 = stage2_generate(model_stage2, prompt[:, start_idx:end_idx], batch_size=current_batch)
+                segments.append(segment)
+            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 array 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=120)
 def generate_music(
     max_new_tokens=5,
     run_n_segments=2,
-    genre_txt=None,
-    lyrics_txt=None,
+    genre_txt="",
+    lyrics_txt="",
     use_audio_prompt=False,
     audio_prompt_path="",
     prompt_start_time=0.0,
     prompt_end_time=30.0,
-    cuda_idx=0,
     rescale=False,
 ):
-    if use_audio_prompt and not audio_prompt_path:
-        raise FileNotFoundError("Please provide an audio prompt filepath when 'use_audio_prompt' is enabled!")
-    max_new_tokens = max_new_tokens * 50  # scaling factor
+    # Scale max_new_tokens (e.g. seconds * 100 tokens per second)
+    max_new_tokens_scaled = max_new_tokens * 100
 
-    with tempfile.TemporaryDirectory() as output_dir:
-        stage1_output_dir = os.path.join(output_dir, "stage1")
+    # Use a temporary directory to store intermediate stage outputs.
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        stage1_output_dir = os.path.join(tmp_dir, "stage1")
+        stage2_output_dir = os.path.join(tmp_dir, "stage2")
         os.makedirs(stage1_output_dir, exist_ok=True)
+        os.makedirs(stage2_output_dir, exist_ok=True)
 
-        # -- 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 ""
+        # ---------------- Stage 1: Text-to-Music Generation ----------------
+        genres = genre_txt.strip()
         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 parameters
+        # Decoding config
         top_p = 0.93
         temperature = 1.0
         repetition_penalty = 1.2
 
-        # 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 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 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]', '')
+        run_n = min(run_n_segments + 1, len(prompt_texts))
+        for i, p in enumerate(tqdm(prompt_texts[:run_n], desc="Stage1 Generation")):
+            section_text = p.replace("[start_of_segment]", "").replace("[end_of_segment]", "")
             guidance_scale = 1.5 if i <= 1 else 1.2
             if i == 0:
-                # 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]")
@@ -207,20 +344,17 @@ 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
 
-            # Enforce maximum context window by slicing if needed
-            max_context = 16384 - max_new_tokens - 1
+            max_context = 16384 - max_new_tokens_scaled - 1
             if input_ids.shape[-1] > max_context:
                 input_ids = input_ids[:, -max_context:]
-
             with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.float16):
                 output_seq = model.generate(
                     input_ids=input_ids,
-                    max_new_tokens=max_new_tokens,
+                    max_new_tokens=max_new_tokens_scaled,
                     min_new_tokens=100,
                     do_sample=True,
                     top_p=top_p,
@@ -233,147 +367,178 @@ 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 raw output codec tokens to temporary files and check token pairs.
+        # Save Stage1 outputs (vocal & instrumental) as npy files.
         ids = raw_output[0].cpu().numpy()
         soa_idx = np.where(ids == soa_token)[0]
         eoa_idx = np.where(ids == eoa_token)[0]
         if len(soa_idx) != len(eoa_idx):
-            raise ValueError(f'Invalid pairs of soa and eoa, Num of soa: {len(soa_idx)}, Num of eoa: {len(eoa_idx)}')
-
+            raise ValueError(f"invalid pairs of soa and eoa: {len(soa_idx)} vs {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]]
+        range_begin = 1 if use_audio_prompt else 0
+        for i in range(range_begin, len(soa_idx)):
+            codec_ids = ids[soa_idx[i]+1:eoa_idx[i]]
             if codec_ids[0] == 32016:
                 codec_ids = codec_ids[1:]
-            # Force even length and reshape into 2 channels.
             codec_ids = codec_ids[:2 * (len(codec_ids) // 2)]
-            codec_ids = np.array(codec_ids)
             reshaped = rearrange(codec_ids, "(n b) -> b n", b=2)
             vocals_list.append(codectool.ids2npy(reshaped[0]))
             instrumentals_list.append(codectool.ids2npy(reshaped[1]))
         vocals = np.concatenate(vocals_list, axis=1)
         instrumentals = np.concatenate(instrumentals_list, axis=1)
-
-        # Save the numpy arrays to temporary files
         vocal_save_path = os.path.join(stage1_output_dir, f"vocal_{str(random_id).replace('.', '@')}.npy")
         inst_save_path = os.path.join(stage1_output_dir, f"instrumental_{str(random_id).replace('.', '@')}.npy")
         np.save(vocal_save_path, vocals)
         np.save(inst_save_path, instrumentals)
         stage1_output_set = [vocal_save_path, inst_save_path]
 
-        print("Converting to Audio...")
-
-        # 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 tracks by decoding codec tokens
-        recons_output_dir = os.path.join(output_dir, "recons")
+        # (Optional) Offload Stage1 model from GPU to free memory.
+        model.cpu()
+        torch.cuda.empty_cache()
+
+        # ---------------- Stage 2: Refinement/Upsampling ----------------
+        print("Stage 2 inference...")
+        model_stage2 = AutoModelForCausalLM.from_pretrained(
+            STAGE2_MODEL,
+            torch_dtype=torch.float16,
+            attn_implementation="flash_attention_2",
+        ).to(device)
+        model_stage2.eval()
+        stage2_result = stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_size=STAGE2_BATCH_SIZE)
+        print("Stage 2 inference completed.")
+
+        # ---------------- Reconstruct Audio from Stage2 Tokens ----------------
+        recons_output_dir = os.path.join(tmp_dir, "recons")
         recons_mix_dir = os.path.join(recons_output_dir, "mix")
         os.makedirs(recons_mix_dir, exist_ok=True)
         tracks = []
-        for npy_path in stage1_output_set:
-            codec_result = np.load(npy_path)
+        for npy in stage2_result:
+            codec_result = np.load(npy)
             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_path))[0] + ".mp3")
+            save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy))[0] + ".mp3")
             tracks.append(save_path)
-            save_audio(decoded_waveform, save_path, sample_rate=16000)
-
-        # Mix vocal and instrumental tracks (using torch to avoid extra I/O if possible)
+            save_audio(decoded_waveform, save_path, 16000, rescale)
+        # Mix vocal and instrumental tracks:
+        mix_audio = None
+        vocal_audio = None
+        instrumental_audio = None
         for inst_path in tracks:
             try:
-                if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) and 'instrumental' in inst_path:
-                    vocal_path = inst_path.replace('instrumental', 'vocal')
+                if (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
-                    # 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))
+                    vocal_data, sr = sf.read(vocal_path)
+                    instrumental_data, _ = sf.read(inst_path)
+                    mix_data = (vocal_data + instrumental_data) / 1.0
+                    recons_mix = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace("instrumental", "mixed"))
+                    sf.write(recons_mix, mix_data, sr)
+                    mix_audio = (sr, (mix_data * 32767).astype(np.int16))
+                    vocal_audio = (sr, (vocal_data * 32767).astype(np.int16))
+                    instrumental_audio = (sr, (instrumental_data * 32767).astype(np.int16))
             except Exception as e:
                 print("Mixing error:", e)
                 return None, None, None
 
-# ------------------ Inference function and Gradio UI ------------------ #
-def infer(genre_txt_content, lyrics_txt_content, num_segments=4, 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
+        # ---------------- Vocoder Upsampling and Post Processing ----------------
+        print("Vocoder upsampling...")
+        vocal_decoder, inst_decoder = build_codec_model(VOCODER_CONFIG_PATH, VOCAL_DECODER_PATH, INST_DECODER_PATH)
+        vocoder_output_dir = os.path.join(tmp_dir, "vocoder")
+        vocoder_stems_dir = os.path.join(vocoder_output_dir, "stems")
+        vocoder_mix_dir = os.path.join(vocoder_output_dir, "mix")
+        os.makedirs(vocoder_stems_dir, exist_ok=True)
+        os.makedirs(vocoder_mix_dir, exist_ok=True)
+        # Process each track with the vocoder (here we process vocal and instrumental separately)
+        if vocal_audio is not None and instrumental_audio is not None:
+            vocal_output = process_audio(
+                stage2_result[0],
+                os.path.join(vocoder_stems_dir, "vocal.mp3"),
+                rescale,
+                None,
+                vocal_decoder,
+                codec_model,
+            )
+            instrumental_output = process_audio(
+                stage2_result[1],
+                os.path.join(vocoder_stems_dir, "instrumental.mp3"),
+                rescale,
+                None,
+                inst_decoder,
+                codec_model,
+            )
+            try:
+                mix_output = instrumental_output + vocal_output
+                vocoder_mix = os.path.join(vocoder_mix_dir, os.path.basename(recons_mix))
+                save_audio(mix_output, vocoder_mix, 44100, rescale)
+                print(f"Created vocoder mix: {vocoder_mix}")
+            except RuntimeError as e:
+                print(e)
+                print("Mixing vocoder outputs failed!")
+        else:
+            print("Missing vocal/instrumental outputs for vocoder stage.")
+
+        # Post-process: Replace low frequency of Stage1 reconstruction with energy-matched vocoder mix.
+        final_mix_path = os.path.join(tmp_dir, "final_mix.mp3")
+        try:
+            replace_low_freq_with_energy_matched(
+                a_file=recons_mix,      # Stage1 mix at 16kHz
+                b_file=vocoder_mix,      # Vocoder mix at 48kHz
+                c_file=final_mix_path,
+                cutoff_freq=5500.0
+            )
+        except Exception as e:
+            print("Post processing error:", e)
+            final_mix_path = recons_mix  # Fall back to Stage1 mix
+
+        # Return final outputs as tuples: (sample_rate, np.int16 audio)
+        final_audio, vocal_audio, instrumental_audio = None, None, None
+        try:
+            final_audio_data, sr = sf.read(final_mix_path)
+            final_audio = (sr, (final_audio_data * 32767).astype(np.int16))
+        except Exception as e:
+            print("Final mix read error:", e)
+        return final_audio, vocal_audio, instrumental_audio
+
+# ----------------------- Gradio Interface -----------------------
 with gr.Blocks() as demo:
     with gr.Column():
-        gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation")
+        gr.Markdown("# YuE: Full-Song Generation (Stage1 + Stage2)")
         gr.HTML(
             """
-            <div style="display:flex;column-gap:4px;">
-                <a href="https://github.com/multimodal-art-projection/YuE">
-                    <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
-                </a>
-                <a href="https://map-yue.github.io">
-                    <img src='https://img.shields.io/badge/Project-Page-green'>
-                </a>
-                <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 style="display:flex; column-gap:4px;">
+              <a href="https://github.com/multimodal-art-projection/YuE"><img src='https://img.shields.io/badge/GitHub-Repo-blue'></a>
+              <a href="https://map-yue.github.io"><img src='https://img.shields.io/badge/Project-Page-green'></a>
             </div>
             """
         )
         with gr.Row():
             with gr.Column():
-                genre_txt = gr.Textbox(label="Genre")
-                lyrics_txt = gr.Textbox(label="Lyrics")
+                genre_txt = gr.Textbox(label="Genre", placeholder="e.g. Bass Metalcore Thrash Metal Furious bright vocal male")
+                lyrics_txt = gr.Textbox(label="Lyrics", placeholder="Paste lyrics with segments such as [verse], [chorus], etc.")
             with gr.Column():
                 num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
-                max_new_tokens = gr.Slider(label="Duration of song", minimum=1, maximum=30, step=1, value=15, interactive=True)
+                max_new_tokens = gr.Slider(label="Duration of song (sec)", minimum=1, maximum=30, step=1, value=15, interactive=True)
+                use_audio_prompt = gr.Checkbox(label="Use Audio Prompt", value=False)
+                audio_prompt_path = gr.Textbox(label="Audio Prompt Filepath (if used)", placeholder="Path to audio file")
                 submit_btn = gr.Button("Submit")
                 music_out = gr.Audio(label="Mixed Audio Result")
-                with gr.Accordion(label="Vocal and Instrumental Result", open=False):
+                with gr.Accordion(label="Vocal and Instrumental Results", open=False):
                     vocal_out = gr.Audio(label="Vocal Audio")
                     instrumental_out = gr.Audio(label="Instrumental Audio")
-
         gr.Examples(
             examples=[
                 [
@@ -421,14 +586,13 @@ Living out my dreams with this mic and a deal
             outputs=[music_out, vocal_out, instrumental_out],
             cache_examples=True,
             cache_mode="eager",
-            fn=infer
+            fn=generate_music
         )
-
     submit_btn.click(
-        fn=infer,
-        inputs=[genre_txt, lyrics_txt, num_segments, max_new_tokens],
+        fn=generate_music,
+        inputs=[max_new_tokens, num_segments, genre_txt, lyrics_txt, use_audio_prompt, audio_prompt_path],
         outputs=[music_out, vocal_out, instrumental_out]
     )
-    gr.Markdown("## Call for Contributions\nIf you find this space interesting please feel free to contribute.")
-    
+    gr.Markdown("## Contributions Welcome\nFeel free to contribute improvements or fixes.")
+
 demo.queue().launch(show_error=True)