app.py

import gradio as gr
import subprocess
import os
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

# Install flash-attn and set environment variable to skip cuda build
print("Installing flash-attn...")
subprocess.run(
    "pip install flash-attn --no-build-isolation",
    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
    shell=True
)

# Download snapshot from huggingface_hub
from huggingface_hub import snapshot_download
folder_path = './xcodec_mini_infer'
if not os.path.exists(folder_path):
    os.mkdir(folder_path)
    print(f"Folder created at: {folder_path}")
else:
    print(f"Folder already exists at: {folder_path}")

snapshot_download(
    repo_id="m-a-p/xcodec_mini_infer",
    local_dir=folder_path
)

# Change working directory to current folder
inference_dir = "."
try:
    os.chdir(inference_dir)
    print(f"Changed working directory to: {os.getcwd()}")
except FileNotFoundError:
    print(f"Directory not found: {inference_dir}")
    exit(1)

# 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'))

# Other imports
from omegaconf import OmegaConf
import torchaudio
from torchaudio.transforms import Resample
import soundfile as sf
from tqdm import tqdm
from einops import rearrange
from codecmanipulator import CodecManipulator
from mmtokenizer import _MMSentencePieceTokenizer
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
import glob
from models.soundstream_hubert_new import SoundStream

# Device setup
device = "cuda:0"

# Load and (optionally) compile the LM model
model = AutoModelForCausalLM.from_pretrained(
    "m-a-p/YuE-s1-7B-anneal-en-cot",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
).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')

# Initialize tokenizer and codec manipulator
mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
codectool = CodecManipulator("xcodec", 0, 1)

# Load codec model config and initialize codec model
model_config = OmegaConf.load(basic_model_config)
# Dynamically create the model from its name in the config.
codec_class = eval(model_config.generator.name)
codec_model = codec_class(**model_config.generator.config).to(device)
parameter_dict = torch.load(resume_path, map_location='cpu')
codec_model.load_state_dict(parameter_dict['codec_model'])
codec_model.eval()
try:
    codec_model = torch.compile(codec_model)
except Exception as e:
    print("torch.compile not used for codec_model:", e)

# Pre-compile the regex pattern for splitting lyrics
LYRICS_PATTERN = re.compile(r"\[(\w+)\](.*?)\n(?=\[|\Z)", re.DOTALL)

