Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/logo/yue.mp3 filter=lfs diff=lfs merge=lfs -text

ORIGINAL_README.md

@@ -0,0 +1,163 @@
+<p align="center">
+    <img src="./assets/logo/白底.png" width="400" />
+</p>
+
+<p align="center">
+    <a href="https://map-yue.github.io/">Demo 🎶</a> &nbsp;|&nbsp; 📑 <a href="">Paper (coming soon)</a>
+    <br>
+    <a href="https://huggingface.co/m-a-p/YuE-s1-7B-anneal-en-cot">YuE-s1-7B-anneal-en-cot 🤗</a> &nbsp;|&nbsp; <a href="https://huggingface.co/m-a-p/YuE-s1-7B-anneal-en-icl">YuE-s1-7B-anneal-en-icl 🤗</a> &nbsp;|&nbsp; <a href="https://huggingface.co/m-a-p/YuE-s1-7B-anneal-jp-kr-cot">YuE-s1-7B-anneal-jp-kr-cot 🤗</a>
+    <br>
+    <a href="https://huggingface.co/m-a-p/YuE-s1-7B-anneal-jp-kr-icl">YuE-s1-7B-anneal-jp-kr-icl 🤗</a> &nbsp;|&nbsp; <a href="https://huggingface.co/m-a-p/YuE-s1-7B-anneal-zh-cot">YuE-s1-7B-anneal-zh-cot 🤗</a> &nbsp;|&nbsp; <a href="https://huggingface.co/m-a-p/YuE-s1-7B-anneal-zh-icl">YuE-s1-7B-anneal-zh-icl 🤗</a>
+    <br>
+    <a href="https://huggingface.co/m-a-p/YuE-s2-1B-general">YuE-s2-1B-general 🤗</a> &nbsp;|&nbsp; <a href="https://huggingface.co/m-a-p/YuE-upsampler">YuE-upsampler 🤗</a>
+</p>
+
+---
+Our model's name is **YuE (乐)**. In Chinese, the word means "music" and "happiness." Some of you may find words that start with Yu hard to pronounce. If so, you can just call it "yeah." We wrote a song with our model's name.
+
+<audio controls src="https://cdn-uploads.huggingface.co/production/uploads/6555e8d8a0c34cd61a6b9ce3/rG-ELxMyzDU7zH-inB9DV.mpga"></audio>
+
+YuE is a groundbreaking series of open-source foundation models designed for music generation, specifically for transforming lyrics into full songs (lyrics2song). It can generate a complete song, lasting several minutes, that includes both a catchy vocal track and complementary accompaniment, ensuring a polished and cohesive result. YuE is capable of modeling diverse genres/vocal styles. Below are examples of songs in the pop and metal genres. For more styles, please visit the demo page.
+
+Pop:Quiet Evening
+<audio controls src="https://cdn-uploads.huggingface.co/production/uploads/640701cb4dc5f2846c91d4eb/gnBULaFjcUyXYzzIwXLZq.mpga"></audio>
+Metal: Step Back
+<audio controls src="https://cdn-uploads.huggingface.co/production/uploads/6555e8d8a0c34cd61a6b9ce3/kmCwl4GRS70UYDEELL-Tn.mpga"></audio>
+
+## News and Updates
+
+* **2025.01.26 🔥**: We have released the **YuE** series.
+
+<br>
+
+## Requirements
+
+Python >=3.8 is recommended.
+
+Install dependencies with the following command:
+
+```
+pip install -r requirements.txt
+```
+
+### **Important: Install FlashAttention 2**
+For saving GPU memory, **FlashAttention 2 is mandatory**. Without it, large sequence lengths will lead to out-of-memory (OOM) errors, especially on GPUs with limited memory. Install it using the following command:
+```
+pip install flash-attn --no-build-isolation
+```
+Before installing FlashAttention, ensure that your CUDA environment is correctly set up. 
+For example, if you are using CUDA 11.8:
+- If using a module system:
+``` module load cuda11.8/toolkit/11.8.0 ```
+- Or manually configure CUDA in your shell:
+```
+    export PATH=/usr/local/cuda-11.8/bin:$PATH
+    export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
+```
+
+---
+
+## GPU Memory Usage and Sessions
+
+YuE requires significant GPU memory for generating long sequences. Below are the recommended configurations:
+
+- **For GPUs with 24GB memory or less**: Run **up to 2 sessions** concurrently to avoid out-of-memory (OOM) errors.
+- **For full song generation** (many sessions, e.g., 4 or more): Use **GPUs with at least 80GB memory**. This can be achieved by combining multiple GPUs and enabling tensor parallelism.
+
+To customize the number of sessions, the interface allows you to specify the desired session count. By default, the model runs **2 sessions** for optimal memory usage.
+
+---
+
+## Quickstart
+
+```
+# Make sure you have git-lfs installed (https://git-lfs.com)
+git lfs install
+git clone https://github.com/multimodal-art-projection/YuE.git
+
+cd YuE/inference/
+git clone https://huggingface.co/m-a-p/xcodec_mini_infer
+```
+
+Here’s a quick guide to help you generate music with **YuE** using 🤗 Transformers. Before running the code, make sure your environment is properly set up, and that all dependencies are installed.
+
+### Running the Script
+
+In the following example, customize the `genres` and `lyrics` in the script, then execute it to generate a song with **YuE**.
+
+Notice: Set `--run_n_segments` to the number of lyric sections if you want to generate a full song. Additionally, you can increase `--stage2_batch_size` based on your available GPU memory.
+
+```bash
+cd YuE/inference/
+python infer.py \
+    --stage1_model m-a-p/YuE-s1-7B-anneal-en-cot \
+    --stage2_model m-a-p/YuE-s2-1B-general \
+    --genre_txt prompt_examples/genre.txt \
+    --lyrics_txt prompt_examples/lyrics.txt \
+    --run_n_segments 2 \
+    --stage2_batch_size 4 \
+    --output_dir ./output \
+    --cuda_idx 0 \
+    --max_new_tokens 3000 
+```
+
+If you want to use audio prompt, enable `--use_audio_prompt`, and provide audio prompt:
+```bash
+cd YuE/inference/
+python infer.py \
+    --stage1_model m-a-p/YuE-s1-7B-anneal-en-icl \
+    --stage2_model m-a-p/YuE-s2-1B-general \
+    --genre_txt prompt_examples/genre.txt \
+    --lyrics_txt prompt_examples/lyrics.txt \
+    --run_n_segments 2 \
+    --stage2_batch_size 4 \
+    --output_dir ./output \
+    --cuda_idx 0 \
+    --max_new_tokens 3000 \
+    --audio_prompt_path {YOUR_AUDIO_FILE} \
+    --prompt_start_time 0 \
+    --prompt_end_time 30 
+```
+
+
+---
+
+### **Execution Time**
+On an **H800 GPU**, generating 30s audio takes **150 seconds**.
+On an **RTX 4090 GPU**, generating 30s audio takes approximately **360 seconds**.  
+
+**Tips:**
+1. `genres` should include details like instruments, genre, mood, vocal timbre, and vocal gender.
+2. The length of `lyrics` segments and the `--max_new_tokens` value should be matched. For example, if `--max_new_tokens` is set to 3000, the maximum duration for a segment is around 30 seconds. Ensure your lyrics fit this time frame.
+3. If using audio prompt，the duration around 30s will be fine.
+---
+
+### Notice
+1. A suitable [Genre] tag consists of five components: genre, instrument, mood, gender, and timbre. All five should be included if possible, separated by spaces. The values of timbre should include "vocal" (e.g., "bright vocal").
+
+2. Although our tags have an open vocabulary, we have provided the 200 most commonly used [tags](./wav_top_200_tags.json). It is recommended to select tags from this list for more stable results.
+
+3. The order of the tags is flexible. For example, a stable genre control string might look like: "[Genre] inspiring female uplifting pop airy vocal electronic bright vocal vocal."
+
+4. Additionally, we have introduced the "Mandarin" and "Cantonese" tags to distinguish between Mandarin and Cantonese, as their lyrics often share similarities.
+
+## License Agreement
+
+Creative Commons Attribution Non Commercial 4.0
+
+---
+
+## Citation
+
+If you find our paper and code useful in your research, please consider giving a star :star: and citation :pencil: :)
+
+```BibTeX
+@misc{yuan2025yue,
+  title={YuE: Open Music Foundation Models for Full-Song Generation},
+  author={Ruibin Yuan and Hanfeng Lin and Shawn Guo and Ge Zhang and Jiahao Pan and Yongyi Zang and Haohe Liu and Xingjian Du and Xeron Du and Zhen Ye and Tianyu Zheng and Yinghao Ma and Minghao Liu and Lijun Yu and Zeyue Tian and Ziya Zhou and Liumeng Xue and Xingwei Qu and Yizhi Li and Tianhao Shen and Ziyang Ma and Shangda Wu and Jun Zhan and Chunhui Wang and Yatian Wang and Xiaohuan Zhou and Xiaowei Chi and Xinyue Zhang and Zhenzhu Yang and Yiming Liang and Xiangzhou Wang and Shansong Liu and Lingrui Mei and Peng Li and Yong Chen and Chenghua Lin and Xie Chen and Gus Xia and Zhaoxiang Zhang and Chao Zhang and Wenhu Chen and Xinyu Zhou and Xipeng Qiu and Roger Dannenberg and Jiaheng Liu and Jian Yang and Stephen Huang and Wei Xue and Xu Tan and Yike Guo}, 
+  howpublished={\url{https://github.com/multimodal-art-projection/YuE}},
+  year={2025},
+  note={GitHub repository}
+}
+```
+<br>

