optimized by deepseek

app.py

@@ -1,12 +1,9 @@
 import gradio as gr
 import subprocess
-import os 
+import os
 import shutil
 import tempfile
 import spaces
-import torch
-import os
-import sys
 
 print("Installing flash-attn...")
 # Install flash attention
@@ -44,394 +41,211 @@ except FileNotFoundError:
 
 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 argparse
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
+import torch
+from huggingface_hub import snapshot_download
+import sys
+import uuid
 import numpy as np
 import json
 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
+import time
 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
 
-is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False
+# Configuration Constants
+MAX_NEW_TOKENS = 3000
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL_NAME = "m-a-p/YuE-s1-7B-anneal-en-cot"
+CODEC_CONFIG_PATH = './xcodec_mini_infer/final_ckpt/config.yaml'
+CODEC_CKPT_PATH = './xcodec_mini_infer/final_ckpt/ckpt_00360000.pth'
 
-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}")
+# Global Initialization
+is_shared_ui = "innova-ai/YuE-music-generator-demo" in os.environ.get('SPACE_ID', '')
 
-# Function to create a temporary file with string content
-def create_temp_file(content, prefix, suffix=".txt"):
-    temp_file = tempfile.NamedTemporaryFile(delete=False, mode="w", prefix=prefix, suffix=suffix)
-    # Ensure content ends with newline and normalize line endings
-    content = content.strip() + "\n\n"  # Add extra newline at end
-    content = content.replace("\r\n", "\n").replace("\r", "\n")
-    temp_file.write(content)
-    temp_file.close()
-    
-    # Debug: Print file contents
-    print(f"\nContent written to {prefix}{suffix}:")
-    print(content)
-    print("---")
+# Preload models and components
+def load_models():
+    print("Initializing models...")
     
-    return temp_file.name
-
-def get_last_mp3_file(output_dir):
-    # List all files in the output directory
-    files = os.listdir(output_dir)
+    # Load main model
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16,
+        attn_implementation="flash_attention_2",
+    ).to(DEVICE).eval()
     
-    # Filter only .mp3 files
-    mp3_files = [file for file in files if file.endswith('.mp3')]
+    # Load tokenizer
+    mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
     
-    if not mp3_files:
-        print("No .mp3 files found in the output folder.")
-        return None
+    # Load codec model
+    model_config = OmegaConf.load(CODEC_CONFIG_PATH)
+    codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(DEVICE)
+    parameter_dict = torch.load(CODEC_CKPT_PATH, map_location='cpu')
+    codec_model.load_state_dict(parameter_dict['codec_model'])
+    codec_model.eval()
     
-    # Get the full path for the mp3 files
-    mp3_files_with_path = [os.path.join(output_dir, file) for file in mp3_files]
+    # Initialize codec tools
+    codectool = CodecManipulator("xcodec", 0, 1)
     
-    # Sort the files based on the modification time (most recent first)
-    mp3_files_with_path.sort(key=lambda x: os.path.getmtime(x), reverse=True)
+    # Precompute token IDs
+    start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
+    end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
     
-    # Return the most recent .mp3 file
-    return mp3_files_with_path[0]
-
-device = "cuda:0"
-
-model = AutoModelForCausalLM.from_pretrained(
-    "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
-    )
-model.to(device)
-model.eval()
+    return model, mmtokenizer, codec_model, codectool, start_of_segment, end_of_segment
+
+# Preload all models and components
+model, mmtokenizer, codec_model, codectool, start_of_segment, end_of_segment = load_models()
+
+# Audio processing cache
+resampler_cache = {}
+def get_resampler(orig_freq, new_freq):
+    key = (orig_freq, new_freq)
+    if key not in resampler_cache:
+        resampler_cache[key] = Resample(orig_freq=orig_freq, new_freq=new_freq).to(DEVICE)
+    return resampler_cache[key]
+
+def load_audio_mono(filepath, sampling_rate=16000):
+    audio, sr = torchaudio.load(filepath)
+    audio = torch.mean(audio, dim=0, keepdim=True).to(DEVICE)
+    if sr != sampling_rate:
+        resampler = get_resampler(sr, sampling_rate)
+        audio = resampler(audio)
+    return audio
 
+@spaces.GPU(duration=120)
 def generate_music(
-    stage1_model="m-a-p/YuE-s1-7B-anneal-en-cot",
-    max_new_tokens=3000,
-    run_n_segments=2,
     genre_txt=None,
     lyrics_txt=None,
+    max_new_tokens=3000,
+    run_n_segments=2,
     use_audio_prompt=False,
     audio_prompt_path="",
     prompt_start_time=0.0,
     prompt_end_time=30.0,
     output_dir="./output",
     keep_intermediate=False,
-    disable_offload_model=False,
-    cuda_idx=0,
-    basic_model_config='./xcodec_mini_infer/final_ckpt/config.yaml',
-    resume_path='./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth',
-    config_path='./xcodec_mini_infer/decoders/config.yaml',
-    vocal_decoder_path='./xcodec_mini_infer/decoders/decoder_131000.pth',
-    inst_decoder_path='./xcodec_mini_infer/decoders/decoder_151000.pth',
     rescale=False,
 ):
