Upload 13 files

.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

.gitignore

@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

assets/logo/yue.mp3

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

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,456 @@
+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")
+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
+    )
+# to device, if gpu is available
+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

@@ -0,0 +1,12 @@
+torch
+omegaconf
+torchaudio
+einops
+numpy
+transformers
+sentencepiece
+tqdm
+tensorboard
+descript-audiotools>=0.7.2
+descript-audio-codec
+scipy==1.10.1

wav_top_200_tags.json

@@ -0,0 +1,830 @@
+{
+  "genre": [
+    "Pop",
+    "rock",
+    "pop",
+    "electronic",
+    "Classical",
+    "R&B",
+    "Electronic",
+    "Rock",
+    "Folk",
+    "rap",
+    "classical",
+    "soundtrack",
+    "country",
+    "indie-rock",
+    "punk",
+    "hiphop",
+    "folk",
+    "jazz",
+    "Country",
+    "hip-hop",
+    "Hip-hop",
+    "experimental",
+    "Hip Hop",
+    "Funk",
+    "blues",
+    "ambient",
+    "Rap",
+    "Jazz",
+    "Ambient",
+    "New Age",
+    "Blues",
+    "experimental pop",
+    "classic rock",
+    "indie rock",
+    "alternative rock",
+    "Reggae",
+    "Electro pop",
+    "K-pop",
+    "Dance",
+    "Soundtrack",
+    "Hip hop",
+    "80s",
+    "Dancehall",
+    "Disco",
+    "House",
+    "Death Metal",
+    "Thrash Metal",
+    "international",
+    "progressive rock",
+    "hard rock",
+    "instrumental",
+    "Lounge",
+    "house",
+    "Latin",
+    "hardcore",
+    "Metalcore",
+    "Soul",
+    "grunge",
+    "Easy listening",
+    "easylistening",
+    "Indian",
+    "ethno",
+    "Hard rock",
+    "hip hop",
+    "Indie Pop",
+    "Electro",
+    "industrial",
+    "grindcore",
+    "post-rock",
+    "Soul-R&B",
+    "Reggaeton",
+    "World",
+    "latin pop",
+    "Classic Rock",
+    "Latin pop",
+    "Deathcore",
+    "soul",
+    "improvisation",
+    "Chinese",
+    "techno",
+    "Salsa",
+    "indie pop",
+    "Hardcore",
+    "拉丁",
+    "Black metal",
+    " Americana",
+    "dance",
+    "rock nacional",
+    "tejano",
+    "indie",
+    "ambient electronic",
+    "world",
+    "Death metal",
+    "Trap",
+    "avant-garde",
+    "Chillout",
+    "Americana",
+    "new wave",
+    "rnb",
+    "pop rock",
+    "post-hardcore",
+    "singer-songwriter",
+    "pop punk",
+    "Power metal",
+    "indie folk",
+    "opera",
+    "Metal",
+    "African",
+    "instrumental rock",
+    "Gospel",
+    "downtempo",
+    "New Wave",
+    "Electro-pop",
+    "rockabilly",
+    "MPB",
+    "goth rock",
+    "soul-R&B",
+    "Black Metal",
+    "Dubstep",
+    "Eurovision",
+    "Bossa Nova",
+    "bossanova",
+    "民谣",
+    "big band",
+    "Synthpop",
+    "死亡金属",
+    "中国传统音乐",
+    "glam rock",
+    "国际音乐",
+    "latin",
+    "operatic",
+    "Melodic Death Metal",
+    "lounge",
+    " Regional Mexican",
+    "instrumental pop",
+    "emo",
+    "旋律死亡金属",
+    "Pop Rock",
+    "popfolk",
+    " Latin",
+    "poprock",
+    "eurovision",
+    "Ska",
+    "Techno",
+    "disco",
+    "基督教音乐",
+    "Indie rock",
+    "Goregrind",
+    "8-bit",
+    "Pop rock",
+    "screamo",
+    "Dance pop",
+    "Guitar",
+    "chillout",
+    "beats",
+    "Big band",
+    "mpb",
+    "Bluegrass",
+    "流行",
+    "Thrash metal",
+    "easy listening",
+    "Samba",
+    "Heavy metal",
+    "Symphonic metal",
+    "Chanson",
+    "Oriental",
+    "synthpop",
+    "Girl group",
+    "Epic",
+    "Celtic",
+    "Screamo",
+    "Espanol",
+    "Middle Eastern",
+    "electro",
+    " Soul-R&B",
+    " Classic Rock",
+    "Heavy Metal",
+    "dubstep",
+    "民乐",
+    "country rock",
+    "funk",
+    "ska",
+    "Indie Rock",
+    "Choral",
+    "J-rock",
+    "shoegaze",
+    "Rockabilly",
+    "grime",
+    "Italian pop",
+    "摇滚",
+    " latin",
+    "Bolero",
+    " orchestral",
+    "experimental hip-hop",
+    "eurodance",
+    "noise rock",
+    "electro pop",
+    "noise",
+    "Crossover Country",
+    "Glitch"
+  ],
+  "instrument": [
+    "Piano",
+    "drums",
+    "guitar",
+    "electric guitar",
+    "Guitar",
+    "synthesizer",
+    "Synthesizer",
+    "Keyboard",
+    "piano",
+    "Drums",
+    "Violin",
+    "bass",
+    "acoustic guitar",
+    "Bass",
+    "violin",
+    "voice",
+    "vocal",
+    "acousticguitar",
+    "Electric guitar",
+    "Acoustic guitar",
+    "electricguitar",
+    "Voice",
+    "keyboard",
+    "saxophone",
+    "beat",
+    "Drum machine",
+    "Cello",
+    "harmonica",
+    "fiddle",
+    "Percussion",
+    "beatboxing",
+    "Vocal",
+    "鼓",
+    "Saxophone",
+    "keys",
+    "harp",
+    "Keyboards",
+    "keyboards",
+    " harmonica",
+    "singing",
+    "吉他",
+    "贝斯",
+    "钢琴",
+    "beats",
+    "flute",
+    "bass guitar",
+    "drum",
+    "brass",
+    "Flute",
+    "Fiddle",
+    "charango",
+    "Sitar",
+    "strings",
+    "trumpet",
+    "Brass",
+    "Vocals",
+    "Trumpet",
+    "string",
+    "Singing",
+    " banjo",
+    "drum machine",
+    "cello",
+    "Acoustic Guitar",
+    "glockenspiel",
+    "computer",
+    "电吉他",
+    "合成器",
+    "键盘",
+    "mallets",
+    "原声吉他",
+    "Drum",
+    "Bass guitar",
+    "Dholak",
+    "congas",
+    "Electric Guitar",
+    "二胡",
+    "鼓机",
+    "synth",
+    "Strings",
+    "小提琴",
+    "Trombone",
+    "percussion",
+    "弦乐",
+    "electricpiano",
+    "风琴",
+    "oboe",
+    "horns",
+    "Erhu",
+    " synthesizer",
+    "acoustic drums",
+    " pedal steel guitar",
+    " Voice",
+    "Tambourine",
+    "singer-songwriter",
+    "Oud",
+    "Qanun",
+    "electronic",
+    " pedal steel",
+    "rapping",
+    "Funky bass",
+    "guitars",
+    "木吉他",
+    "Alto saxophone",
+    "Ukulele",
+    "扬琴",
+    "oud",
+    "sitar",
+    "打击乐器",
+    "Synth",
+    "organ",
+    "Kanun",
+    "人声",
+    "古筝",
+    " accordion",
+    "bandura",
+    "banjo",
+    "长笛",
+    "pandeira",
+    "turntables",
+    "Alto Saxophone",
+    " slideguitar",
+    " electricguitar",
+    "rap",
+    "harpsichord",
+    "萨克斯管",
+    "maracas",
+    "口琴",
+    "Guitars",
+    "Dobro guitar",
+    "vocals",
+    "choir",
+    "Ableton",
+    "Horns",
+    "AcousticGuitar",
+    "笛子",
+    "synth drums",
+    "Glockenspiel",
+    "Harp",
+    "zither",
+    "Dobro",
+    "Musical instrument",
+    "electric piano",
+    "竖琴",
+    "Horn",
+    "手风琴",
+    "None",
+    "Choir",
+    "铜管乐器",
+    "String",
+    "vocal samples",
+    "trombone",
+    "班卓琴",
+    "hu lu si",
+    "Pandeira",
+    "采样器",
+    " Banjo",
+    "Synth bass",
+    "synth bass",
+    "mallet",
+    " tabla",
+    "dulcimer",
+    "声乐",
+    "Cavaquinho",
+    "大提琴",
+    "toms",
+    "ney",
+    " trumpet",
+    " voice",
+    "低音",
+    "Zither",
+    "shakuhachi",
+    "主唱",
+    " electric guitar",
+    "tambourine",
+    "Turntables",
+    "lyrics",
+    " concertina",
+    " piano",
+    " steel guitar",
+    "Bongos",
+    "Koto",
+    "808 bass",
+    "Marimba",
+    " drums",
+    "Dance",
+    "萨克斯风",
+    "木琴",
+    " bass",
+    "ukulele",
+    "Steel pan",
+    "女声",
+    "键盘乐器",
+    "whistle",
+    "soprano saxophone",
+    "Nylon string guitar",
+    "synth_lead",
+    "电脑",
+    "Shakuhachi",
+    "oboes",
+    "Rap"
+  ],
+  "mood": [
+    "Uplifting",
+    "emotional",
+    "uplifting",
+    "happy",
+    "Inspiring",
+    "romantic",
+    "sad",
+    "Love",
+    "melancholic",
+    "dark",
+    "Upbeat",
+    "Energetic",
+    "Romantic",
+    "Melancholic",
+    "Nostalgic",
+    "Calm",
+    "Hopeful",
+    "melodic",
+    "relaxing",
+    "Romance",
+    "Emotional",
+    "Dreamy",
+    "energetic",
+    "rebellious",
+    "Dance",
+    "inspiring",
+    " introspective",
+    "Confident",
+    "aggressive",
+    "Positive",
+    "calm",
+    "cool",
+    "Happy",
+    "hopeful",
+    "beautiful",
+    "advertising",
+    "angry",
+    "Sad",
+    "relaxed",
+    "Celebratory",
+    "Angry",
+    "Bold",
+    "Introspective",
+    "Optimistic",
+    "sentimental",
+    "optimistic",
+    "Tough",
+    "motivational",
+    "Heartfelt",
+    "Funky",
+    "communication",
+    "Danceable",
+    "vivacious",
+    "love",
+    "commercial",
+    "Vivacious",
+    "heavy",
+    "ballad",
+    "thoughtful",
+    "fast-paced",
+    "Futuristic",
+    "Joyful",
+    "emotion",
+    "Soulful",
+    "attitude",
+    "positive",
+    "epic",
+    "Festive",
+    "Melodic",
+    "Dancy",
+    "Aggressive",
+    "soft",
+    "Calming",
+    "exciting",
+    "dreamy",
+    "Epic",
+    "nostalgic",
+    "powerful",
+    "adventure",
+    "passionate",
+    "Determined",
+    "沟通",
+    "Sensual",
+    "Playful",
+    "street",
+    "heartfelt",
+    "Rebellious",
+    "intense",
+    "Sentimental",
+    "inspirational",
+    "travel",
+    "Adventurous",
+    "atmospheric",
+    "summer",
+    "easygoing",
+    "Cheerful",
+    "Cool",
+    "Dark",
+    "rock",
+    "Inspiration",
+    "Chill",
+    "Intense",
+    "confident",
+    "empowering",
+    "Violent",
+    "Intimate",
+    "longing",
+    " meditative",
+    "Attitude",
+    "romance",
+    "experimental",
+    "at sea",
+    "放松",
+    "chill",
+    "Exciting",
+    "Soothing",
+    "Empowering",
+    "暴力",
+    "Brawny",
+    "cheerful",
+    "Motivational",
+    "Vibraphone",
+    "tough",
+    "determined",
+    "hardcore",
+    "Reflective",
+    "funny",
+    "Peaceful",
+    "loud",
+    "Pensive",
+    "向上",
+    "playful",
+    "Furious",
+    "时尚",
+    "希望",
+    "rough",
+    "Intimacy",
+    "dance",
+    "Vibrant",
+    "Relaxed",
+    "soundscape",
+    "Brutal",
+    "thought-provoking",
+    "success",
+    "sleepy",
+    "Elegant",
+    "children",
+    "intimate",
+    "残酷",
+    "怀旧",
+    "improvisational",
+    "浪漫",
+    "Ambient",
+    "Affectionate",
+    "Gory",
+    "Dramatic",
+    "enthusiastic",
+    "感性",
+    "ambient",
+    "Gentle",
+    "愤怒",
+    "快乐",
+    "黑暗",
+    "brawny",
+    "Seductive",
+    "Dancing",
+    "introspective",
+    "instrumental",
+    "Satisfied",
+    "hard",
+    "史诗",
+    " documentary",
+    " dreamy",
+    "Lively",
+    "child",
+    "sassy",
+    "dissonant",
+    "Emotive",
+    "electronic",
+    "抒情",
+    "meditative",
+    "Gloomy",
+    "groovy",
+    " film",
+    "adventure, emotion",
+    "ambitious",
+    "Spiritual",
+    "christmas",
+    "reminiscent",
+    "saloon",
+    "vintage",
+    "梦幻",
+    "爱",
+    "fast_decay",
+    "Comedy",
+    "Asian",
+    "侵略��",
+    "Admirative",
+    " communication",
+    "忧郁"
+  ],
+  "gender": [
+    "male",
+    "female",
+    "singing",
+    "soprano",
+    "child",
+    "human",
+    "human female voice",
+    "unspecified",
+    "screamo",
+    "mezzo-soprano",
+    "human voice",
+    "not specified",
+    "tenor",
+    "rapping",
+    "singing voice",
+    "squeaky",
+    "童声",
+    "children"
+  ],
+  "timbre": [
+    "bright",
+    "full",
+    "airy",
+    "clear",
+    "mellow",
+    "dark",
+    "rich",
+    "reverb",
+    "light",
+    "crisp",
+    "broad",
+    "powerful",
+    "piercing",
+    "high-pitched",
+    "bass",
+    "deep",
+    "not applicable",
+    "baritone",
+    "not specified",
+    "vibrant",
+    "boomy",
+    "varied",
+    "bouncy",
+    "range",
+    "harsh",
+    " airy",
+    "round",
+    "uplifting",
+    "soft",
+    "husky",
+    "tenor",
+    "pontificate",
+    "aggressive",
+    "neat",
+    "high",
+    "exuberant",
+    "open",
+    "full bodied",
+    "strong",
+    "grainy",
+    "vocal fry",
+    "gravelly",
+    "low",
+    "long_release",
+    "polished",
+    "velvet",
+    "placid",
+    "plastic",
+    "sharp",
+    "robust",
+    "muffled",
+    "distortion",
+    "crunchy",
+    "resonant",
+    "pure",
+    "年轻",
+    "preenched",
+    "gruff",
+    "raspy",
+    "passionate",
+    "nonlinear_env",
+    "high pitched",
+    "athletic",
+    "reedy",
+    "shimmering",
+    "charismatic",
+    "gliding",
+    "raw",
+    "plucky",
+    "loud",
+    "youthful",
+    "thin",
+    "soulful",
+    "smooth",
+    "flat",
+    "tempo-synced",
+    "opulent",
+    "variable",
+    "happy",
+    "prettily",
+    "percussive",
+    "singing voice",
+    "barrel",
+    "breezy",
+    "vocal",
+    "honeyed",
+    "vivacious",
+    "full-bodied",
+    "persuasive",
+    "tender",
+    "potent",
+    "preppy",
+    " raspy",
+    "narrow",
+    "fruity",
+    "whiny",
+    "hollow",
+    "singing",
+    "rapping",
+    "flexible",
+    " alto",
+    "sweet",
+    "agitated",
+    "shaky",
+    "dainty",
+    "明亮",
+    "soprano",
+    "vocal range",
+    "rough",
+    "有力",
+    "成熟",
+    "sultry",
+    "barren",
+    "bulky",
+    "prevalent",
+    "bellowing",
+    "dusty",
+    "elevated",
+    "wide",
+    "rumbly",
+    "shrill",
+    "prettily produced",
+    "projected",
+    "low pitched",
+    "bold",
+    "grassy",
+    "plush",
+    "glorious",
+    "elevated pitch",
+    "whispery",
+    "long",
+    "nasal",
+    "preened",
+    "squeaky",
+    "hellosing",
+    "commanding",
+    "textural",
+    "noble",
+    "frustrated",
+    "warm",
+    "punchy",
+    "pretty",
+    "changeable",
+    "mushy",
+    "vocalist",
+    "gritty",
+    "barking",
+    "human",
+    "bass heavy",
+    "dulcet",
+    " smooth",
+    "young",
+    "rhythmic",
+    "vocals",
+    "helmet",
+    "screamy",
+    "hoarse",
+    "rebellious",
+    "soothing",
+    "童声",
+    "bitter",
+    "为了让声乐更加生动，使用了混响效果。",
+    "barrel-shaped",
+    "reed",
+    "强有力",
+    "低沉",
+    "whimsical",
+    "exaggerated",
+    "温暖",
+    "low-pitched",
+    "emotional",
+    "graceful",
+    "breakable",
+    "screechy",
+    "muddy",
+    "breathy",
+    "柔和",
+    "weathered",
+    "roaring",
+    "青春",
+    "pensive",
+    "textured",
+    "清脆",
+    "melodic",
+    "helmeted",
+    " velvety",
+    "充满活力",
+    "圆润",
+    "preteen",
+    "rhythm",
+    "treble",
+    "shouty",
+    " husky",
+    "medium",
+    "blue",
+    "screeching",
+    "multiphonic",
+    "quaint",
+    "rhytmic",
+    "轻盈"
+  ]
+}