Delete inference

inference/codecmanipulator.py

@@ -1,203 +0,0 @@
-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

@@ -1,456 +0,0 @@
-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

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

inference/mmtokenizer.py

@@ -1,367 +0,0 @@
-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

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

inference/prompt_examples/lyrics.txt

@@ -1,39 +0,0 @@
-[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