-    if use_audio_prompt and not audio_prompt_path:
-        raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
-    
-    model = stage1_model
-    cuda_idx = cuda_idx
-    max_new_tokens = max_new_tokens
-    stage1_output_dir = os.path.join(output_dir, f"stage1")
+    # Create output directories once
+    os.makedirs(output_dir, exist_ok=True)
+    stage1_output_dir = os.path.join(output_dir, "stage1")
     os.makedirs(stage1_output_dir, exist_ok=True)
 
-    mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-    
-    codectool = CodecManipulator("xcodec", 0, 1)
-    model_config = OmegaConf.load(basic_model_config)
-    codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
-    parameter_dict = torch.load(resume_path, map_location='cpu')
-    codec_model.load_state_dict(parameter_dict['codec_model'])
-    codec_model.to(device)
-    codec_model.eval()
-    
-    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 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: str):
-        pattern = r"\[(\w+)\](.*?)\n(?=\[|\Z)"
-        segments = re.findall(pattern, lyrics, re.DOTALL)
-        structured_lyrics = [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
-        return structured_lyrics
-
-    # Call the function and print the result
-    stage1_output_set = []
-
+    # Process inputs
     genres = genre_txt.strip()
     lyrics = split_lyrics(lyrics_txt+"\n")
-    # intruction
     full_lyrics = "\n".join(lyrics)
-    prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"]
-    prompt_texts += lyrics
-
-
+    prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"] + 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
+    # Audio prompt processing
+    audio_prompt_codec_ids = []
+    if use_audio_prompt:
+        if not audio_prompt_path:
+            raise FileNotFoundError("Audio prompt path required when using audio prompt!")
+        
+        audio_prompt = load_audio_mono(audio_prompt_path)
+        with torch.inference_mode():
+            raw_codes = codec_model.encode(audio_prompt.unsqueeze(0), target_bw=0.5)
+            raw_codes = raw_codes.transpose(0, 1).cpu().numpy().astype(np.int16)
+        
+        code_ids = codectool.npy2ids(raw_codes[0])
+        audio_prompt_codec = code_ids[int(prompt_start_time*50):int(prompt_end_time*50)]
+        audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
 
-    # Format text prompt
+    # Generation loop optimization
     run_n_segments = min(run_n_segments+1, len(lyrics))
-
-    print(list(enumerate(tqdm(prompt_texts[:run_n_segments]))))
-
-    for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
-        section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
-        guidance_scale = 1.5 if i <=1 else 1.2
-        if i==0:
-            continue
-        if i==1:
-            if use_audio_prompt:
-                audio_prompt = load_audio_mono(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(prompt_start_time *50): int(prompt_end_time *50)] # 50 is tps of xcodec
-                audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
-                sentence_ids = mmtokenizer.tokenize("[start_of_reference]") +  audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
-                head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
+    output_seq = None
+    
+    with torch.inference_mode():
+        for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
+            if i == 0: continue  # Skip system prompt
+            
+            # Prepare prompt
+            section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
+            guidance_scale = 1.5 if i <= 1 else 1.2
+            
+            if i == 1:
+                prompt_ids = mmtokenizer.tokenize(prompt_texts[0])
+                if use_audio_prompt:
+                    prompt_ids += mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
+                prompt_ids += start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_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 = 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-1
-        if input_ids.shape[-1] > max_context:
-            print(f'Section {i}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
-            input_ids = input_ids[:, -(max_context):]
-        with torch.no_grad():
+            # Process input sequence
+            prompt_ids = torch.tensor(prompt_ids, device=DEVICE).unsqueeze(0)
+            input_ids = torch.cat([output_seq, prompt_ids], dim=1) if i > 1 else prompt_ids
+            
+            # Generate sequence
             output_seq = model.generate(
-                input_ids=input_ids, 
-                max_new_tokens=max_new_tokens, 
-                min_new_tokens=100, 
-                do_sample=True, 
-                top_p=top_p,
-                temperature=temperature, 
-                repetition_penalty=repetition_penalty, 
+                input_ids=input_ids,
+                max_new_tokens=max_new_tokens,
+                min_new_tokens=100,
+                do_sample=True,
+                top_p=0.93,
+                temperature=1.0,
+                repetition_penalty=1.2,
                 eos_token_id=mmtokenizer.eoa,
                 pad_token_id=mmtokenizer.eoa,
-                logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
+                logits_processor=LogitsProcessorList([
+                    BlockTokenRangeProcessor(0, 32002),
+                    BlockTokenRangeProcessor(32016, 32016)
+                ]),
                 guidance_scale=guidance_scale,
-                )
-            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)}')
+    # Post-processing optimization
+    ids = output_seq[0].cpu().numpy()
+    soa_idx = np.where(ids == mmtokenizer.soa)[0]
+    eoa_idx = np.where(ids == mmtokenizer.eoa)[0]
+    
+    # Vectorized audio processing
+    vocals, instrumentals = process_audio_segments(ids, soa_idx, eoa_idx, codectool)
+    
+    # Save and mix audio
+    return save_and_mix_audio(vocals, instrumentals, genres, random_id, output_dir)
 