README.md

@@ -1,14 +1,12 @@
 ---
-title: YuE Music Generator Demo
-emoji: 🌍
-colorFrom: yellow
+title: YuE
+emoji: 👩‍🎤
+colorFrom: pink
 colorTo: green
 sdk: gradio
 sdk_version: 5.13.1
 app_file: app.py
 pinned: false
-license: apache-2.0
-short_description: Suno lvl open source music generator
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,333 +1,252 @@
-import os
-import subprocess
-
-# Install flash attention
-subprocess.run(
-    "pip install flash-attn --no-build-isolation",
-    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-    shell=True,
-)
-
-import spaces
-import os
-import torch
-import numpy as np
-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 gradio as gr
-import re
-from collections import Counter
-from codecmanipulator import CodecManipulator
-from mmtokenizer import _MMSentencePieceTokenizer
-from transformers import AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
-# 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
-
-# Initialize global variables and models
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-codectool = CodecManipulator("xcodec", 0, 1)
-codectool_stage2 = CodecManipulator("xcodec", 0, 8)
+import subprocess
+import os 
+import shutil
+import tempfile
+
+is_shared_ui = True if "fffiloni/YuE" in os.environ['SPACE_ID'] else False
+
+# Install required package
+def install_flash_attn():
+    try:
+        print("Installing flash-attn...")
+        subprocess.run(
+            ["pip", "install", "flash-attn", "--no-build-isolation"], 
+            check=True
+        )
+        print("flash-attn installed successfully!")
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to install flash-attn: {e}")
+        exit(1)
 
-# Load models once at startup
-def load_models():
-    # Stage 1 Model
-    stage1_model = AutoModelForCausalLM.from_pretrained(
-        "m-a-p/YuE-s1-7B-anneal-en-cot",
-        torch_dtype=torch.bfloat16,
-        attn_implementation="flash_attention_2"
-    ).to(device)
-    stage1_model.eval()
+# Install flash-attn
+install_flash_attn()
 
-    # Stage 2 Model
-    stage2_model = AutoModelForCausalLM.from_pretrained(
-        "m-a-p/YuE-s2-1B-general",
-        torch_dtype=torch.float16,
-        attn_implementation="flash_attention_2"
-    ).to(device)
-    stage2_model.eval()
+from huggingface_hub import snapshot_download 
 
-    # Codec Model
-    model_config = OmegaConf.load('./xcodec_mini_infer/final_ckpt/config.yaml')
-    codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
-    parameter_dict = torch.load('./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', map_location='cpu')
-    codec_model.load_state_dict(parameter_dict['codec_model'])
-    codec_model.eval()
+# Create xcodec_mini_infer folder
+folder_path = './inference/xcodec_mini_infer'
 
-    return stage1_model, stage2_model, codec_model
+# Create the folder if it doesn't exist
+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}")
 
-stage1_model, stage2_model, codec_model = load_models()
+snapshot_download(
+    repo_id = "m-a-p/xcodec_mini_infer",
+    local_dir = "./inference/xcodec_mini_infer"
+)
 
-# Helper functions
-def split_lyrics(lyrics):
-    pattern = r"\[(\w+)\](.*?)\n(?=\[|\Z)"
-    segments = re.findall(pattern, lyrics, re.DOTALL)
-    return [f"[{seg[0]}]\n{seg[1].strip()}\n\n" for seg in segments]
+# Change to the "inference" directory
+inference_dir = "./inference"
+try:
+    os.chdir(inference_dir)
+    print(f"Changed working directory to: {os.getcwd()}")
+except FileNotFoundError:
+    print(f"Directory not found: {inference_dir}")
+    exit(1)
+
+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}")
+
+# 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("---")
+    
+    return temp_file.name
 
-def load_audio_mono(filepath, sampling_rate=16000):
-    audio, sr = torchaudio.load(filepath)
-    audio = torch.mean(audio, 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 get_last_mp3_file(output_dir):
+    # List all files in the output directory
+    files = os.listdir(output_dir)
+    
+    # Filter only .mp3 files
+    mp3_files = [file for file in files if file.endswith('.mp3')]
+    
+    if not mp3_files:
+        print("No .mp3 files found in the output folder.")
+        return None
+    
+    # Get the full path for the mp3 files
+    mp3_files_with_path = [os.path.join(output_dir, file) for file in mp3_files]
+    
+    # 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)
+    
+    # Return the most recent .mp3 file
+    return mp3_files_with_path[0]
 
-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)
+def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
+    # Create temporary files
+    genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
+    lyrics_txt_path = create_temp_file(lyrics_txt_content, prefix="lyrics_")
 
-# Stage 1 Generation
-def stage1_generate(genres, lyrics_text, use_audio_prompt, audio_prompt_path, prompt_start_time, prompt_end_time):
-    structured_lyrics = split_lyrics(lyrics_text)
-    full_lyrics = "\n".join(structured_lyrics)
-    prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"] + structured_lyrics
+    print(f"Genre TXT path: {genre_txt_path}")
+    print(f"Lyrics TXT path: {lyrics_txt_path}")
 
-    random_id = str(uuid.uuid4())
-    output_dir = os.path.join("./output", random_id)
+    # 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)
+ 
+    # Command and arguments with optimized settings
+    command = [
+        "python", "infer.py",
+        "--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot",
+        "--stage2_model", "m-a-p/YuE-s2-1B-general",
+        "--genre_txt", f"{genre_txt_path}",
+        "--lyrics_txt", f"{lyrics_txt_path}",
+        "--run_n_segments", f"{num_segments}",
+        "--stage2_batch_size", "4",
+        "--output_dir", f"{output_dir}",
+        "--cuda_idx", "0",
+        "--max_new_tokens", f"{max_new_tokens}",
+        "--disable_offload_model"
+    ]
 
