Upload 6 files

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

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
+)

mm_tokenizer_v0.2_hf/tokenizer.model

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

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

prompt_examples/genre.txt

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

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