-    vocals = []
-    instrumentals = []
-    range_begin = 1 if use_audio_prompt else 0
+def process_audio_segments(ids, soa_idx, eoa_idx, codectool):
+    vocals, instrumentals = [], []
+    range_begin = 1 if len(soa_idx) > len(eoa_idx) 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}_vocal_{random_id}".replace('.', '@')+'.npy')
-    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}_instrumental_{random_id}".replace('.', '@')+'.npy')
-    np.save(vocal_save_path, vocals)
-    np.save(inst_save_path, instrumentals)
-    stage1_output_set.append(vocal_save_path)
-    stage1_output_set.append(inst_save_path)
-
-
-    # offload model
-    if not disable_offload_model:
-        model.cpu()
-        del model
-        torch.cuda.empty_cache()
-
-    print("Converting to Audio...")
+        codec_ids = codec_ids[:2 * (len(codec_ids) // 2)]
+        
+        # Vectorized processing
+        arr = rearrange(codec_ids, "(n b) -> b n", b=2)
+        vocals.append(codectool.ids2npy(arr[0]))
+        instrumentals.append(codectool.ids2npy(arr[1]))
+        
+    return np.concatenate(vocals, axis=1), np.concatenate(instrumentals, axis=1)
 
-    # 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(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:
-        codec_result = np.load(npy)
-        decodec_rlt=[]
-        with torch.no_grad():
-            decoded_waveform = codec_model.decode(torch.as_tensor(codec_result.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device))
-        decoded_waveform = decoded_waveform.cpu().squeeze(0)
-        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)
+def save_and_mix_audio(vocals, instrumentals, genres, random_id, output_dir):
+    # Save directly to memory buffers
+    vocal_buf = torch.as_tensor(vocals.astype(np.int16), device=DEVICE)
+    inst_buf = torch.as_tensor(instrumentals.astype(np.int16), device=DEVICE)
     
-
-    # vocoder to upsample audios
-    vocal_decoder, inst_decoder = build_codec_model(config_path, vocal_decoder_path, inst_decoder_path)
-    vocoder_output_dir = os.path.join(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
-    vocal_output = None
-    for npy in stage1_output_set:
-        if 'instrumental' in npy:
-            # Process instrumental
-            instrumental_output = process_audio(
-                npy,
-                os.path.join(vocoder_stems_dir, 'instrumental.mp3'),
-                rescale,
-                argparse.Namespace(**locals()), # Convert local variables to argparse.Namespace
-                inst_decoder,
-                codec_model
-            )
-        else:
-            # Process vocal
-            vocal_output = process_audio(
-                npy,
-                os.path.join(vocoder_stems_dir, 'vocal.mp3'),
-                rescale,
-                 argparse.Namespace(**locals()), # Convert local variables to argparse.Namespace
-                vocal_decoder,
-                codec_model
-            )
-    # mix tracks
-    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 mix: {vocoder_mix}")
-    except RuntimeError as e:
-        print(e)
-        print(f"mix {vocoder_mix} failed! inst: {instrumental_output.shape}, vocal: {vocal_output.shape}")
-
-    # Post process
-    replace_low_freq_with_energy_matched(
-        a_file=recons_mix,     # 16kHz
-        b_file=vocoder_mix,     # 48kHz
-        c_file=os.path.join(output_dir, os.path.basename(recons_mix)),
-        cutoff_freq=5500.0
-    )
-    print("All process Done")
-    return recons_mix
-
-
-@spaces.GPU(duration=120)
-def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=200):
-
-    # 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)
+    with torch.inference_mode():
+        vocal_wav = codec_model.decode(vocal_buf.unsqueeze(0).permute(1, 0, 2))
+        inst_wav = codec_model.decode(inst_buf.unsqueeze(0).permute(1, 0, 2))
     
-    # Execute the command
-    try:
-        music = generate_music(stage1_model=model, genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments, output_dir=output_dir, cuda_idx=0, max_new_tokens=max_new_tokens)
-        
-        return music
-    except subprocess.CalledProcessError as e:
-        print(f"Error occurred: {e}")
-        return None
-    finally:
-        # Clean up temporary files
-        print("Temporary files deleted.")
+    # Mix directly in GPU memory
+    mixed = (vocal_wav + inst_wav) / 2
+    mixed = mixed.squeeze(0).cpu().numpy()
+    
+    # Save final output
+    output_path = os.path.join(output_dir, f"mixed_{genres}_{random_id}.mp3")
+    sf.write(output_path, mixed.T, 16000)
+    
+    return output_path
 
 # Gradio