# ------------------ GPU decorated generation function ------------------ #
@spaces.GPU(duration=120)
def generate_music(
    max_new_tokens=5,
    run_n_segments=2,
    genre_txt=None,
    lyrics_txt=None,
    use_audio_prompt=False,
    audio_prompt_path="",
    prompt_start_time=0.0,
    prompt_end_time=30.0,
    cuda_idx=0,
    rescale=False,
):
    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 * 100  # scaling factor

    with tempfile.TemporaryDirectory() as output_dir:
        stage1_output_dir = os.path.join(output_dir, "stage1")
        os.makedirs(stage1_output_dir, exist_ok=True)

        # -- In-place logits processor that blocks token ranges --
        class BlockTokenRangeProcessor(LogitsProcessor):
            def __init__(self, start_id, end_id):
                # Pre-create a tensor for indices if possible
                self.blocked_token_ids = list(range(start_id, end_id))
            def __call__(self, input_ids, scores):
                scores[:, self.blocked_token_ids] = -float("inf")
                return scores

        # -- Audio processing utility --
        def load_audio_mono(filepath, sampling_rate=16000):
            audio, sr = torchaudio.load(filepath)
            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 parameters
        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 the global prompt (first element)
        global_prompt_ids = mmtokenizer.tokenize(prompt_texts[0])
        run_n_segments = min(run_n_segments + 1, len(prompt_texts))
        
        # Loop over segments. (Note: Each segment is processed sequentially.)
        for i, p in enumerate(tqdm(prompt_texts[:run_n_segments], desc="Generating segments")):
            # Remove any spurious tokens in the text
            section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
            guidance_scale = 1.5 if i <= 1 else 1.2
            if i == 0:
                # Skip generation on the instruction segment.
                continue

            # Build prompt IDs differently depending on whether audio prompt is enabled.
            if i == 1:
                if use_audio_prompt:
                    audio_prompt = load_audio_mono(audio_prompt_path)
                    audio_prompt = 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]")
                    head_id = global_prompt_ids + sentence_ids
                else:
                    head_id = global_prompt_ids
                prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [soa_token] + codectool.sep_ids
            else:
                prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [soa_token] + codectool.sep_ids

            prompt_ids_tensor = torch.as_tensor(prompt_ids, device=device).unsqueeze(0)
            if raw_output is not None:
                # Concatenate previous outputs with the new prompt
                input_ids = torch.cat([raw_output, prompt_ids_tensor], dim=1)
            else:
                input_ids = prompt_ids_tensor

            # 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,
                    min_new_tokens=100,
                    do_sample=True,
                    top_p=top_p,
                    temperature=temperature,
                    repetition_penalty=repetition_penalty,
                    eos_token_id=eoa_token,
                    pad_token_id=eoa_token,
                    logits_processor=LogitsProcessorList([
                        BlockTokenRangeProcessor(0, 32002),
                        BlockTokenRangeProcessor(32016, 32016)
                    ]),
                    guidance_scale=guidance_scale,
                    use_cache=True
                )
                # 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.
        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)}')

        vocals_list = []
        instrumentals_list = []
        # If using an audio prompt, skip the first pair (it may be reference)
        start_idx = 1 if use_audio_prompt else 0
        for i in range(start_idx, len(soa_idx)):
            codec_ids = ids[soa_idx[i] + 1: eoa_idx[i]]
            if codec_ids[0] == 32016:
                codec_ids = codec_ids[1:]
            # 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")
        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)
            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")
            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)
        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 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))
            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=2, max_new_tokens=15):
    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: Open Music Foundation Models for Full-Song Generation")
        gr.HTML(
            """
            <div style="display:flex;column-gap:4px;">
                <a href="https://github.com/multimodal-art-projection/YuE">
                    <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
                </a>
                <a href="https://map-yue.github.io">
                    <img src='https://img.shields.io/badge/Project-Page-green'>
                </a>
                <a href="https://huggingface.co/spaces/innova-ai/YuE-music-generator-demo?duplicate=true">
                    <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
                </a>
            </div>
            """
        )
        with gr.Row():
            with gr.Column():
                genre_txt = gr.Textbox(label="Genre")
                lyrics_txt = gr.Textbox(label="Lyrics")
            with gr.Column():
                num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
                max_new_tokens = gr.Slider(label="Duration of song", minimum=1, maximum=30, step=1, value=15, interactive=True)
                submit_btn = gr.Button("Submit")
                music_out = gr.Audio(label="Mixed Audio Result")
                with gr.Accordion(label="Vocal and Instrumental Result", open=False):
                    vocal_out = gr.Audio(label="Vocal Audio")
                    instrumental_out = gr.Audio(label="Instrumental Audio")

        gr.Examples(
            examples=[
                [
                    "Bass Metalcore Thrash Metal Furious bright vocal male Angry aggressive vocal Guitar",
                    """[verse]
Step back cause I'll ignite
Won't quit without a fight
No escape, gear up, it's a fierce fight
Brace up, raise your hands up and light
Fear the might. Step back cause I'll ignite
Won't back down without a fight
It keeps going and going, the heat is on.

[chorus]
Hot flame. Hot flame.
Still here, still holding aim
I don't care if I'm bright or dim: nah.
I've made it clear, I'll make it again
All I want is my crew and my gain.
I'm feeling wild, got a bit of rebel style.
Locked inside my mind, hot flame.
                    """
                ],
                [
                    "rap piano street tough piercing vocal hip-hop synthesizer clear vocal male",
                    """[verse]
Woke up in the morning, sun is shining bright
Chasing all my dreams, gotta get my mind right
City lights are fading, but my vision's clear
Got my team beside me, no room for fear
Walking through the streets, beats inside my head
Every step I take, closer to the bread
People passing by, they don't understand
Building up my future with my own two hands

[chorus]
This is my life, and I'mma keep it real
Never gonna quit, no, I'm never gonna stop
Through the highs and lows, I'mma keep it real
Living out my dreams with this mic and a deal
                    """
                ]
            ],
            inputs=[genre_txt, lyrics_txt],
            outputs=[music_out, vocal_out, instrumental_out],
            cache_examples=True,
            cache_mode="eager",
            fn=infer
        )

    submit_btn.click(
        fn=infer,
        inputs=[genre_txt, lyrics_txt, num_segments, max_new_tokens],
        outputs=[music_out, vocal_out, instrumental_out]
    )
    gr.Markdown("## Call for Contributions\nIf you find this space interesting please feel free to contribute.")
    
demo.queue().launch(show_error=True)