-    stage1_output_set = []
-    for i, p in enumerate(tqdm(prompt_texts)):
-        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 and 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).cpu().numpy().astype(np.int16)
-            audio_prompt_codec = codectool.npy2ids(raw_codes[0])[int(prompt_start_time * 50): int(prompt_end_time * 50)]
-            audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
-            sentence_ids = mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
-            head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
-        else:
-            head_id = mmtokenizer.tokenize(prompt_texts[0])
-
-        prompt_ids = head_id + mmtokenizer.tokenize("[start_of_segment]") + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
-        prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device)
-
-        with torch.no_grad():
-            output_seq = stage1_model.generate(
-                input_ids=prompt_ids,
-                max_new_tokens=3000,
-                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,
-            )
-
-        if i > 1:
-            raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, prompt_ids.shape[-1]:]], dim=1)
+    # Set up environment variables for CUDA with optimized settings
+    env = os.environ.copy()
+    env.update({
+        "CUDA_VISIBLE_DEVICES": "0",
+        "CUDA_HOME": "/usr/local/cuda",
+        "PATH": f"/usr/local/cuda/bin:{env.get('PATH', '')}",
+        "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}"
+    })
+    
+    # Execute the command
+    try:
+        subprocess.run(command, check=True, env=env)
+        print("Command executed successfully!")
+        
+        # Check and print the contents of the output folder
+        output_files = os.listdir(output_dir)
+        if output_files:
+            print("Output folder contents:")
+            for file in output_files:
+                print(f"- {file}")
+
+            last_mp3 = get_last_mp3_file(output_dir)
+
+            if last_mp3:
+                print("Last .mp3 file:", last_mp3)
+                return last_mp3
+            else:
+                return None
         else:
-            raw_output = output_seq
-
-    # Save Stage 1 outputs
-    ids = raw_output[0].cpu().numpy()
-    soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
-    eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
-
-    vocals = []
-    instrumentals = []
-    for i in range(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(output_dir, f"vocal_{random_id}.npy")
-    inst_save_path = os.path.join(output_dir, f"instrumental_{random_id}.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)
-
-    return stage1_output_set, output_dir
-
-# Stage 2 Generation
-def stage2_generate(model, prompt, batch_size=16):
-    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)
-
-    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 = np.concatenate(codec_list, axis=0)
-        prompt_ids = np.concatenate(
-            [
-                np.tile([mmtokenizer.soa, mmtokenizer.stage_1], (batch_size, 1)),
-                codec_ids,
-                np.tile([mmtokenizer.stage_2], (batch_size, 1)),
-            ],
-            axis=1
+            print("Output folder is empty.")
+            return None
+    except subprocess.CalledProcessError as e:
+        print(f"Error occurred: {e}")
+        return None
+    finally:
+        # Clean up temporary files
+        os.remove(genre_txt_path)
+        os.remove(lyrics_txt_path)
+        print("Temporary files deleted.")
+
+# Gradio 
+
+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/fffiloni/YuE?duplicate=true">
+                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
+            </a>
+        </div>
+        """)
+        with gr.Row():
+            with gr.Column():
+                genre_txt = gr.Textbox(label="Genre")
+                lyrics_txt = gr.Textbox(label="Lyrics")
+                
+            with gr.Column():
+                if is_shared_ui:
+                    num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
+                    max_new_tokens = gr.Slider(label="Max New Tokens", minimum=500, maximum="3000", step=500, value=1500, interactive=True)
+                else:
+                    num_segments = gr.Number(label="Number of Song Segments", value=2, interactive=True)
+                    max_new_tokens = gr.Slider(label="Max New Tokens", minimum=500, maximum="24000", step=500, value=3000, interactive=True)
+                submit_btn = gr.Button("Submit")
+                music_out = gr.Audio(label="Audio Result")
+
+        gr.Examples(
+            examples = [
+                [
+                    "female blues airy vocal bright vocal piano sad romantic guitar jazz",
+                    """[verse]
+In the quiet of the evening, shadows start to fall
+Whispers of the night wind echo through the hall
+Lost within the silence, I hear your gentle voice
+Guiding me back homeward, making my heart rejoice
+
+[chorus]
+Don't let this moment fade, hold me close tonight
+With you here beside me, everything's alright
+Can't imagine life alone, don't want to let you go
+Stay with me forever, let our love just flow
+                    """
+                ],
+                [
+                    "rap piano street tough piercing vocal hip-hop synthesizer clear vocal male",
+                    """[verse]
+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'm aiming for the top
+Never gonna quit, no, I'm never gonna stop
+Through the highs and lows, I'mma keep it real
+Living out my dreams with this mic and a deal
+                    """
+                ]
+            ], 
+             inputs = [genre_txt, lyrics_txt]
         )
-    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 = torch.as_tensor(codec_ids).to(device)
-    prompt_ids = torch.as_tensor(prompt_ids).to(device)
-    len_prompt = prompt_ids.shape[-1]
-
-    block_list = LogitsProcessorList([BlockTokenRangeProcessor(0, 46358), BlockTokenRangeProcessor(53526, mmtokenizer.vocab_size)])
-
-    for frames_idx in range(codec_ids.shape[1]):
-        cb0 = codec_ids[:, frames_idx:frames_idx + 1]
-        prompt_ids = torch.cat([prompt_ids, cb0], dim=1)
-        input_ids = prompt_ids
-
-        with torch.no_grad():
-            stage2_output = model.generate(
-                input_ids=input_ids,
-                min_new_tokens=7,
-                max_new_tokens=7,
-                eos_token_id=mmtokenizer.eoa,
-                pad_token_id=mmtokenizer.eoa,
-                logits_processor=block_list,
-            )
-
-        assert stage2_output.shape[1] - prompt_ids.shape[1] == 7, f"output new tokens={stage2_output.shape[1] - prompt_ids.shape[1]}"
-        prompt_ids = stage2_output
-
-    if batch_size > 1:
-        output = prompt_ids.cpu().numpy()[:, len_prompt:]
-        output_list = [output[i] for i in range(batch_size)]
-        output = np.concatenate(output_list, axis=0)
-    else:
-        output = prompt_ids[0].cpu().numpy()[len_prompt:]
-
-    return output
-
-def stage2_inference(model, stage1_output_set, output_dir, batch_size=4):
-    stage2_result = []
-    for i in tqdm(range(len(stage1_output_set))):
-        output_filename = os.path.join(output_dir, os.path.basename(stage1_output_set[i]))
-        if os.path.exists(output_filename):
-            continue
-
-        prompt = np.load(stage1_output_set[i]).astype(np.int32)
-        output_duration = prompt.shape[-1] // 50 // 6 * 6
-        num_batch = output_duration // 6
-
-        if num_batch <= batch_size:
-            output = stage2_generate(model, 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_size = batch_size if seg != num_segments - 1 or num_batch % batch_size == 0 else num_batch % batch_size
-                segment = stage2_generate(model, prompt[:, start_idx:end_idx], batch_size=current_batch_size)
-                segments.append(segment)
-            output = np.concatenate(segments, axis=0)
-
-        if output_duration * 50 != prompt.shape[-1]:
-            ending = stage2_generate(model, prompt[:, output_duration * 50:], batch_size=1)
-            output = np.concatenate([output, ending], axis=0)
-        output = codectool_stage2.ids2npy(output)
-
-        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_frequant = sorted(counter.items(), key=lambda x: x[1], reverse=True)[0][0]
-                    fixed_output[i, j] = most_frequant
-        np.save(output_filename, fixed_output)
-        stage2_result.append(output_filename)
-    return stage2_result
-
-# Main Gradio function
-@spaces.GPU()
-def generate_music(genres, lyrics_text, use_audio_prompt, audio_prompt, start_time, end_time, progress=gr.Progress()):
-    progress(0.1, "Running Stage 1 Generation...")
-    stage1_output_set, output_dir = stage1_generate(genres, lyrics_text, use_audio_prompt, audio_prompt, start_time, end_time)
-
-    progress(0.6, "Running Stage 2 Refinement...")
-    stage2_result = stage2_inference(stage2_model, stage1_output_set, output_dir)
-
-    progress(0.8, "Processing Audio...")
-    vocal_decoder, inst_decoder = build_codec_model('./xcodec_mini_infer/decoders/config.yaml', './xcodec_mini_infer/decoders/decoder_131000.pth', './xcodec_mini_infer/decoders/decoder_151000.pth')
-    vocoder_output_dir = os.path.join(output_dir, "vocoder")
-    os.makedirs(vocoder_output_dir, exist_ok=True)
-
-    for npy in stage2_result:
-        if 'instrumental' in npy:
-            process_audio(npy, os.path.join(vocoder_output_dir, 'instrumental.mp3'), False, None, inst_decoder, codec_model)
-        else:
-            process_audio(npy, os.path.join(vocoder_output_dir, 'vocal.mp3'), False, None, vocal_decoder, codec_model)
-
-    return [
-        os.path.join(vocoder_output_dir, 'instrumental.mp3'),
-        os.path.join(vocoder_output_dir, 'vocal.mp3')
-    ]
-
-# Gradio UI
-with gr.Blocks(title="AI Music Generation") as demo:
-    gr.Markdown("# 🎵 AI Music Generation Pipeline")
     
-    with gr.Row():
-        with gr.Column():
-            genre_input = gr.Textbox(label="Genre Tags", placeholder="e.g., Pop, Happy, Female Vocal")
-            lyrics_input = gr.Textbox(label="Lyrics", lines=10, placeholder="Enter lyrics with segments...")
-            use_audio_prompt = gr.Checkbox(label="Use Audio Prompt")
-            audio_input = gr.Audio(label="Reference Audio", type="filepath", visible=False)
-            start_time = gr.Number(label="Start Time (sec)", value=0.0, visible=False)
-            end_time = gr.Number(label="End Time (sec)", value=30.0, visible=False)
-            
-            generate_btn = gr.Button("Generate Music", variant="primary")
-            
-        with gr.Column():
-            vocal_output = gr.Audio(label="Vocal Track", interactive=False)
-            inst_output = gr.Audio(label="Instrumental Track", interactive=False)
-
-    use_audio_prompt.change(
-        lambda x: [gr.update(visible=x), gr.update(visible=x), gr.update(visible=x)],
-        inputs=use_audio_prompt,
-        outputs=[audio_input, start_time, end_time]
+    submit_btn.click(
+        fn = infer, 
+        inputs = [genre_txt, lyrics_txt, num_segments, max_new_tokens],
+        outputs = [music_out]
     )
-
-    generate_btn.click(
-        generate_music,
-        inputs=[genre_input, lyrics_input, use_audio_prompt, audio_input, start_time, end_time],
-        outputs=[vocal_output, inst_output]
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+demo.queue().launch(show_api=False, show_error=True)

inference/codecmanipulator.py

@@ -0,0 +1,203 @@
+import json
+import numpy as np
+import einops
+
+
+class CodecManipulator(object):
+    r"""
+    **mm tokenizer v0.1**
+    see codeclm/hf/mm_tokenizer_v0.1_hf/id2vocab.json
+
+    text tokens: 
+        llama tokenizer 0~31999
+    
+    special tokens: "32000": "<EOD>", "32001": "<SOA>", "32002": "<EOA>", "32003": "<SOI>", "32004": "<EOI>", "32005": "<SOV>", "32006": "<EOV>", "32007": "<s_local>", "32008": "<e_local>", "32009": "<s_global>", "32010": "<e_global>", "32011": "<semantic>", "32012": "<acoustic>", "32013": "<low_level>", "32014": "<dac_16k>", "32015": "<dac_44k>", "32016": "<xcodec>", "32017": "<placeholder>", "32018": "<semantic_mert>", "32019": "<semantic_hubert>", "32020": "<visual>", "32021": "<semanticodec>"
+
+    mm tokens:
+        dac_16k: 4 codebook, 1024 vocab, 32022 - 36117
+        dac_44k: 9 codebook, 1024 vocab, 36118 - 45333
+        xcodec: 12 codebook, 1024 vocab, 45334 - 57621
+        semantic mert: 1024, 57622 - 58645
+        semantic hubert: 512, 58646 - 59157
+        visual: 64000, not included in v0.1
+        semanticodec 100tps 16384: semantic=16384, 59158 - 75541, acoustic=8192, 75542 - 83733
+    """
+    def __init__(self, codec_type, quantizer_begin=None, n_quantizer=None, teacher_forcing=False, data_feature="codec"):
+        self.codec_type = codec_type
+        self.mm_v0_2_cfg = {
+            "dac16k": {"codebook_size": 1024, "num_codebooks": 4, "global_offset": 32022, "sep": ["<dac_16k>"], "fps": 50},
+            "dac44k": {"codebook_size": 1024, "num_codebooks": 9, "global_offset": 36118, "sep": ["<dac_44k>"]},
+            "xcodec": {"codebook_size": 1024, "num_codebooks": 12, "global_offset": 45334, "sep": ["<xcodec>"], "fps": 50},
+            "mert": {"codebook_size": 1024, "global_offset": 57622, "sep": ["<semantic_mert>"]},
+            "hubert": {"codebook_size": 512, "global_offset": 58646, "sep": ["<semantic_hubert>"]},
+            "semantic/s": {"codebook_size": 16384, "num_codebooks": 1, "global_offset": 59158, "sep": ["<semanticodec>", "<semantic>"]},
+            "semantic/a": {"codebook_size": 8192, "num_codebooks": 1, "global_offset": 75542, "sep": ["<semanticodec>", "<acoustic>"]},
+            "semanticodec": {"codebook_size": [16384, 8192], "num_codebooks": 2, "global_offset": 59158, "sep": ["<semanticodec>"], "fps": 50},
+            "special_tokens": {
+                '<EOD>': 32000, '<SOA>': 32001, '<EOA>': 32002, '<SOI>': 32003, '<EOI>': 32004, '<SOV>': 32005, '<EOV>': 32006, '<s_local>': 32007, '<e_local>': 32008, '<s_global>': 32009, '<e_global>': 32010, '<semantic>': 32011, '<acoustic>': 32012, '<stage_1>': 32013, '<dac_16k>': 32014, '<dac_44k>': 32015, '<xcodec>': 32016, '<stage_2>': 32017, '<semantic_mert>': 32018, '<semantic_hubert>': 32019, '<visual>': 32020, '<semanticodec>': 32021
+            },
+            "metadata": {
+                "len": 83734,
+                "text_range": [0, 31999],
+                "special_range": [32000, 32021],
+                "mm_range": [32022, 83733]
+            },
+            "codec_range": {
+                "dac16k": [32022, 36117],
+                "dac44k": [36118, 45333],
+                "xcodec": [45334, 57621],
+                # "hifi16k": [53526, 57621],
+                "mert": [57622, 58645],
+                "hubert": [58646, 59157],
+                "semantic/s": [59158, 75541],
+                "semantic/a": [75542, 83733],
+                "semanticodec": [59158, 83733]
+            }
+        }
+        self.sep = self.mm_v0_2_cfg[self.codec_type]["sep"]
+        self.sep_ids = [self.mm_v0_2_cfg["special_tokens"][s] for s in self.sep]
+        self.codebook_size = self.mm_v0_2_cfg[self.codec_type]["codebook_size"]
+        self.num_codebooks = self.mm_v0_2_cfg[self.codec_type]["num_codebooks"]
+        self.global_offset = self.mm_v0_2_cfg[self.codec_type]["global_offset"]
+        self.fps = self.mm_v0_2_cfg[self.codec_type]["fps"] if "fps" in self.mm_v0_2_cfg[self.codec_type] else None
+
+        self.quantizer_begin = quantizer_begin if quantizer_begin is not None else 0
+        self.n_quantizer = n_quantizer if n_quantizer is not None else self.num_codebooks  
+        self.teacher_forcing = teacher_forcing 
+        self.data_feature = data_feature
+
+
+    def offset_tok_ids(self, x, global_offset=0, codebook_size=2048, num_codebooks=4):
+        """
+        x: (K, T)
+        """
+        if isinstance(codebook_size, int):
+            assert x.max() < codebook_size, f"max(x)={x.max()}, codebook_size={codebook_size}"
+        elif isinstance(codebook_size, list):
+            for i, cs in enumerate(codebook_size):
+                assert x[i].max() < cs, f"max(x)={x[i].max()}, codebook_size={cs}, layer_id={i}"
+        else:
+            raise ValueError(f"codebook_size={codebook_size}")
+        assert x.min() >= 0, f"min(x)={x.min()}"
+        assert x.shape[0] == num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={num_codebooks}, n_quantizer={self.n_quantizer}"
+
+        _x = x.copy()
+        _x = _x.astype(np.uint32)
+        cum_offset = 0
+        quantizer_begin = self.quantizer_begin
+        quantizer_end = quantizer_begin+self.n_quantizer
+        for k in range(self.quantizer_begin, quantizer_end): # k: quantizer_begin to quantizer_end - 1
+            if isinstance(codebook_size, int):
+                _x[k] += global_offset + k * codebook_size
+            elif isinstance(codebook_size, list):
+                _x[k] += global_offset + cum_offset
+                cum_offset += codebook_size[k]
+            else:
+                raise ValueError(f"codebook_size={codebook_size}")
+        return _x[quantizer_begin:quantizer_end]
+
+    def unoffset_tok_ids(self, x, global_offset=0, codebook_size=2048, num_codebooks=4):
+        """
+        x: (K, T)
+        """
+        if isinstance(codebook_size, int):
+            assert x.max() < global_offset + codebook_size * num_codebooks, f"max(x)={x.max()}, codebook_size={codebook_size}"
+        elif isinstance(codebook_size, list):
+            assert x.max() < global_offset + sum(codebook_size), f"max(x)={x.max()}, codebook_size={codebook_size}"
+        assert x.min() >= global_offset, f"min(x)={x.min()}, global_offset={global_offset}"
+        assert x.shape[0] == num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={num_codebooks}, n_quantizer={self.n_quantizer}"
+        
+        _x = x.copy()
+        _x = _x.astype(np.uint32)
+        cum_offset = 0
+        quantizer_begin = self.quantizer_begin
+        quantizer_end = quantizer_begin+self.n_quantizer
+        for k in range(quantizer_begin, quantizer_end):
+            if isinstance(codebook_size, int):
+                _x[k-quantizer_begin] -= global_offset + k * codebook_size
+            elif isinstance(codebook_size, list):
+                _x[k-quantizer_begin] -= global_offset + cum_offset
+                cum_offset += codebook_size[k]
+            else:
+                raise ValueError(f"codebook_size={codebook_size}")
+        return _x
+
+    def flatten(self, x):
+        if len(x.shape) > 2:
+            x = x.squeeze()
+        assert x.shape[0] == self.num_codebooks or x.shape[0] == self.n_quantizer, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={self.num_codebooks}, n_quantizer={self.n_quantizer}"
+        return einops.rearrange(x, 'K T -> (T K)')
+
+    def unflatten(self, x, n_quantizer=None):
+        x = x.squeeze()
+        assert len(x.shape) == 1
+        assert x.shape[0] % self.num_codebooks == 0 or x.shape[0] % self.n_quantizer == 0, \
+            f"x.shape[0]={x.shape[0]}, num_codebooks={self.num_codebooks}, n_quantizer={self.n_quantizer}"
+        if n_quantizer!=self.num_codebooks:
+            return einops.rearrange(x, '(T K) -> K T', K=n_quantizer)
+        return einops.rearrange(x, '(T K) -> K T', K=self.num_codebooks)
+    
+    # def check_codec_type_from_path(self, path):
+    #     if self.codec_type == "hifi16k":
+    #         assert "academicodec_hifi_16k_320d_large_uni" in path
+    
+    def get_codec_type_from_range(self, ids):
+        ids_range = [ids.min(), ids.max()]
+        codec_range = self.mm_v0_2_cfg["codec_range"]
+        for codec_type, r in codec_range.items():
+            if ids_range[0] >= r[0] and ids_range[1] <= r[1]:
+                return codec_type
+        raise ValueError(f"ids_range={ids_range}, codec_range={codec_range}")
+
+    def npy2ids(self, npy):
+        if isinstance(npy, str):
+            data = np.load(npy)
+        elif isinstance(npy, np.ndarray):
+            data = npy
+        else:
+            raise ValueError(f"not supported type: {type(npy)}")
+        # data = data.squeeze()
+
+        assert len(data.shape)==2,  f'data shape: {data.shape} is not (n_codebook, seq_len)'
+        data = self.offset_tok_ids(
+            data, 
+            global_offset=self.global_offset, 
+            codebook_size=self.codebook_size, 
+            num_codebooks=self.num_codebooks, 
+        )
+        data = self.flatten(data)
+        codec_range = self.get_codec_type_from_range(data)
+        assert codec_range == self.codec_type, f"get_codec_type_from_range(data)={codec_range}, self.codec_type={self.codec_type}"
+        data = data.tolist()
+        return data
+    
+    def ids2npy(self, token_ids):
+        # make sure token_ids starts with codebook 0
+        if isinstance(self.codebook_size, int):
+            codebook_0_range = (self.global_offset + self.quantizer_begin*self.codebook_size, self.global_offset + (self.quantizer_begin+1)*self.codebook_size)
+        elif isinstance(self.codebook_size, list):
+            codebook_0_range = (self.global_offset, self.global_offset + self.codebook_size[0])
+        assert token_ids[0] >= codebook_0_range[0] \
+            and token_ids[0] < codebook_0_range[1], f"token_ids[0]={token_ids[self.quantizer_begin]}, codebook_0_range={codebook_0_range}"
+        data = np.array(token_ids)
+        data = self.unflatten(data, n_quantizer=self.n_quantizer)
+        data = self.unoffset_tok_ids(
+            data, 
+            global_offset=self.global_offset, 
+            codebook_size=self.codebook_size, 
+            num_codebooks=self.num_codebooks, 
+        )
+        return data
+
+    def npy_to_json_str(self, npy_path):
+        data = self.npy2ids(npy_path)
+        return json.dumps({"text": data, "src": npy_path, "codec": self.codec_type})
+    
+    def sep(self):
+        return ''.join(self.sep)
+    
+    def sep_ids(self):
+        return self.sep_ids

inference/infer.py

@@ -0,0 +1,459 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
+import argparse
+import torch
+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
+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
+
+
+parser = argparse.ArgumentParser()
+# Model Configuration:
+parser.add_argument("--stage1_model", type=str, default="m-a-p/YuE-s1-7B-anneal-en-cot", help="The model checkpoint path or identifier for the Stage 1 model.")
+parser.add_argument("--stage2_model", type=str, default="m-a-p/YuE-s2-1B-general", help="The model checkpoint path or identifier for the Stage 2 model.")
+parser.add_argument("--max_new_tokens", type=int, default=3000, help="The maximum number of new tokens to generate in one pass during text generation.")
+parser.add_argument("--run_n_segments", type=int, default=2, help="The number of segments to process during the generation.")
+parser.add_argument("--stage2_batch_size", type=int, default=4, help="The batch size used in Stage 2 inference.")
+# Prompt
+parser.add_argument("--genre_txt", type=str, required=True, help="The file path to a text file containing genre tags that describe the musical style or characteristics (e.g., instrumental, genre, mood, vocal timbre, vocal gender). This is used as part of the generation prompt.")
+parser.add_argument("--lyrics_txt", type=str, required=True, help="The file path to a text file containing the lyrics for the music generation. These lyrics will be processed and split into structured segments to guide the generation process.")
+parser.add_argument("--use_audio_prompt", action="store_true", help="If set, the model will use an audio file as a prompt during generation. The audio file should be specified using --audio_prompt_path.")
+parser.add_argument("--audio_prompt_path", type=str, default="", help="The file path to an audio file to use as a reference prompt when --use_audio_prompt is enabled.")
+parser.add_argument("--prompt_start_time", type=float, default=0.0, help="The start time in seconds to extract the audio prompt from the given audio file.")
+parser.add_argument("--prompt_end_time", type=float, default=30.0, help="The end time in seconds to extract the audio prompt from the given audio file.")
+# Output 
+parser.add_argument("--output_dir", type=str, default="./output", help="The directory where generated outputs will be saved.")
+parser.add_argument("--keep_intermediate", action="store_true", help="If set, intermediate outputs will be saved during processing.")
+parser.add_argument("--disable_offload_model", action="store_true", help="If set, the model will not be offloaded from the GPU to CPU after Stage 1 inference.")
+parser.add_argument("--cuda_idx", type=int, default=0)
+# Config for xcodec and upsampler
+parser.add_argument('--basic_model_config', default='./xcodec_mini_infer/final_ckpt/config.yaml', help='YAML files for xcodec configurations.')
+parser.add_argument('--resume_path', default='./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', help='Path to the xcodec checkpoint.')
+parser.add_argument('--config_path', type=str, default='./xcodec_mini_infer/decoders/config.yaml', help='Path to Vocos config file.')
+parser.add_argument('--vocal_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_131000.pth', help='Path to Vocos decoder weights.')
+parser.add_argument('--inst_decoder_path', type=str, default='./xcodec_mini_infer/decoders/decoder_151000.pth', help='Path to Vocos decoder weights.')
+parser.add_argument('-r', '--rescale', action='store_true', help='Rescale output to avoid clipping.')
+
+
+args = parser.parse_args()
+if args.use_audio_prompt and not args.audio_prompt_path:
+    raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
+stage1_model = args.stage1_model
+stage2_model = args.stage2_model
+cuda_idx = args.cuda_idx
+max_new_tokens = args.max_new_tokens
+stage1_output_dir = os.path.join(args.output_dir, f"stage1")
+stage2_output_dir = stage1_output_dir.replace('stage1', 'stage2')
+os.makedirs(stage1_output_dir, exist_ok=True)
+os.makedirs(stage2_output_dir, exist_ok=True)
+
+# load tokenizer and model
+device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
+
+# Now you can use `device` to move your tensors or models to the GPU (if available)
+print(f"Using device: {device}")
+
+mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
+model = AutoModelForCausalLM.from_pretrained(
+    stage1_model, 
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2", # To enable flashattn, you have to install flash-attn
+    )
+model.to(device)
+model.eval()
+
+codectool = CodecManipulator("xcodec", 0, 1)
+codectool_stage2 = CodecManipulator("xcodec", 0, 8)
+model_config = OmegaConf.load(args.basic_model_config)
+codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
+parameter_dict = torch.load(args.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):
+    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 = []
+# Tips:
+# genre tags support instrumental，genre，mood，vocal timbr and vocal gender
+# all kinds of tags are needed
+with open(args.genre_txt) as f:
+    genres = f.read().strip()
+with open(args.lyrics_txt) as f:
+    lyrics = split_lyrics(f.read())
+# intruction
+full_lyrics = "\n".join(lyrics)
+prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"]
+prompt_texts += lyrics
+
+
+random_id = uuid.uuid4()
+output_seq = None
+# Here is suggested decoding config
+top_p = 0.93
+temperature = 1.0
+repetition_penalty = 1.2
+# special tokens
+start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
+end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
+# Format text prompt
+run_n_segments = min(args.run_n_segments+1, len(lyrics))
+for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
+    section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
+    guidance_scale = 1.5 if i <=1 else 1.2
+    if i==0:
+        continue
+    if i==1:
+        if args.use_audio_prompt:
+            audio_prompt = load_audio_mono(args.audio_prompt_path)
+            audio_prompt.unsqueeze_(0)
+            with torch.no_grad():
+                raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=0.5)
+            raw_codes = raw_codes.transpose(0, 1)
+            raw_codes = raw_codes.cpu().numpy().astype(np.int16)
+            # Format audio prompt
+            code_ids = codectool.npy2ids(raw_codes[0])
+            audio_prompt_codec = code_ids[int(args.prompt_start_time *50): int(args.prompt_end_time *50)] # 50 is tps of xcodec
+            audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [mmtokenizer.eoa]
+            sentence_ids = mmtokenizer.tokenize("[start_of_reference]") +  audio_prompt_codec_ids + mmtokenizer.tokenize("[end_of_reference]")
+            head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
+        else:
+            head_id = mmtokenizer.tokenize(prompt_texts[0])
+        prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+    else:
+        prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
+
+    prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device) 
+    input_ids = torch.cat([raw_output, prompt_ids], dim=1) if i > 1 else prompt_ids
+    # Use window slicing in case output sequence exceeds the context of model
+    max_context = 16384-max_new_tokens-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():
+        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=mmtokenizer.eoa,
+            pad_token_id=mmtokenizer.eoa,
+            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
+
+# save raw output and check sanity
+ids = raw_output[0].cpu().numpy()
+soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
+eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
+if len(soa_idx)!=len(eoa_idx):
+    raise ValueError(f'invalid pairs of soa and eoa, Num of soa: {len(soa_idx)}, Num of eoa: {len(eoa_idx)}')
+
+vocals = []
+instrumentals = []
+range_begin = 1 if args.use_audio_prompt else 0
+for i in range(range_begin, len(soa_idx)):
+    codec_ids = ids[soa_idx[i]+1:eoa_idx[i]]
+    if codec_ids[0] == 32016:
+        codec_ids = codec_ids[1:]
+    codec_ids = codec_ids[:2 * (codec_ids.shape[0] // 2)]
+    vocals_ids = codectool.ids2npy(rearrange(codec_ids,"(n b) -> b n", b=2)[0])
+    vocals.append(vocals_ids)
+    instrumentals_ids = codectool.ids2npy(rearrange(codec_ids,"(n b) -> b n", b=2)[1])
+    instrumentals.append(instrumentals_ids)
+vocals = np.concatenate(vocals, axis=1)
+instrumentals = np.concatenate(instrumentals, axis=1)
+vocal_save_path = os.path.join(stage1_output_dir, f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_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 args.disable_offload_model:
+    model.cpu()
+    del model
+    torch.cuda.empty_cache()
+
+print("Stage 2 inference...")
+model_stage2 = AutoModelForCausalLM.from_pretrained(
+    stage2_model, 
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2"
+    )
+model_stage2.to(device)
+model_stage2.eval()
+
+def stage2_generate(model, prompt, batch_size=16):
+    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)
+    
+    # Prepare prompt_ids based on batch size or single input
+    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 = np.concatenate(codec_list, axis=0)
+        prompt_ids = np.concatenate(
+            [
+                np.tile([mmtokenizer.soa, mmtokenizer.stage_1], (batch_size, 1)),
+                codec_ids,
+                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(),  # Flatten the 2D array to 1D
+            np.array([mmtokenizer.stage_2])
+        ]).astype(np.int32)
+        prompt_ids = prompt_ids[np.newaxis, ...]
+
+    codec_ids = torch.as_tensor(codec_ids).to(device)
+    prompt_ids = torch.as_tensor(prompt_ids).to(device)
+    len_prompt = prompt_ids.shape[-1]
+    
+    block_list = LogitsProcessorList([BlockTokenRangeProcessor(0, 46358), BlockTokenRangeProcessor(53526, mmtokenizer.vocab_size)])
+
+    # Teacher forcing generate loop
+    for frames_idx in range(codec_ids.shape[1]):
+        cb0 = codec_ids[:, frames_idx:frames_idx+1]
+        prompt_ids = torch.cat([prompt_ids, cb0], dim=1)
+        input_ids = prompt_ids
+
+        with torch.no_grad():
+            stage2_output = model.generate(input_ids=input_ids, 
+                min_new_tokens=7,
+                max_new_tokens=7,
+                eos_token_id=mmtokenizer.eoa,
+                pad_token_id=mmtokenizer.eoa,
+                logits_processor=block_list,
+            )
+        
+        assert stage2_output.shape[1] - prompt_ids.shape[1] == 7, f"output new tokens={stage2_output.shape[1]-prompt_ids.shape[1]}"
+        prompt_ids = stage2_output
+
+    # Return output based on batch size
+    if batch_size > 1:
+        output = prompt_ids.cpu().numpy()[:, len_prompt:]
+        output_list = [output[i] for i in range(batch_size)]
+        output = np.concatenate(output_list, axis=0)
+    else:
+        output = prompt_ids[0].cpu().numpy()[len_prompt:]
+
+    return output
+
+def stage2_inference(model, stage1_output_set, stage2_output_dir, batch_size=4):
+    stage2_result = []
+    for i in tqdm(range(len(stage1_output_set))):
+        output_filename = os.path.join(stage2_output_dir, os.path.basename(stage1_output_set[i]))
+        
+        if os.path.exists(output_filename):
+            print(f'{output_filename} stage2 has done.')
+            continue
+        
+        # Load the prompt
+        prompt = np.load(stage1_output_set[i]).astype(np.int32)
+        
+        # Only accept 6s segments
+        output_duration = prompt.shape[-1] // 50 // 6 * 6
+        num_batch = output_duration // 6
+        
+        if num_batch <= batch_size:
+            # If num_batch is less than or equal to batch_size, we can infer the entire prompt at once
+            output = stage2_generate(model, prompt[:, :output_duration*50], batch_size=num_batch)
+        else:
+            # If num_batch is greater than batch_size, process in chunks of batch_size
+            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
+                # Ensure the end_idx does not exceed the available length
+                end_idx = min((seg + 1) * batch_size * 300, output_duration*50)  # Adjust the last segment
+                current_batch_size = batch_size if seg != num_segments-1 or num_batch % batch_size == 0 else num_batch % batch_size
+                segment = stage2_generate(
+                    model,
+                    prompt[:, start_idx:end_idx],
+                    batch_size=current_batch_size
+                )
+                segments.append(segment)
+
+            # Concatenate all the segments
+            output = np.concatenate(segments, axis=0)
+        
+        # Process the ending part of the prompt
+        if output_duration*50 != prompt.shape[-1]:
+            ending = stage2_generate(model, prompt[:, output_duration*50:], batch_size=1)
+            output = np.concatenate([output, ending], axis=0)
+        output = codectool_stage2.ids2npy(output)
+
+        # Fix invalid codes (a dirty solution, which may harm the quality of audio)
+        # We are trying to find better one
+        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_frequant = sorted(counter.items(), key=lambda x: x[1], reverse=True)[0][0]
+                    fixed_output[i, j] = most_frequant
+        # save output
+        np.save(output_filename, fixed_output)
+        stage2_result.append(output_filename)
+    return stage2_result
+
+stage2_result = stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_size=args.stage2_batch_size)
+print(stage2_result)
+print('Stage 2 DONE.\n')
+# convert audio tokens to audio
+def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
+    folder_path = os.path.dirname(path)
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path)
+    limit = 0.99
+    max_val = wav.abs().max()
+    wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
+    torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)
+# reconstruct tracks
+recons_output_dir = os.path.join(args.output_dir, "recons")
+recons_mix_dir = os.path.join(recons_output_dir, 'mix')
+os.makedirs(recons_mix_dir, exist_ok=True)
+tracks = []
+for npy in stage2_result:
+    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)
+
+# vocoder to upsample audios
+vocal_decoder, inst_decoder = build_codec_model(args.config_path, args.vocal_decoder_path, args.inst_decoder_path)
+vocoder_output_dir = os.path.join(args.output_dir, 'vocoder')
+vocoder_stems_dir = os.path.join(vocoder_output_dir, 'stems')
+vocoder_mix_dir = os.path.join(vocoder_output_dir, 'mix')
+os.makedirs(vocoder_mix_dir, exist_ok=True)
+os.makedirs(vocoder_stems_dir, exist_ok=True)
+for npy in stage2_result:
+    if 'instrumental' in npy:
+        # Process instrumental
+        instrumental_output = process_audio(
+            npy,
+            os.path.join(vocoder_stems_dir, 'instrumental.mp3'),
+            args.rescale,
+            args,
+            inst_decoder,
+            codec_model
+        )
+    else:
+        # Process vocal
+        vocal_output = process_audio(
+            npy,
+            os.path.join(vocoder_stems_dir, 'vocal.mp3'),
+            args.rescale,
+            args,
+            vocal_decoder,
+            codec_model
+        )
+# mix tracks
+try:
+    mix_output = instrumental_output + vocal_output
+    vocoder_mix = os.path.join(vocoder_mix_dir, os.path.basename(recons_mix))
+    save_audio(mix_output, vocoder_mix, 44100, args.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(args.output_dir, os.path.basename(recons_mix)),
+    cutoff_freq=5500.0
+)

inference/mm_tokenizer_v0.2_hf/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee5c7cbf32da93989f14d9ba635e3e1d1ab2cc88a92908a5ed0f149375f6ee49
+size 1761962

inference/mmtokenizer.py

@@ -0,0 +1,367 @@
+from abc import ABC
+from abc import abstractmethod
+
+
+class AbstractTokenizer(ABC):
+    """Abstract class for tokenizer."""
+
+    def __init__(self, name):
+        self.name = name
+        super().__init__()
+
+    @property
+    @abstractmethod
+    def vocab_size(self):
+        pass
+
+    @property
+    @abstractmethod
+    def vocab(self):
+        """Dictionary from vocab text token to id token."""
+        pass
+
+    @property
+    @abstractmethod
+    def inv_vocab(self):
+        """Dictionary from vocab id token to text token."""
+        pass
+
+    @abstractmethod
+    def tokenize(self, text):
+        pass
+
+    def detokenize(self, token_ids):
+        raise NotImplementedError('detokenizer is not implemented for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def cls(self):
+        raise NotImplementedError('CLS is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def sep(self):
+        raise NotImplementedError('SEP is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def pad(self):
+        raise NotImplementedError('PAD is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def eod(self):
+        raise NotImplementedError('EOD is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+    @property
+    def mask(self):
+        raise NotImplementedError('MASK is not provided for {} '
+                                  'tokenizer'.format(self.name))
+
+
+class _SentencePieceTokenizer(AbstractTokenizer):
+    """SentencePieceTokenizer-Megatron wrapper"""
+
+    def __init__(self, model_file, vocab_extra_ids=0):
+        name = 'SentencePieceTokenizer'
+        super().__init__(name)
+
+        import sentencepiece
+        self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file)
+        self._initalize(vocab_extra_ids)
+
+    def _populate_vocab(self):
+        self._vocab = {}
+        self._inv_vocab = {}
+
+        for i in range(len(self.tokenizer)):
+            t = self.tokenizer.id_to_piece(i)
+            self._inv_vocab[i] = t
+            self._vocab[t] = i
+
+    def _initalize(self, vocab_extra_ids):
+        self._populate_vocab()
+        self._special_tokens = {}
+        self._inv_special_tokens = {}
+
+        self._t5_tokens = []
+
+        def _add_special_token(t):
+            if t not in self._vocab:
+                next_id = len(self._vocab)
+                self._vocab[t] = next_id
+                self._inv_vocab[next_id] = t
+            self._special_tokens[t] = self._vocab[t]
+            self._inv_special_tokens[self._vocab[t]] = t
+
+        _add_special_token('<CLS>')
+        self._cls_id = self._vocab['<CLS>']
+        _add_special_token('<SEP>')
+        self._sep_id = self._vocab['<SEP>']
+        _add_special_token('<EOD>')
+        self._eod_id = self._vocab['<EOD>']
+        _add_special_token('<MASK>')
+        self._mask_id = self._vocab['<MASK>']
+
+        pad_id = self.tokenizer.pad_id()
+        try:
+            pad_token = self.tokenizer.id_to_piece(pad_id)
+        except IndexError:
+            pad_token = '<PAD>'
+        _add_special_token(pad_token)
+        self._pad_id = self._vocab[pad_token]
+
+        bos_id = self.tokenizer.bos_id()
+        try:
+            bos_token = self.tokenizer.id_to_piece(bos_id)
+        except IndexError:
+            bos_token = '<BOS>'
+        _add_special_token(bos_token)
+        self._bos_id = self._vocab[bos_token]
+
+        eos_id = self.tokenizer.eos_id()
+        try:
+            eos_token = self.tokenizer.id_to_piece(eos_id)
+        except IndexError:
+            eos_token = '<EOS>'
+        _add_special_token(eos_token)
+        self._eos_id = self._vocab[eos_token]
+
+        for i in range(vocab_extra_ids):
+            t = "<extra_id_{}>".format(i)
+            _add_special_token(t)
+            self._t5_tokens += [t]
+
+    @property
+    def vocab_size(self):
+        return len(self._vocab)
+
+    @property
+    def vocab(self):
+        return self._vocab
+
+    @property
+    def inv_vocab(self):
+        return self._inv_vocab
+
+    @property
+    def decoder(self):
+        return self._inv_vocab
+
+    @property
+    def encoder(self):
+        return self._vocab
+
+    # From:
+    # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L89
+    def tokenize(self, text):
+        ids = []
+        idx = 0
+
+        while 1:
+            indices = {}
+            for token in self._special_tokens:
+                try:
+                    indices[token] = text[idx:].index(token)
+                except ValueError:
+                    continue
+            if len(indices) == 0:
+                break
+
+            next_token = min(indices, key=indices.get)
+            next_idx = idx + indices[next_token]
+
+            ids.extend(self.tokenizer.encode_as_ids(text[idx:next_idx]))
+            ids.append(self._special_tokens[next_token])
+            idx = next_idx + len(next_token)
+
+        ids.extend(self.tokenizer.encode_as_ids(text[idx:]))
+        return ids
+
+    # From:
+    # https://github.com/NVIDIA/NeMo/blob/c8fa217e811d60d11d014827c7f3845ff6c99ae7/nemo/collections/common/tokenizers/sentencepiece_tokenizer.py#L125
+    def detokenize(self, ids):
+        text = ""
+        last_i = 0
+
+        for i, id in enumerate(ids):
+            if id in self._inv_special_tokens:
+                text += self.tokenizer.decode_ids(ids[last_i:i]) + " "
+                text += self._inv_special_tokens[id] + " "
+                last_i = i + 1
+
+        text += self.tokenizer.decode_ids(ids[last_i:])
+        return text
+
+    @property
+    def cls(self):
+        return self._cls_id
+
+    @property
+    def sep(self):
+        return self._sep_id
+
+    @property
+    def pad(self):
+        return self._pad_id
+
+    @property
+    def bos_token_id(self):
+        return self._bos_id
+
+    @property
+    def bos(self):
+        return self._bos_id
+
+    @property
+    def eod(self):
+        return self._eod_id
+
+    @property
+    def eos_token_id(self):
+        return self._eos_id
+
+    @property
+    def eos(self):
+        return self._eos_id
+
+    @property
+    def mask(self):
+        return self._mask_id
+
+    @property
+    def additional_special_tokens_ids(self):
+        return [self.vocab[k] for k in self._t5_tokens]
+
+class _MMSentencePieceTokenizer(_SentencePieceTokenizer):
+    """SentencePieceTokenizer-Megatron wrapper"""
+
+    def __init__(self, model_file, vocab_extra_ids=0):
+        super().__init__(model_file, vocab_extra_ids)
+
+
+    def _initalize(self, vocab_extra_ids):
+        self._populate_vocab()
+        self._special_tokens = {}
+        self._inv_special_tokens = {}
+
+        self._t5_tokens = []
+
+        def _add_special_token(t):
+            if t not in self._vocab:
+                next_id = len(self._vocab)
+                self._vocab[t] = next_id
+                self._inv_vocab[next_id] = t
+            self._special_tokens[t] = self._vocab[t]
+            self._inv_special_tokens[self._vocab[t]] = t
+
+        _add_special_token('<CLS>')
+        self._cls_id = self._vocab['<CLS>']
+        _add_special_token('<SEP>')
+        self._sep_id = self._vocab['<SEP>']
+        _add_special_token('<EOD>')
+        self._eod_id = self._vocab['<EOD>']
+        _add_special_token('<MASK>')
+        self._mask_id = self._vocab['<MASK>']
+
+        _add_special_token('<SOA>')
+        self._soa_id = self._vocab['<SOA>']
+        _add_special_token('<EOA>')
+        self._eoa_id = self._vocab['<EOA>']
+        _add_special_token('<SOV>')
+        self._sov_id = self._vocab['<SOV>']
+        _add_special_token('<EOV>')
+        self._eov_id = self._vocab['<EOV>']
+        _add_special_token('<SOI>')
+        self._soi_id = self._vocab['<SOI>']
+        _add_special_token('<EOI>')
+        self._eoi_id = self._vocab['<EOI>']
+        _add_special_token('<s_local>')
+        self._s_local_id = self._vocab['<s_local>']
+        _add_special_token('<e_local>')
+        self._e_local_id = self._vocab['<e_local>']
+        _add_special_token('<s_global>')
+        self._s_global_id = self._vocab['<s_global>']
+        _add_special_token('<e_global>')
+        self._e_global_id = self._vocab['<e_global>']
+        _add_special_token('<stage_1>')
+        self._stage_1_id = self._vocab['<stage_1>']
+        _add_special_token('<stage_2>')
+        self._stage_2_id = self._vocab['<stage_2>']
+        pad_id = self.tokenizer.pad_id()
+        try:
+            pad_token = self.tokenizer.id_to_piece(pad_id)
+        except IndexError:
+            pad_token = '<PAD>'
+        _add_special_token(pad_token)
+        self._pad_id = self._vocab[pad_token]
+
+        bos_id = self.tokenizer.bos_id()
+        try:
+            bos_token = self.tokenizer.id_to_piece(bos_id)
+        except IndexError:
+            bos_token = '<BOS>'
+        _add_special_token(bos_token)
+        self._bos_id = self._vocab[bos_token]
+
+        eos_id = self.tokenizer.eos_id()
+        try:
+            eos_token = self.tokenizer.id_to_piece(eos_id)
+        except IndexError:
+            eos_token = '<EOS>'
+        _add_special_token(eos_token)
+        self._eos_id = self._vocab[eos_token]
+
+        for i in range(vocab_extra_ids):
+            t = "<extra_id_{}>".format(i)
+            _add_special_token(t)
+            self._t5_tokens += [t]
+    
+    @property
+    def soa(self):
+        return self._soa_id
+    
+    @property
+    def eoa(self):
+        return self._eoa_id
+    
+    @property
+    def sov(self):
+        return self._sov_id
+    
+    @property
+    def eov(self):
+        return self._eov_id
+    
+    @property
+    def soi(self):
+        return self._soi_id
+    
+    @property
+    def eoi(self):
+        return self._eoi_id
+    
+    @property
+    def s_local(self):
+        return self._s_local_id
+    
+    @property
+    def e_local(self):
+        return self._e_local_id
+    
+    @property
+    def s_global(self):
+        return self._s_global_id
+
+    @property
+    def e_global(self):
+        return self._e_global_id
+    
+    @property
+    def stage_1(self):
+        return self._stage_1_id
+
+    @property
+    def stage_2(self):
+        return self._stage_2_id

inference/prompt_examples/genre.txt

@@ -0,0 +1 @@
+inspiring female uplifting pop airy vocal electronic bright vocal vocal

inference/prompt_examples/lyrics.txt

@@ -0,0 +1,39 @@
+[verse]
+Staring at the sunset, colors paint the sky
+Thoughts of you keep swirling, can't deny
+I know I let you down, I made mistakes
+But I'm here to mend the heart I didn't break
+
+[chorus]
+Every road you take, I'll be one step behind
+Every dream you chase, I'm reaching for the light
+You can't fight this feeling now
+I won't back down
+You know you can't deny it now
+I won't back down
+
+[verse]
+They might say I'm foolish, chasing after you
+But they don't feel this love the way we do
+My heart beats only for you, can't you see?
+I won't let you slip away from me
+
+[chorus]
+Every road you take, I'll be one step behind
+Every dream you chase, I'm reaching for the light
+You can't fight this feeling now
+I won't back down
+You know you can't deny it now
+I won't back down
+
+[bridge]
+No, I won't back down, won't turn around
+Until you're back where you belong
+I'll cross the oceans wide, stand by your side
+Together we are strong
+
+[outro]
+Every road you take, I'll be one step behind
+Every dream you chase, love's the tie that binds
+You can't fight this feeling now
+I won't back down

requirements.txt

@@ -1,8 +1,8 @@
-torch
+torch==2.2.0
+torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu118
 omegaconf
-torchaudio
 einops
-numpy
+numpy<2
 transformers
 sentencepiece
 tqdm
@@ -10,3 +10,6 @@ tensorboard
 descript-audiotools>=0.7.2
 descript-audio-codec
 scipy==1.10.1
+huggingface-hub==0.25.2
+wheel
+#https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu11torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl