init

.gitattributes

@@ -20,6 +20,7 @@
 *.pickle filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
+*.th filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text

data/gigaspeech.py

@@ -0,0 +1,156 @@
+import os
+import torch
+import random
+import copy
+import logging
+import shutil
+
+class dataset(torch.utils.data.Dataset):
+    def __init__(self, args, split):
+        super().__init__()
+        self.args = args
+        self.split = split
+        assert self.split in ['train', 'validation', 'test']
+        manifest_fn = os.path.join(self.args.dataset_dir, self.args.manifest_name, self.split+".txt")
+
+        with open(manifest_fn, "r") as rf:
+            data = [l.strip().split("\t") for l in rf.readlines()]
+        lengths_list = [int(item[-1]) for item in data]
+        self.data = []
+        self.lengths_list = []
+        for d, l in zip(data, lengths_list):
+            if l >= self.args.encodec_sr*self.args.audio_min_length:
+                if self.args.drop_long and l > self.args.encodec_sr*self.args.audio_max_length:
+                    continue
+                self.data.append(d)
+                self.lengths_list.append(l)
+        logging.info(f"number of data points for {self.split} split: {len(self.lengths_list)}")
+
+        # phoneme vocabulary
+        vocab_fn = os.path.join(self.args.dataset_dir,"vocab.txt")
+        shutil.copy(vocab_fn, os.path.join(self.args.exp_dir, "vocab.txt"))
+        with open(vocab_fn, "r") as f:
+            temp = [l.strip().split(" ") for l in f.readlines() if len(l) != 0]
+            self.phn2num = {item[1]:int(item[0]) for item in temp}
+        
+        self.symbol_set = set(["<SIL>", "<MUSIC>", "<NOISE>", "<OTHER>"])
+    
+    def __len__(self):
+        return len(self.lengths_list)
+    
+    def _load_phn_enc(self, index):
+        item = self.data[index]
+        pf = os.path.join(self.args.dataset_dir, self.args.phn_folder_name, item[1]+".txt")
+        ef = os.path.join(self.args.dataset_dir, self.args.encodec_folder_name, item[1]+".txt")
+        try:
+            with open(pf, "r") as p, open(ef, "r") as e:
+                phns = [l.strip() for l in p.readlines()]
+                assert len(phns) == 1, phns
+                x = [self.phn2num[item] for item in phns[0].split(" ") if item not in self.symbol_set] # drop ["<SIL>", "<MUSIC>", "<NOISE>", "<OTHER>"], as they are not in training set annotation
+                encos = [l.strip().split() for k, l in enumerate(e.readlines()) if k < self.args.n_codebooks]
+                
+                assert len(encos) == self.args.n_codebooks, ef
+                if self.args.special_first:
+                    y = [[int(n)+self.args.n_special for n in l] for l in encos]
+                else:
+                    y = [[int(n) for n in l] for l in encos]
+        except Exception as e:
+            logging.info(f"loading failed for {pf} and {ef}, maybe files don't exist or are corrupted")
+            logging.info(f"error message: {e}")
+            return [], [[]]
+
+        return x, y
+
+    def __getitem__(self, index):
+        x, y = self._load_phn_enc(index)
+        x_len, y_len = len(x), len(y[0])
+
+        if x_len == 0 or y_len == 0:
+            return {
+            "x": None, 
+            "x_len": None, 
+            "y": None, 
+            "y_len": None, 
+            "y_mask_interval": None, # index y_mask_interval[1] is the position of start_of_continue token
+            "extra_mask_start": None # this is only used in VE1
+            }
+        while y_len < self.args.encodec_sr*self.args.audio_min_length:
+            assert not self.args.dynamic_batching
+            index = random.choice(range(len(self))) # regenerate an index
+            x, y = self._load_phn_enc(index)
+            x_len, y_len = len(x), len(y[0])
+        if self.args.drop_long:
+            while x_len > self.args.text_max_length or y_len > self.args.encodec_sr*self.args.audio_max_length:
+                index = random.choice(range(len(self))) # regenerate an index
+                x, y = self._load_phn_enc(index)
+                x_len, y_len = len(x), len(y[0])
+
+        ### padding and cropping below ###
+        ### padding and cropping below ###
+        # adjust the length of encodec codes, pad to max_len or randomly crop
+        orig_y_len = copy.copy(y_len)
+        max_len = int(self.args.audio_max_length * self.args.encodec_sr)
+        if y_len > max_len:
+            audio_start = random.choice(range(0, y_len-max_len))
+            for i in range(len(y)):
+                y[i] = y[i][audio_start:(audio_start+max_len)]
+            y_len = max_len
+        else:
+            audio_start = 0
+            if not self.args.dynamic_batching:
+                pad = [0] * (max_len - y_len) if self.args.sep_special_token else [self.args.audio_pad_token] * (max_len - y_len)
+                for i in range(len(y)):
+                    y[i] = y[i] + pad
+        
+        # adjust text
+        # if audio is cropped, and text is longer than max, crop max based on how audio is cropped
+        if audio_start > 0 and len(x) > self.args.text_max_length: # if audio is longer than max and text is long than max, start text the way audio started
+            x = x[int(len(x)*audio_start/orig_y_len):]
+            if len(x) > self.args.text_max_length: # if text is still longer than max, cut the end
+                x = x[:self.args.text_max_length]
+        
+        x_len = len(x)
+        if x_len > self.args.text_max_length:
+            text_start = random.choice(range(0, x_len - self.args.text_max_length))
+            x = x[text_start:text_start+self.args.text_max_length]
+            x_len = self.args.text_max_length
+        elif self.args.pad_x and x_len <= self.args.text_max_length:
+            pad = [0] * (self.args.text_max_length - x_len) if self.args.sep_special_token else [self.args.text_pad_token] * (self.args.text_max_length - x_len)
+            x = x + pad
+        ### padding and cropping above ###
+        ### padding and cropping above ###
+
+        return {
+            "x": torch.LongTensor(x), 
+            "x_len": x_len, 
+            "y": torch.LongTensor(y), 
+            "y_len": y_len
+            }
+            
+
+    def collate(self, batch):
+        out = {key:[] for key in batch[0]}
+        for item in batch:
+            if item['x'] == None: # deal with load failure
+                continue
+            for key, val in item.items():
+                out[key].append(val)
+        res = {}
+        if self.args.pad_x:
+            res["x"] = torch.stack(out["x"], dim=0)
+        else:
+            res["x"] = torch.nn.utils.rnn.pad_sequence(out["x"], batch_first=True, padding_value=self.args.text_pad_token)
+        res["x_lens"] = torch.LongTensor(out["x_len"])
+        if self.args.dynamic_batching:
+            if out['y'][0].ndim==2:
+                res['y'] = torch.nn.utils.rnn.pad_sequence([item.transpose(1,0) for item in out['y']],padding_value=self.args.audio_pad_token)
+                res['y'] = res['y'].permute(1,2,0) # T B K -> B K T
+            else:
+                assert out['y'][0].ndim==1, out['y'][0].shape
+                res['y'] = torch.nn.utils.rnn.pad_sequence(out['y'], batch_first=True, padding_value=self.args.audio_pad_token)
+        else:
+            res['y'] = torch.stack(out['y'], dim=0)
+        res["y_lens"] = torch.LongTensor(out["y_len"])
+        res["text_padding_mask"] = torch.arange(res['x'][0].shape[-1]).unsqueeze(0) >= res['x_lens'].unsqueeze(1)
+        res["audio_padding_mask"] = torch.arange(res['y'][0].shape[-1]).unsqueeze(0) >= res['y_lens'].unsqueeze(1)
+        return res

data/phonemize_encodec_encode_hf.py

@@ -0,0 +1,206 @@
+import argparse
+def parse_args():
+    parser = argparse.ArgumentParser(description="encode the librilight dataset using encodec model")
+    parser.add_argument("--dataset_size", type=str, default='xs', help='sizes of gigaspeech, xs, s, m, l, xl. we use xl for VoiceCraft training, xs is good for debugging')
+    parser.add_argument('--download_to', type=str, default="/data/scratch/pyp/datasets/gigaspeech_debug", help="dir where you want the huggingface gigaspeech dataset to be downloaded to")
+    parser.add_argument('--save_dir', type=str, default="/data/scratch/pyp/datasets/gigaspeech_phn_enc_manifest_debug", help="path to the manifest, phonemes, and encodec codes dirs")
+    parser.add_argument('--encodec_model_path', type=str, default="/data/scratch/pyp/exp_pyp/audiocraft/encodec/xps/6f79c6a8/checkpoint.th")
+    parser.add_argument('--n_workers', type=int, default=4, help="Number of parallel worker processes")
+    parser.add_argument('--mega_batch_size', type=int, default=100, help="Number of samples in each mega batch for multiprocess dataloading")
+    parser.add_argument('--batch_size', type=int, default=4, help="batch size for encodec encoding, decrease it if OOM. This is the sum of batch size *over each gpu*, so increase it if you are using more gpus")
+    parser.add_argument('--model_sr', type=int, default=16000, help='encodec input audio sample rate')
+    parser.add_argument('--downsample_rate', type=int, default=320, help='encodec downsample rate')
+    parser.add_argument('--model_code_sr', type=int, default=50, help='encodec model code sample rate')
+    parser.add_argument('--len_cap', type=float, default=35.0, help='will drop audios that are longer than this number')
+    parser.add_argument('--max_len', type=int, default=30000, help='max length of audio in samples, if exceed, will cut a batch into half to process, decrease this number if OOM on your machine')
+    return parser.parse_args()
+if __name__ == "__main__":
+    import logging
+    formatter = (
+        "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d || %(message)s"
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO)
+    args = parse_args()
+
+    import os
+    import numpy as np
+    import torch
+    import tqdm
+    import time
+    from datasets import load_dataset, DownloadConfig
+
+    from tokenizer import TextTokenizer, tokenize_text
+    
+    # get the path
+    phn_save_root = os.path.join(args.save_dir, args.dataset_size, "phonemes")
+    codes_save_root = os.path.join(args.save_dir, args.dataset_size, "encodec_16khz_4codebooks")
+    vocab_fn = os.path.join(args.save_dir, args.dataset_size, "vocab.txt")
+    os.makedirs(phn_save_root, exist_ok=True)
+    os.makedirs(codes_save_root, exist_ok=True)
+
+
+    def sort_by_audio_len(lens):
+        inds = np.argsort(lens).tolist()
+        logging.info(f"longest: {lens[inds[-1]]*args.model_code_sr} encodec codes, {lens[inds[-1]]:.2f} sec.")
+        logging.info(f"shortest: {lens[inds[0]]*args.model_code_sr} encodec codes, {lens[inds[0]]:.2f} sec.")
+        logging.info(f"median: {lens[inds[len(inds)//2]]*args.model_code_sr} encodec codes, {lens[inds[len(inds)//2]]:.2f} sec.")
+        logging.info(f"95 percentile longest: {lens[inds[int(len(inds)*0.95)]]*args.model_code_sr} encodec codes, {lens[inds[int(len(inds)*0.95)]]:.2f} sec.")
+        return inds[::-1]
+    
+    def write_array_to_txt_file(array, filename):
+        with open(filename, 'w') as f:
+            for a in array[:-1]:
+                f.write(' '.join(map(str, a))+'\n')
+            f.write(' '.join(map(str, array[-1])))
+    
+
+    ### phonemization
+    # load tokenizer
+    # load the encodec model
+    from audiocraft.solvers import CompressionSolver
+    model = CompressionSolver.model_from_checkpoint(args.encodec_model_path)
+    model = model.cuda()
+    model = model.eval()
+    text_tokenizer = TextTokenizer()
+
+
+    # https://github.com/SpeechColab/GigaSpeech
+    # there are only four different punctuations
+    # need to check whether there are other < started strings
+    punc2sym = {" <COMMA>": ",", " <PERIOD>": ".", " <QUESTIONMARK>": "?", " <EXCLAMATIONPOINT>": "!"} # note the space in front of each punc name
+    gar2sym = {"<SIL>": "#%#", "<MUSIC>": "##%", "<NOISE>": "%%#", "<OTHER>":"%#%"} # so that they are savely keep as the original sym when using tokenize_text
+    punc2sym.update(gar2sym)
+
+    word2sym = { "h æ ʃ h ɐ ʃ p ɚ s ɛ n t": "<MUSIC>", "h æ ʃ p ɚ s ɛ n t h æ ʃ": "<SIL>", "p ɚ s ɛ n t h ɐ ʃ p ɚ s ɛ n t": "<OTHER>", "p ɚ s ɛ n t p ɚ s ɛ n t h æ ʃ": "<NOISE>"}
+    forbidden_words = set(['#%#', '##%', '%%#', '%#%'])
+
+    dc = DownloadConfig(cache_dir=args.download_to)
+    stime = time.time()
+    logging.info("loading the dataset...")
+    gs = load_dataset("speechcolab/gigaspeech", args.dataset_size, use_auth_token=True, cache_dir = args.download_to, download_config=dc)
+    logging.info(f"time spend on loading the dataset: {time.time() - stime:.2f} seconds")
+
+    splits = ['validation', 'test', 'train']
+    
+    logging.info(f"gigaspeech dataset {args.dataset_size} info: {gs}")
+    logging.info(f"phonemizing...")
+    phn_vocab = set()
+    all_lens = []
+    
+    # you will see a ton of [WARNING] words_mismatch.py:88......, it's not a issue
+    for split in tqdm.tqdm(splits):
+        skip = 0
+        logging.info(f"now processing split {split}...")
+        for item in tqdm.tqdm(gs[split]):
+            save_fn = os.path.join(phn_save_root, item['segment_id']+".txt")
+            text = item['text']
+            if sum(word in forbidden_words for word in text.split(" ")):
+                logging.info(f"skip {item['segment_id']}, because it contains forbiden words. It's transcript: {text}")
+                skip += 1
+                continue
+            for k, v in punc2sym.items():
+                text = text.replace(k, v)
+            phn = tokenize_text(text_tokenizer, text)
+            phn_seq = " ".join(phn)
+            for k, v in word2sym.items():
+                phn_seq = phn_seq.replace(k, v)
+            phn_vocab.update(phn_seq.split(" "))
+            all_lens.append(len(phn_seq.split(" ")))
+            with open(save_fn, "w") as f:
+                f.write(phn_seq)
+        logging.info(f"split {split} has {len(gs[split])} samples in total, skipped {skip} due to forbiden words")
+
+    print(f"phn vocab size: {len(list(phn_vocab))}")
+    print("phn sequence stats: ")
+    print(f"longest: {max(all_lens)}")
+    print(f"shortest: {min(all_lens)}")
+    print(f"median: {np.quantile(all_lens, 0.5)}")
+    print(f"95 percentile longest: {np.quantile(all_lens, 0.95)}")
+    print("write vocabulary to ", vocab_fn)
+    with open(vocab_fn, "w") as f:
+        for i, phn in enumerate(list(phn_vocab)):
+            if i < len(list(phn_vocab)) - 1:
+                f.write(f"{str(i)} {phn}\n")
+            else:
+                f.write(f"{str(i)} {phn}")
+
+    class mydataset(torch.utils.data.Dataset):
+        def __init__(self, split):
+            super().__init__()
+            self.data = gs[split]
+        def __len__(self):
+            return len(self.data)
+        def __getitem__(self, ind):
+            try:
+                segment_id, audio, sr, text, begin_time, end_time = self.data[ind]['segment_id'], torch.from_numpy(self.data[ind]['audio']['array']).float(), self.data[ind]['audio']['sampling_rate'], self.data[ind]['text'], self.data[ind]['begin_time'], self.data[ind]['end_time']
+            except:
+                return None, None, None, None, None, None
+            
+            return segment_id, audio, sr, text, begin_time, end_time
+        def collate(self, batch):
+            res = {'segment_id': [], "audio": [], "sr": [], "text": [], "begin_time": [], "end_time": []}
+            for item in batch:
+                if item[0] != None:
+                    res['segment_id'].append(item[0])
+                    res['audio'].append(item[1])
+                    res['sr'].append(item[2])
+                    res['text'].append(item[3])
+                    res['begin_time'].append(item[4])
+                    res['end_time'].append(item[5])
+            return res
+
+
+    ## encodec codes extraction
+    logging.info("encodec encoding...")
+    train_dataset = mydataset('train')
+    train_loader = torch.torch.utils.data.DataLoader(train_dataset, batch_size=args.mega_batch_size, shuffle=False, drop_last=False, num_workers=args.n_workers, collate_fn=train_dataset.collate)
+    validation_dataset = mydataset('validation')
+    validation_loader = torch.torch.utils.data.DataLoader(validation_dataset, batch_size=args.mega_batch_size, shuffle=False, drop_last=False, num_workers=args.n_workers, collate_fn=validation_dataset.collate)
+    test_dataset = mydataset('test')
+    test_loader = torch.torch.utils.data.DataLoader(test_dataset, batch_size=args.mega_batch_size, shuffle=False, drop_last=False, num_workers=args.n_workers, collate_fn=test_dataset.collate)
+    splits = ['validation', 'test', 'train']
+    loaders = [validation_loader, test_loader, train_loader]
+    # splits = ['validation'] # for debug
+    # loaders = [validation_loader]
+    for split, loader in zip(splits, loaders):
+        skip = 0
+        logging.info(f"now processing split {split}...")
+        mega_n_steps = int(np.ceil(len(gs[split]) / args.mega_batch_size))
+        logging.info(f"partition the split {split} into {mega_n_steps} parts, each has {args.mega_batch_size} samples")
+        for m, mega_batch in enumerate(loader):
+            logging.info(f"====================================")
+            logging.info(f"====================================")
+            logging.info(f"now processing mega step {m+1}/{mega_n_steps}")
+            lengths = np.array(mega_batch['end_time']) - np.array(mega_batch['begin_time'])
+            sorted_inds = sort_by_audio_len(lengths)
+            for j in range(len(sorted_inds))[::-1]:
+                if lengths[sorted_inds[j]] < 0.2 or lengths[sorted_inds[j]] > args.len_cap: # skip samples that are too short (shorter than 0.2s), or too big (bigger than 80s)
+                    skip += 1
+                    del sorted_inds[j]
+            
+            n_steps = int(np.ceil(len(sorted_inds) / args.batch_size))
+            for n in tqdm.tqdm(range(n_steps), disable=True):
+                inds_used = sorted_inds[n*args.batch_size:(n+1)*args.batch_size]
+                audio_batch = [mega_batch['audio'][id] for id in inds_used]
+                sr_batch = [mega_batch['sr'][id] for id in inds_used]
+                segment_id_batch = [mega_batch['segment_id'][id] for id in inds_used]
+                text_batch = [mega_batch['text'][id] for id in inds_used]
+                padded_wav = torch.nn.utils.rnn.pad_sequence(audio_batch, batch_first=True).unsqueeze(1) # [B, T] -> [B, 1, T]
+                all_lens = [lengths[id] for id in inds_used]
+                with torch.no_grad():
+                    if max(all_lens) > args.max_len and len(all_lens) > 1: # NOTE decrease args.max_len if OOM, or chunk it into more than 2 forward passes
+                        codes = []
+                        inwav = padded_wav.cuda()
+                        codes.append(model.encode(inwav[:len(inwav)//2])[0].cpu())
+                        codes.append(model.encode(inwav[len(inwav)//2:])[0].cpu())
+                        codes = torch.cat(codes, dim=0)
+                    else:
+                        encoded_frames = model.encode(padded_wav.cuda())
+                        # logging.info(f"encoded_frames: {encoded_frames[0].shape}")
+                        codes = encoded_frames[0].cpu()
+
+                for i, length in enumerate(all_lens):
+                    save_fn = os.path.join(codes_save_root, segment_id_batch[i]+".txt")
+                    actual_len = round(length * args.model_code_sr) # 320 is downsample rate for this model
+                    cur_code = codes[i].tolist() if type(codes) == list else codes[i, :, :actual_len].tolist()
+                    write_array_to_txt_file(cur_code, save_fn)

data/tokenizer.py

@@ -0,0 +1,149 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/data/tokenizer.py
+# Copyright    2023                            (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Pattern, Union
+
+import numpy as np
+import torch
+import torchaudio
+# from lhotse.features import FeatureExtractor
+# from lhotse.utils import Seconds, compute_num_frames
+from phonemizer.backend import EspeakBackend
+from phonemizer.backend.espeak.language_switch import LanguageSwitch
+from phonemizer.backend.espeak.words_mismatch import WordMismatch
+from phonemizer.punctuation import Punctuation
+from phonemizer.separator import Separator
+
+
+
+class TextTokenizer:
+    """Phonemize Text."""
+
+    def __init__(
+        self,
+        language="en-us",
+        backend="espeak",
+        separator=Separator(word="_", syllable="-", phone="|"),
+        preserve_punctuation=True,
+        punctuation_marks: Union[str, Pattern] = Punctuation.default_marks(),
+        with_stress: bool = False,
+        tie: Union[bool, str] = False,
+        language_switch: LanguageSwitch = "keep-flags",
+        words_mismatch: WordMismatch = "ignore",
+    ) -> None:
+        phonemizer = EspeakBackend(
+            language,
+            punctuation_marks=punctuation_marks,
+            preserve_punctuation=preserve_punctuation,
+            with_stress=with_stress,
+            tie=tie,
+            language_switch=language_switch,
+            words_mismatch=words_mismatch,
+        )
+        
+        self.backend = phonemizer
+        self.separator = separator
+
+    def to_list(self, phonemized: str) -> List[str]:
+        fields = []
+        for word in phonemized.split(self.separator.word):
+            # "ɐ    m|iː|n?"    ɹ|ɪ|z|ɜː|v; h|ɪ|z.
+            pp = re.findall(r"\w+|[^\w\s]", word, re.UNICODE)
+            fields.extend(
+                [p for p in pp if p != self.separator.phone]
+                + [self.separator.word]
+            )
+        assert len("".join(fields[:-1])) == len(phonemized) - phonemized.count(
+            self.separator.phone
+        )
+        return fields[:-1]
+
+    def __call__(self, text, strip=True) -> List[List[str]]:
+        if isinstance(text, str):
+            text = [text]
+
+        phonemized = self.backend.phonemize(
+            text, separator=self.separator, strip=strip, njobs=1
+        )
+        return [self.to_list(p) for p in phonemized]
+
+
+def tokenize_text(tokenizer: TextTokenizer, text: str) -> List[str]:
+    phonemes = tokenizer([text.strip()])
+    return phonemes[0]  # k2symbols
+
+def convert_audio(wav: torch.Tensor, sr: int, target_sr: int, target_channels: int):
+    assert wav.shape[0] in [1, 2], "Audio must be mono or stereo."
+    if target_channels == 1:
+        wav = wav.mean(0, keepdim=True)
+    elif target_channels == 2:
+        *shape, _, length = wav.shape
+        wav = wav.expand(*shape, target_channels, length)
+    elif wav.shape[0] == 1:
+        wav = wav.expand(target_channels, -1)
+    wav = torchaudio.transforms.Resample(sr, target_sr)(wav)
+    return wav
+
+class AudioTokenizer:
+    """EnCodec audio."""
+
+    def __init__(
+        self,
+        device: Any = None,
+        signature = None
+    ) -> None:
+        from audiocraft.solvers import CompressionSolver
+        model = CompressionSolver.model_from_checkpoint(signature)
+        self.sample_rate = model.sample_rate
+        self.channels = model.channels
+        
+        if not device:
+            device = torch.device("cpu")
+            if torch.cuda.is_available():
+                device = torch.device("cuda:0")
+
+        self._device = device
+
+        self.codec = model.to(device)
+
+    @property
+    def device(self):
+        return self._device
+
+    def encode(self, wav: torch.Tensor) -> torch.Tensor:
+        codes = self.codec.encode(wav.to(self.device))
+        return [(codes[0], None)]
+
+    def decode(self, frames: torch.Tensor) -> torch.Tensor:
+        frames = frames[0][0] # [1,4,T]
+        return self.codec.decode(frames)
+    
+
+
+def tokenize_audio(tokenizer: AudioTokenizer, audio_path: str, offset = -1, num_frames=-1):
+    # Load and pre-process the audio waveform
+    if offset != -1 and num_frames!=-1:
+        wav, sr = torchaudio.load(audio_path, frame_offset=offset, num_frames=num_frames)
+    else:
+        wav, sr = torchaudio.load(audio_path)
+    wav = convert_audio(wav, sr, tokenizer.sample_rate, tokenizer.channels)
+    wav = wav.unsqueeze(0)
+
+    # Extract discrete codes from EnCodec
+    with torch.no_grad():
+        encoded_frames = tokenizer.encode(wav)
+    return encoded_frames

demo/temp/84_121550_000074_000000.txt

@@ -0,0 +1 @@
+But when I had approached so near to them The common object, which the sense deceives, Lost not by distance any of its marks,

demo/temp/mfa_alignments/84_121550_000074_000000.csv

@@ -0,0 +1,109 @@
+Begin,End,Label,Type,Speaker
+0.03,0.18,but,words,temp
+0.18,0.32,when,words,temp
+0.32,0.48,i,words,temp
+0.48,0.64,had,words,temp
+0.64,1.19,approached,words,temp
+1.22,1.58,so,words,temp
+1.58,1.91,near,words,temp
+1.91,2.07,to,words,temp
+2.07,2.42,them,words,temp
+2.53,2.61,the,words,temp
+2.61,3.01,common,words,temp
+3.05,3.62,object,words,temp
+3.68,3.93,which,words,temp
+3.93,4.02,the,words,temp
+4.02,4.34,sense,words,temp
+4.34,4.97,deceives,words,temp
+5.04,5.54,lost,words,temp
+5.54,6.0,not,words,temp
+6.0,6.14,by,words,temp
+6.14,6.67,distance,words,temp
+6.79,7.05,any,words,temp
+7.05,7.18,of,words,temp
+7.18,7.34,its,words,temp
+7.34,7.87,marks,words,temp
+0.03,0.06,B,phones,temp
+0.06,0.09,AH1,phones,temp
+0.09,0.18,T,phones,temp
+0.18,0.23,W,phones,temp
+0.23,0.27,EH1,phones,temp
+0.27,0.32,N,phones,temp
+0.32,0.48,AY1,phones,temp
+0.48,0.49,HH,phones,temp
+0.49,0.6,AE1,phones,temp
+0.6,0.64,D,phones,temp
+0.64,0.7,AH0,phones,temp
+0.7,0.83,P,phones,temp
+0.83,0.88,R,phones,temp
+0.88,0.99,OW1,phones,temp
+0.99,1.12,CH,phones,temp
+1.12,1.19,T,phones,temp
+1.22,1.4,S,phones,temp
+1.4,1.58,OW1,phones,temp
+1.58,1.7,N,phones,temp
+1.7,1.84,IH1,phones,temp
+1.84,1.91,R,phones,temp
+1.91,2.01,T,phones,temp
+2.01,2.07,AH0,phones,temp
+2.07,2.13,DH,phones,temp
+2.13,2.3,EH1,phones,temp
+2.3,2.42,M,phones,temp
+2.53,2.55,DH,phones,temp
+2.55,2.61,AH0,phones,temp
+2.61,2.73,K,phones,temp
+2.73,2.85,AA1,phones,temp
+2.85,2.9,M,phones,temp
+2.9,2.95,AH0,phones,temp
+2.95,3.01,N,phones,temp
+3.05,3.22,AA1,phones,temp
+3.22,3.27,B,phones,temp
+3.27,3.34,JH,phones,temp
+3.34,3.48,EH0,phones,temp
+3.48,3.54,K,phones,temp
+3.54,3.62,T,phones,temp
+3.68,3.69,HH,phones,temp
+3.69,3.76,W,phones,temp
+3.76,3.8,IH1,phones,temp
+3.8,3.93,CH,phones,temp
+3.93,3.95,DH,phones,temp
+3.95,4.02,AH0,phones,temp
+4.02,4.12,S,phones,temp
+4.12,4.21,EH1,phones,temp
+4.21,4.27,N,phones,temp
+4.27,4.34,S,phones,temp
+4.34,4.42,D,phones,temp
+4.42,4.45,IH0,phones,temp
+4.45,4.59,S,phones,temp
+4.59,4.79,IY1,phones,temp
+4.79,4.87,V,phones,temp
+4.87,4.97,Z,phones,temp
+5.04,5.12,L,phones,temp
+5.12,5.33,AO1,phones,temp
+5.33,5.42,S,phones,temp
+5.42,5.54,T,phones,temp
+5.54,5.7,N,phones,temp
+5.7,5.89,AA1,phones,temp
+5.89,6.0,T,phones,temp
+6.0,6.05,B,phones,temp
+6.05,6.14,AY1,phones,temp
+6.14,6.24,D,phones,temp
+6.24,6.3,IH1,phones,temp
+6.3,6.38,S,phones,temp
+6.38,6.45,T,phones,temp
+6.45,6.51,AH0,phones,temp
+6.51,6.57,N,phones,temp
+6.57,6.67,S,phones,temp
+6.79,6.89,EH1,phones,temp
+6.89,6.95,N,phones,temp
+6.95,7.05,IY0,phones,temp
+7.05,7.13,AH0,phones,temp
+7.13,7.18,V,phones,temp
+7.18,7.22,IH0,phones,temp
+7.22,7.29,T,phones,temp
+7.29,7.34,S,phones,temp
+7.34,7.39,M,phones,temp
+7.39,7.5,AA1,phones,temp
+7.5,7.58,R,phones,temp
+7.58,7.7,K,phones,temp
+7.7,7.87,S,phones,temp

gradio_app.py

@@ -0,0 +1,528 @@
+import os
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "5" # these are only used if developping locally
+import gradio as gr
+import torch
+import torchaudio
+from data.tokenizer import (
+    AudioTokenizer,
+    TextTokenizer,
+)
+from models import voicecraft
+import io
+import numpy as np
+import random
+import spaces
+
+
+whisper_model, voicecraft_model = None, None
+
+@spaces.GPU(duration=20)
+def seed_everything(seed):
+    if seed != -1:
+        os.environ['PYTHONHASHSEED'] = str(seed)
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+
+@spaces.GPU(duration=120)
+def load_models(whisper_model_choice, voicecraft_model_choice):
+    global whisper_model, voicecraft_model
+
+    if whisper_model_choice is not None:
+        import whisper
+        from whisper.tokenizer import get_tokenizer
+        whisper_model = {
+            "model": whisper.load_model(whisper_model_choice),
+            "tokenizer": get_tokenizer(multilingual=False)
+        }
+
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    voicecraft_name = f"{voicecraft_model_choice}.pth"
+    ckpt_fn = f"./pretrained_models/{voicecraft_name}"
+    encodec_fn = "./pretrained_models/encodec_4cb2048_giga.th"
+    if not os.path.exists(ckpt_fn):
+        os.system(f"wget https://huggingface.co/pyp1/VoiceCraft/resolve/main/{voicecraft_name}\?download\=true")
+        os.system(f"mv {voicecraft_name}\?download\=true ./pretrained_models/{voicecraft_name}")
+    if not os.path.exists(encodec_fn):
+        os.system(f"wget https://huggingface.co/pyp1/VoiceCraft/resolve/main/encodec_4cb2048_giga.th")
+        os.system(f"mv encodec_4cb2048_giga.th ./pretrained_models/encodec_4cb2048_giga.th")
+
+    ckpt = torch.load(ckpt_fn, map_location="cpu")
+    model = voicecraft.VoiceCraft(ckpt["config"])
+    model.load_state_dict(ckpt["model"])
+    model.to(device)
+    model.eval()
+    voicecraft_model = {
+        "ckpt": ckpt,
+        "model": model,
+        "text_tokenizer": TextTokenizer(backend="espeak"),
+        "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
+    }
+
+    return gr.Accordion()
+
+@spaces.GPU(duration=60)
+def transcribe(seed, audio_path):
+    if whisper_model is None:
+        raise gr.Error("Whisper model not loaded")
+    seed_everything(seed)
+    
+    number_tokens = [
+        i
+        for i in range(whisper_model["tokenizer"].eot)
+        if all(c in "0123456789" for c in whisper_model["tokenizer"].decode([i]).removeprefix(" "))
+    ]
+    result = whisper_model["model"].transcribe(audio_path, suppress_tokens=[-1] + number_tokens, word_timestamps=True)
+    words = [word_info for segment in result["segments"] for word_info in segment["words"]]
+    
+    transcript = result["text"]
+    transcript_with_start_time = " ".join([f"{word['start']} {word['word']}" for word in words])
+    transcript_with_end_time = " ".join([f"{word['word']} {word['end']}" for word in words])
+
+    choices = [f"{word['start']} {word['word']} {word['end']}" for word in words]
+
+    return [
+        transcript, transcript_with_start_time, transcript_with_end_time,
+        gr.Dropdown(value=choices[-1], choices=choices, interactive=True), # prompt_to_word
+        gr.Dropdown(value=choices[0], choices=choices, interactive=True), # edit_from_word
+        gr.Dropdown(value=choices[-1], choices=choices, interactive=True), # edit_to_word
+        words
+    ]
+
+
+def get_output_audio(audio_tensors, codec_audio_sr):
+    result = torch.cat(audio_tensors, 1)
+    buffer = io.BytesIO()
+    torchaudio.save(buffer, result, int(codec_audio_sr), format="wav")
+    buffer.seek(0)
+    return buffer.read()
+
+@spaces.GPU(duration=90)
+def run(seed, left_margin, right_margin, codec_audio_sr, codec_sr, top_k, top_p, temperature,
+        stop_repetition, sample_batch_size, kvcache, silence_tokens,
+        audio_path, word_info, transcript, smart_transcript,
+        mode, prompt_end_time, edit_start_time, edit_end_time,
+        split_text, selected_sentence, previous_audio_tensors):
+    if voicecraft_model is None:
+        raise gr.Error("VoiceCraft model not loaded")
+    if smart_transcript and (word_info is None):
+        raise gr.Error("Can't use smart transcript: whisper transcript not found")
+
+    seed_everything(seed)
+    if mode == "Long TTS":
+        if split_text == "Newline":
+            sentences = transcript.split('\n')
+        else:
+            from nltk.tokenize import sent_tokenize
+            sentences = sent_tokenize(transcript.replace("\n", " "))
+    elif mode == "Rerun":
+        colon_position = selected_sentence.find(':')
+        selected_sentence_idx = int(selected_sentence[:colon_position])
+        sentences = [selected_sentence[colon_position + 1:]]
+    else:
+        sentences = [transcript.replace("\n", " ")]
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    info = torchaudio.info(audio_path)
+    audio_dur = info.num_frames / info.sample_rate
+
+    audio_tensors = []
+    inference_transcript = ""
+    for sentence in sentences:
+        decode_config = {"top_k": top_k, "top_p": top_p, "temperature": temperature, "stop_repetition": stop_repetition,
+                         "kvcache": kvcache, "codec_audio_sr": codec_audio_sr, "codec_sr": codec_sr,
+                         "silence_tokens": silence_tokens, "sample_batch_size": sample_batch_size}
+        if mode != "Edit":
+            from inference_tts_scale import inference_one_sample
+
+            if smart_transcript:                
+                target_transcript = ""
+                for word in word_info:
+                    if word["end"] < prompt_end_time:
+                        target_transcript += word["word"]
+                    elif (word["start"] + word["end"]) / 2 < prompt_end_time:
+                        # include part of the word it it's big, but adjust prompt_end_time
+                        target_transcript += word["word"]
+                        prompt_end_time = word["end"]
+                        break
+                    else:
+                        break
+                target_transcript += f" {sentence}"
+            else:
+                target_transcript = sentence
+
+            inference_transcript += target_transcript + "\n"
+
+            prompt_end_frame = int(min(audio_dur, prompt_end_time) * info.sample_rate)
+            _, gen_audio = inference_one_sample(voicecraft_model["model"],
+                                                voicecraft_model["ckpt"]["config"],
+                                                voicecraft_model["ckpt"]["phn2num"],
+                                                voicecraft_model["text_tokenizer"], voicecraft_model["audio_tokenizer"],
+                                                audio_path, target_transcript, device, decode_config,
+                                                prompt_end_frame)
+        else:
+            from inference_speech_editing_scale import inference_one_sample
+
+            if smart_transcript:
+                target_transcript = ""
+                for word in word_info:
+                    if word["start"] < edit_start_time:
+                        target_transcript += word["word"]
+                    else:
+                        break
+                target_transcript += f" {sentence}"
+                for word in word_info:
+                    if word["end"] > edit_end_time:
+                        target_transcript += word["word"]
+            else:
+                target_transcript = sentence
+
+            inference_transcript += target_transcript + "\n"
+
+            morphed_span = (max(edit_start_time - left_margin, 1 / codec_sr), min(edit_end_time + right_margin, audio_dur))
+            mask_interval = [[round(morphed_span[0]*codec_sr), round(morphed_span[1]*codec_sr)]]
+            mask_interval = torch.LongTensor(mask_interval)
+            
+            _, gen_audio = inference_one_sample(voicecraft_model["model"],
+                                                voicecraft_model["ckpt"]["config"],
+                                                voicecraft_model["ckpt"]["phn2num"],
+                                                voicecraft_model["text_tokenizer"], voicecraft_model["audio_tokenizer"],
+                                                audio_path, target_transcript, mask_interval, device, decode_config)
+        gen_audio = gen_audio[0].cpu()
+        audio_tensors.append(gen_audio)
+
+    if mode != "Rerun":
+        output_audio = get_output_audio(audio_tensors, codec_audio_sr)
+        sentences = [f"{idx}: {text}" for idx, text in enumerate(sentences)]
+        component = gr.Dropdown(choices=sentences, value=sentences[0])
+        return output_audio, inference_transcript, component, audio_tensors
+    else:
+        previous_audio_tensors[selected_sentence_idx] = audio_tensors[0]
+        output_audio = get_output_audio(previous_audio_tensors, codec_audio_sr)
+        sentence_audio = get_output_audio(audio_tensors, codec_audio_sr)
+        return output_audio, inference_transcript, sentence_audio, previous_audio_tensors
+    
+    
+def update_input_audio(audio_path):
+    if audio_path is None:
+        return 0, 0, 0
+    
+    info = torchaudio.info(audio_path)
+    max_time = round(info.num_frames / info.sample_rate, 2)
+    return [
+        gr.Slider(maximum=max_time, value=max_time),
+        gr.Slider(maximum=max_time, value=0),
+        gr.Slider(maximum=max_time, value=max_time),
+    ]
+
+    
+def change_mode(mode):
+    tts_mode_controls, edit_mode_controls, edit_word_mode, split_text, long_tts_sentence_editor
+    return [
+        gr.Group(visible=mode != "Edit"),
+        gr.Group(visible=mode == "Edit"),
+        gr.Radio(visible=mode == "Edit"),
+        gr.Radio(visible=mode == "Long TTS"),
+        gr.Group(visible=mode == "Long TTS"),
+    ]
+
+
+def load_sentence(selected_sentence, codec_audio_sr, audio_tensors):
+    if selected_sentence is None:
+        return None
+    colon_position = selected_sentence.find(':')
+    selected_sentence_idx = int(selected_sentence[:colon_position])
+    return get_output_audio([audio_tensors[selected_sentence_idx]], codec_audio_sr)
+
+
+def update_bound_word(is_first_word, selected_word, edit_word_mode):
+    if selected_word is None:
+        return None
+
+    word_start_time = float(selected_word.split(' ')[0])
+    word_end_time = float(selected_word.split(' ')[-1])
+    if edit_word_mode == "Replace half":
+        bound_time = (word_start_time + word_end_time) / 2
+    elif is_first_word:
+        bound_time = word_start_time
+    else:
+        bound_time = word_end_time
+
+    return bound_time
+
+
+def update_bound_words(from_selected_word, to_selected_word, edit_word_mode):
+    return [
+        update_bound_word(True, from_selected_word, edit_word_mode),
+        update_bound_word(False, to_selected_word, edit_word_mode),
+    ]
+
+
+smart_transcript_info = """
+If enabled, the target transcript will be constructed for you:</br>
+ - In TTS and Long TTS mode just write the text you want to synthesize.</br>
+ - In Edit mode just write the text to replace selected editing segment.</br>
+If disabled, you should write the target transcript yourself:</br>
+ - In TTS mode write prompt transcript followed by generation transcript.</br>
+ - In Long TTS select split by newline (<b>SENTENCE SPLIT WON'T WORK</b>) and start each line with a prompt transcript.</br>
+ - In Edit mode write full prompt</br>
+"""
+
+demo_original_transcript = " But when I had approached so near to them, the common object, which the sense deceives, lost not by distance any of its marks."
+
+demo_text = {
+    "TTS": {
+        "smart": "I cannot believe that the same model can also do text to speech synthesis as well!",
+        "regular": "But when I had approached so near to them, the common I cannot believe that the same model can also do text to speech synthesis as well!"
+    },
+    "Edit": {
+        "smart": "saw the mirage of the lake in the distance,",
+        "regular": "But when I saw the mirage of the lake in the distance, which the sense deceives, Lost not by distance any of its marks,"
+    },
+    "Long TTS": {
+        "smart": "You can run generation on a big text!\n"
+                 "Just write it line-by-line. Or sentence-by-sentence.\n"
+                 "If some sentences sound odd, just rerun generation on them, no need to generate the whole text again!",
+        "regular": "But when I had approached so near to them, the common You can run generation on a big text!\n"
+                   "But when I had approached so near to them, the common Just write it line-by-line. Or sentence-by-sentence.\n"
+                   "But when I had approached so near to them, the common If some sentences sound odd, just rerun generation on them, no need to generate the whole text again!"
+    }
+}
+
+all_demo_texts = {vv for k, v in demo_text.items() for kk, vv in v.items()}
+
+demo_words = [
+    "0.03  but 0.18",
+    "0.18  when 0.32",
+    "0.32  i 0.48",
+    "0.48  had 0.64",
+    "0.64  approached 1.19",
+    "1.22  so 1.58",
+    "1.58  near 1.91",
+    "1.91  to 2.07",
+    "2.07  them 2.42",
+    "2.53  the 2.61",
+    "2.61  common 3.01",
+    "3.05  object 3.62",
+    "3.68  which 3.93",
+    "3.93  the 4.02",
+    "4.02  sense 4.34",
+    "4.34  deceives 4.97",
+    "5.04  lost 5.54",
+    "5.54  not 6.00",
+    "6.00  by 6.14",
+    "6.14  distance 6.67",
+    "6.79  any 7.05",
+    "7.05  of 7.18",
+    "7.18  its 7.34",
+    "7.34  marks 7.87"
+]
+
+demo_word_info = [
+    {"word": "but", "start": 0.03, "end": 0.18},
+    {"word": "when", "start": 0.18, "end": 0.32},
+    {"word": "i", "start": 0.32, "end": 0.48},
+    {"word": "had", "start": 0.48, "end": 0.64},
+    {"word": "approached", "start": 0.64, "end": 1.19},
+    {"word": "so", "start": 1.22, "end": 1.58},
+    {"word": "near", "start": 1.58, "end": 1.91},
+    {"word": "to", "start": 1.91, "end": 2.07},
+    {"word": "them", "start": 2.07, "end": 2.42},
+    {"word": "the", "start": 2.53, "end": 2.61},
+    {"word": "common", "start": 2.61, "end": 3.01},
+    {"word": "object", "start": 3.05, "end": 3.62},
+    {"word": "which", "start": 3.68, "end": 3.93},
+    {"word": "the", "start": 3.93, "end": 4.02},
+    {"word": "sense", "start": 4.02, "end": 4.34},
+    {"word": "deceives", "start": 4.34, "end": 4.97},
+    {"word": "lost", "start": 5.04, "end": 5.54},
+    {"word": "not", "start": 5.54, "end": 6.0},
+    {"word": "by", "start": 6.0, "end": 6.14},
+    {"word": "distance", "start": 6.14, "end": 6.67},
+    {"word": "any", "start": 6.79, "end": 7.05},
+    {"word": "of", "start": 7.05, "end": 7.18},
+    {"word": "its", "start": 7.18, "end": 7.34},
+    {"word": "marks", "start": 7.34, "end": 7.87}
+]
+
+
+def update_demo(mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word):
+    if transcript not in all_demo_texts:
+        return transcript, edit_from_word, edit_to_word
+    
+    replace_half = edit_word_mode == "Replace half"
+    change_edit_from_word = edit_from_word == demo_words[2] or edit_from_word == demo_words[3]
+    change_edit_to_word = edit_to_word == demo_words[11] or edit_to_word == demo_words[12]
+    demo_edit_from_word_value = demo_words[2] if replace_half else demo_words[3]
+    demo_edit_to_word_value = demo_words[12] if replace_half else demo_words[11]
+    return [
+        demo_text[mode]["smart" if smart_transcript else "regular"],
+        demo_edit_from_word_value if change_edit_from_word else edit_from_word,
+        demo_edit_to_word_value if change_edit_to_word else edit_to_word,
+    ]
+
+
+with gr.Blocks() as app:
+    with gr.Row():
+        with gr.Column(scale=2):
+            load_models_btn = gr.Button(value="Load models")
+        with gr.Column(scale=5):
+            with gr.Accordion("Select models", open=False) as models_selector:
+                with gr.Row():
+                    voicecraft_model_choice = gr.Radio(label="VoiceCraft model", value="giga830M", choices=["giga330M", "giga830M"])
+                    whisper_model_choice = gr.Radio(label="Whisper model", value="base.en",
+                                                    choices=[None, "tiny.en", "base.en", "small.en", "medium.en", "large"])
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            input_audio = gr.Audio(value="./demo/84_121550_000074_000000.wav", label="Input Audio", type="filepath")
+            with gr.Group():
+                original_transcript = gr.Textbox(label="Original transcript", lines=5, value=demo_original_transcript, interactive=False,
+                                                 info="Use whisper model to get the transcript. Fix it if necessary.")
+                with gr.Accordion("Word start time", open=False):
+                    transcript_with_start_time = gr.Textbox(label="Start time", lines=5, interactive=False, info="Start time before each word")
+                with gr.Accordion("Word end time", open=False):
+                    transcript_with_end_time = gr.Textbox(label="End time", lines=5, interactive=False, info="End time after each word")
+
+                transcribe_btn = gr.Button(value="Transcribe")
+            
+        with gr.Column(scale=3):
+            with gr.Group():
+                transcript = gr.Textbox(label="Text", lines=7, value=demo_text["TTS"]["smart"])
+                with gr.Row():
+                    smart_transcript = gr.Checkbox(label="Smart transcript", value=True)
+                    with gr.Accordion(label="?", open=False):
+                        info = gr.Markdown(value=smart_transcript_info)
+
+                with gr.Row():
+                    mode = gr.Radio(label="Mode", choices=["TTS", "Edit", "Long TTS"], value="TTS")
+                    split_text = gr.Radio(label="Split text", choices=["Newline", "Sentence"], value="Newline",
+                                          info="Split text into parts and run TTS for each part.", visible=False)
+                    edit_word_mode = gr.Radio(label="Edit word mode", choices=["Replace half", "Replace all"], value="Replace half",
+                                              info="What to do with first and last word", visible=False)
+                
+                with gr.Group() as tts_mode_controls:
+                    prompt_to_word = gr.Dropdown(label="Last word in prompt", choices=demo_words, value=demo_words[10], interactive=True)
+                    prompt_end_time = gr.Slider(label="Prompt end time", minimum=0, maximum=7.93, step=0.01, value=3.01)
+
+                with gr.Group(visible=False) as edit_mode_controls:
+                    with gr.Row():
+                        edit_from_word = gr.Dropdown(label="First word to edit", choices=demo_words, value=demo_words[2], interactive=True)
+                        edit_to_word = gr.Dropdown(label="Last word to edit", choices=demo_words, value=demo_words[12], interactive=True)
+                    with gr.Row():
+                        edit_start_time = gr.Slider(label="Edit from time", minimum=0, maximum=7.93, step=0.01, value=0.35)
+                        edit_end_time = gr.Slider(label="Edit to time", minimum=0, maximum=7.93, step=0.01, value=3.75)
+
+                run_btn = gr.Button(value="Run")
+
+        with gr.Column(scale=2):
+            output_audio = gr.Audio(label="Output Audio")
+            with gr.Accordion("Inference transcript", open=False):
+                inference_transcript = gr.Textbox(label="Inference transcript", lines=5, interactive=False,
+                                                  info="Inference was performed on this transcript.")
+            with gr.Group(visible=False) as long_tts_sentence_editor:
+                sentence_selector = gr.Dropdown(label="Sentence", value=None,
+                                                info="Select sentence you want to regenerate")
+                sentence_audio = gr.Audio(label="Sentence Audio", scale=2)
+                rerun_btn = gr.Button(value="Rerun")
+
+    with gr.Row():
+        with gr.Accordion("VoiceCraft config", open=False):
+            seed = gr.Number(label="seed", value=-1, precision=0)
+            left_margin = gr.Number(label="left_margin", value=0.08)
+            right_margin = gr.Number(label="right_margin", value=0.08)
+            codec_audio_sr = gr.Number(label="codec_audio_sr", value=16000)
+            codec_sr = gr.Number(label="codec_sr", value=50)
+            top_k = gr.Number(label="top_k", value=0)
+            top_p = gr.Number(label="top_p", value=0.8)
+            temperature = gr.Number(label="temperature", value=1)
+            stop_repetition = gr.Radio(label="stop_repetition", choices=[-1, 1, 2, 3], value=3,
+                                       info="if there are long silence in the generated audio, reduce the stop_repetition to 3, 2 or even 1, -1 = disabled")
+            sample_batch_size = gr.Number(label="sample_batch_size", value=4, precision=0,
+                                          info="generate this many samples and choose the shortest one")
+            kvcache = gr.Radio(label="kvcache", choices=[0, 1], value=1,
+                                info="set to 0 to use less VRAM, but with slower inference")
+            silence_tokens = gr.Textbox(label="silence tokens", value="[1388,1898,131]")
+
+    
+    audio_tensors = gr.State()
+    word_info = gr.State(value=demo_word_info)
+
+    
+    mode.change(fn=update_demo,
+                inputs=[mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word],
+                outputs=[transcript, edit_from_word, edit_to_word])
+    edit_word_mode.change(fn=update_demo,
+                          inputs=[mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word],
+                          outputs=[transcript, edit_from_word, edit_to_word])
+    smart_transcript.change(fn=update_demo,
+                            inputs=[mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word],
+                            outputs=[transcript, edit_from_word, edit_to_word])
+    
+    load_models_btn.click(fn=load_models,
+                          inputs=[whisper_model_choice, voicecraft_model_choice],
+                          outputs=[models_selector])
+    
+    input_audio.upload(fn=update_input_audio,
+                       inputs=[input_audio],
+                       outputs=[prompt_end_time, edit_start_time, edit_end_time])
+    transcribe_btn.click(fn=transcribe,
+                         inputs=[seed, input_audio],
+                         outputs=[original_transcript, transcript_with_start_time, transcript_with_end_time,
+                                  prompt_to_word, edit_from_word, edit_to_word, word_info])
+
+    mode.change(fn=change_mode,
+                inputs=[mode],
+                outputs=[tts_mode_controls, edit_mode_controls, edit_word_mode, split_text, long_tts_sentence_editor])
+
+    run_btn.click(fn=run,
+                  inputs=[
+                      seed, left_margin, right_margin,
+                      codec_audio_sr, codec_sr,
+                      top_k, top_p, temperature,
+                      stop_repetition, sample_batch_size,
+                      kvcache, silence_tokens,
+                      input_audio, word_info, transcript, smart_transcript,
+                      mode, prompt_end_time, edit_start_time, edit_end_time,
+                      split_text, sentence_selector, audio_tensors
+                  ],
+                  outputs=[output_audio, inference_transcript, sentence_selector, audio_tensors])
+    
+    sentence_selector.change(fn=load_sentence,
+                             inputs=[sentence_selector, codec_audio_sr, audio_tensors],
+                             outputs=[sentence_audio])
+    rerun_btn.click(fn=run,
+                    inputs=[
+                        seed, left_margin, right_margin,
+                        codec_audio_sr, codec_sr,
+                        top_k, top_p, temperature,
+                        stop_repetition, sample_batch_size,
+                        kvcache, silence_tokens,
+                        input_audio, word_info, transcript, smart_transcript,
+                        gr.State(value="Rerun"), prompt_end_time, edit_start_time, edit_end_time,
+                        split_text, sentence_selector, audio_tensors
+                    ],
+                    outputs=[output_audio, inference_transcript, sentence_audio, audio_tensors])
+    
+    prompt_to_word.change(fn=update_bound_word,
+                          inputs=[gr.State(False), prompt_to_word, gr.State("Replace all")],
+                          outputs=[prompt_end_time])
+    edit_from_word.change(fn=update_bound_word,
+                          inputs=[gr.State(True), edit_from_word, edit_word_mode],
+                          outputs=[edit_start_time])
+    edit_to_word.change(fn=update_bound_word,
+                        inputs=[gr.State(False), edit_to_word, edit_word_mode],
+                        outputs=[edit_end_time])
+    edit_word_mode.change(fn=update_bound_words,
+                          inputs=[edit_from_word, edit_to_word, edit_word_mode],
+                          outputs=[edit_start_time, edit_end_time])
+
+
+if __name__ == "__main__":
+    app.launch()

inference_speech_editing_scale.py

@@ -0,0 +1,226 @@
+import argparse, pickle
+import logging
+import os, random
+import numpy as np
+import torch
+import torchaudio
+
+from data.tokenizer import (
+    AudioTokenizer,
+    TextTokenizer,
+    tokenize_audio,
+    tokenize_text
+)
+
+from models import voicecraft
+import argparse, time, tqdm
+
+# this script only works for the musicgen architecture
+def get_args():
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--manifest_fn", type=str, default="path/to/eval_metadata_file")
+    parser.add_argument("--audio_root", type=str, default="path/to/audio_folder")
+    parser.add_argument("--exp_dir", type=str, default="path/to/model_folder")
+    parser.add_argument("--left_margin", type=float, default=0.08, help="extra space on the left to the word boundary")
+    parser.add_argument("--right_margin", type=float, default=0.08, help="extra space on the right to the word boundary")
+    parser.add_argument("--seed", type=int, default=1)
+    parser.add_argument("--codec_audio_sr", type=int, default=16000, help='the sample rate of audio that the codec is trained for')
+    parser.add_argument("--codec_sr", type=int, default=50, help='the sample rate of the codec codes')
+    parser.add_argument("--top_k", type=int, default=-1, help="sampling param")
+    parser.add_argument("--top_p", type=float, default=0.8, help="sampling param")
+    parser.add_argument("--temperature", type=float, default=1.0, help="sampling param")
+    parser.add_argument("--output_dir", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--signature", type=str, default=None, help="path to the encodec model")
+    parser.add_argument("--stop_repetition", type=int, default=2, help="used for inference, when the number of consecutive repetition of a token is bigger than this, stop it")
+    parser.add_argument("--kvcache", type=int, default=1, help='if true, use kv cache, which is 4-8x faster than without')
+    parser.add_argument("--silence_tokens", type=str, default="[1388,1898,131]", help="note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
+    return parser.parse_args()
+
+@torch.no_grad()
+def inference_one_sample(model, model_args, phn2num, text_tokenizer, audio_tokenizer, audio_fn, target_text, mask_interval, device, decode_config):
+    # phonemize
+    text_tokens = [phn2num[phn] for phn in
+            tokenize_text(
+                text_tokenizer, text=target_text.strip()
+            ) if phn in phn2num
+        ]
+    text_tokens = torch.LongTensor(text_tokens).unsqueeze(0)
+    text_tokens_lens = torch.LongTensor([text_tokens.shape[-1]])
+
+    encoded_frames = tokenize_audio(audio_tokenizer, audio_fn)
+    original_audio = encoded_frames[0][0].transpose(2,1) # [1,T,K]
+    assert original_audio.ndim==3 and original_audio.shape[0] == 1 and original_audio.shape[2] == model_args.n_codebooks, original_audio.shape
+    logging.info(f"with direct encodec encoding before input, original audio length: {original_audio.shape[1]} codec frames, which is {original_audio.shape[1]/decode_config['codec_sr']:.2f} sec.")
+
+    # forward
+    stime = time.time()
+    encoded_frames = model.inference(
+        text_tokens.to(device),
+        text_tokens_lens.to(device),
+        original_audio[...,:model_args.n_codebooks].to(device), # [1,T,8]
+        mask_interval=mask_interval.unsqueeze(0).to(device),
+        top_k=decode_config['top_k'],
+        top_p=decode_config['top_p'],
+        temperature=decode_config['temperature'],
+        stop_repetition=decode_config['stop_repetition'],
+        kvcache=decode_config['kvcache'],
+        silence_tokens=eval(decode_config['silence_tokens']) if type(decode_config['silence_tokens']) == str else decode_config['silence_tokens'],
+    ) # output is [1,K,T]
+    logging.info(f"inference on one sample take: {time.time() - stime:.4f} sec.")
+    if type(encoded_frames) == tuple:
+        encoded_frames = encoded_frames[0]
+    logging.info(f"generated encoded_frames.shape: {encoded_frames.shape}, which is {encoded_frames.shape[-1]/decode_config['codec_sr']} sec.")
+    
+
+    # decode (both original and generated)
+    original_sample = audio_tokenizer.decode(
+        [(original_audio.transpose(2,1), None)] # [1,T,8] -> [1,8,T]
+    )
+    generated_sample = audio_tokenizer.decode(
+        [(encoded_frames, None)]
+    )
+
+    return original_sample, generated_sample
+
+def get_model(exp_dir, device=None):
+    with open(os.path.join(exp_dir, "args.pkl"), "rb") as f:
+        model_args = pickle.load(f)
+
+    logging.info("load model weights...")
+    model = voicecraft.VoiceCraft(model_args)
+    ckpt_fn = os.path.join(exp_dir, "best_bundle.pth")
+    ckpt = torch.load(ckpt_fn, map_location='cpu')['model']
+    phn2num = torch.load(ckpt_fn, map_location='cpu')['phn2num']
+    model.load_state_dict(ckpt)
+    del ckpt
+    logging.info("done loading weights...")
+    if device == None:
+        device = torch.device("cpu")
+        if torch.cuda.is_available():
+            device = torch.device("cuda:0")
+    model.to(device)
+    model.eval()
+    return model, model_args, phn2num
+
+
+def get_mask_interval(ali_fn, word_span_ind, editType):
+    with open(ali_fn, "r") as rf:
+        data = [l.strip().split(",") for l in rf.readlines()]
+        data = data[1:]
+    tmp = word_span_ind.split(",")
+    s, e = int(tmp[0]), int(tmp[-1])
+    start = None
+    for j, item in enumerate(data):
+        if j == s and item[3] == "words":
+            if editType == 'insertion':
+                start = float(item[1])
+            else:
+                start = float(item[0])
+        if j == e and item[3] == "words":
+            if editType == 'insertion':
+                end = float(item[0])
+            else:
+                end = float(item[1])
+            assert start != None
+            break
+    return (start, end)
+
+if __name__ == "__main__":
+    def seed_everything(seed):
+        os.environ['PYTHONHASHSEED'] = str(seed)
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+    formatter = (
+        "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d || %(message)s"
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO)
+    args = get_args()
+    # args.device = 'cpu'
+    args.allowed_repeat_tokens = eval(args.allowed_repeat_tokens)
+    seed_everything(args.seed)
+
+    # load model
+    stime = time.time()
+    logging.info(f"loading model from {args.exp_dir}")
+    model, model_args, phn2num = get_model(args.exp_dir)
+    if not os.path.isfile(model_args.exp_dir):
+        model_args.exp_dir = args.exp_dir
+    logging.info(f"loading model done, took {time.time() - stime:.4f} sec")
+
+    # setup text and audio tokenizer
+    text_tokenizer = TextTokenizer(backend="espeak")
+    audio_tokenizer = AudioTokenizer(signature=args.signature) # will also put the neural codec model on gpu
+
+    with open(args.manifest_fn, "r") as rf:
+        manifest = [l.strip().split("\t") for l in rf.readlines()]
+    manifest = manifest[1:]
+    
+    # wav_fn	txt_fn	alingment_fn	num_words	word_span_ind
+    audio_fns = []
+    target_texts = []
+    mask_intervals = []
+    edit_types = []
+    new_spans = []
+    orig_spans = []
+    os.makedirs(args.output_dir, exist_ok=True)
+    if args.crop_concat:
+        mfa_temp = f"{args.output_dir}/mfa_temp"
+        os.makedirs(mfa_temp, exist_ok=True)
+    for item in manifest:
+        audio_fn = os.path.join(args.audio_root, item[0])
+        temp = torchaudio.info(audio_fn)
+        audio_dur = temp.num_frames/temp.sample_rate
+        audio_fns.append(audio_fn)
+        target_text = item[2].split("|")[-1]
+        edit_types.append(item[5].split("|"))
+        new_spans.append(item[4].split("|"))
+        orig_spans.append(item[3].split("|"))
+        target_texts.append(target_text) # the last transcript is the target
+        # mi needs to be created from word_ind_span and alignment_fn, along with args.left_margin and args.right_margin
+        mis = []
+        all_ind_intervals = item[3].split("|")
+        editTypes = item[5].split("|")
+        smaller_indx = []
+        alignment_fn = os.path.join(args.audio_root, "aligned", item[0].replace(".wav", ".csv"))
+        if not os.path.isfile(alignment_fn):
+            alignment_fn = alignment_fn.replace("/aligned/", "/aligned_csv/")
+            assert os.path.isfile(alignment_fn), alignment_fn
+        for ind_inter,editType in zip(all_ind_intervals, editTypes):
+            # print(ind_inter)
+            mi = get_mask_interval(alignment_fn, ind_inter, editType)
+            mi = (max(mi[0] - args.left_margin, 1/args.codec_sr), min(mi[1] + args.right_margin, audio_dur)) # in seconds
+            mis.append(mi)
+            smaller_indx.append(mi[0])
+        ind = np.argsort(smaller_indx)
+        mis = [mis[id] for id in ind]
+        mask_intervals.append(mis)
+
+
+
+    for i, (audio_fn, target_text, mask_interval) in enumerate(tqdm.tqdm(zip(audio_fns, target_texts, mask_intervals))):
+        orig_mask_interval = mask_interval
+        mask_interval = [[round(cmi[0]*args.codec_sr), round(cmi[1]*args.codec_sr)] for cmi in mask_interval]
+        # logging.info(f"i: {i}, mask_interval: {mask_interval}")
+        mask_interval = torch.LongTensor(mask_interval) # [M,2]
+        orig_audio, new_audio = inference_one_sample(model, model_args, phn2num, text_tokenizer, audio_tokenizer, audio_fn, target_text, mask_interval, args.device, vars(args))
+        
+        # save segments for comparison
+        orig_audio, new_audio = orig_audio[0].cpu(), new_audio[0].cpu()
+        # logging.info(f"length of the resynthesize orig audio: {orig_audio.shape}")
+
+        save_fn_new = f"{args.output_dir}/{os.path.basename(audio_fn)[:-4]}_new_seed{args.seed}.wav"
+        
+        torchaudio.save(save_fn_new, new_audio, args.codec_audio_sr)
+
+        save_fn_orig = f"{args.output_dir}/{os.path.basename(audio_fn)[:-4]}_orig.wav"
+        if not os.path.isfile(save_fn_orig):
+            orig_audio, orig_sr = torchaudio.load(audio_fn)
+            if orig_sr != args.codec_audio_sr:
+                orig_audio = torchaudio.transforms.Resample(orig_sr, args.codec_audio_sr)(orig_audio)
+            torchaudio.save(save_fn_orig, orig_audio, args.codec_audio_sr)
+

inference_tts_scale.py

@@ -0,0 +1,190 @@
+import argparse, pickle
+import logging
+import os, random
+import numpy as np
+import torch
+import torchaudio
+
+from data.tokenizer import (
+    AudioTokenizer,
+    TextTokenizer,
+    tokenize_audio,
+    tokenize_text
+)
+
+from models import voicecraft
+import argparse, time, tqdm
+
+
+# this script only works for the musicgen architecture
+def get_args():
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--manifest_fn", type=str, default="path/to/eval_metadata_file")
+    parser.add_argument("--audio_root", type=str, default="path/to/audio_folder")
+    parser.add_argument("--exp_dir", type=str, default="path/to/model_folder")
+    parser.add_argument("--seed", type=int, default=1)
+    parser.add_argument("--codec_audio_sr", type=int, default=16000, help='the sample rate of audio that the codec is trained for')
+    parser.add_argument("--codec_sr", type=int, default=50, help='the sample rate of the codec codes')
+    parser.add_argument("--top_k", type=int, default=0, help="sampling param")
+    parser.add_argument("--top_p", type=float, default=0.8, help="sampling param")
+    parser.add_argument("--temperature", type=float, default=1.0, help="sampling param")
+    parser.add_argument("--output_dir", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--signature", type=str, default=None, help="path to the encodec model")
+    parser.add_argument("--crop_concat", type=int, default=0)
+    parser.add_argument("--stop_repetition", type=int, default=-1, help="used for inference, when the number of consecutive repetition of a token is bigger than this, stop it")
+    parser.add_argument("--kvcache", type=int, default=1, help='if true, use kv cache, which is 4-8x faster than without')
+    parser.add_argument("--sample_batch_size", type=int, default=1, help="batch size for sampling, NOTE that it's not running inference for several samples, but duplicate one input sample batch_size times, and during inference, we only return the shortest generation")
+    parser.add_argument("--silence_tokens", type=str, default="[1388,1898,131]", help="note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
+    return parser.parse_args()
+
+
+@torch.no_grad()
+def inference_one_sample(model, model_args, phn2num, text_tokenizer, audio_tokenizer, audio_fn, target_text, device, decode_config, prompt_end_frame):
+    # phonemize
+    text_tokens = [phn2num[phn] for phn in
+            tokenize_text(
+                text_tokenizer, text=target_text.strip()
+            ) if phn in phn2num
+        ]
+    text_tokens = torch.LongTensor(text_tokens).unsqueeze(0)
+    text_tokens_lens = torch.LongTensor([text_tokens.shape[-1]])
+
+    # encode audio
+    encoded_frames = tokenize_audio(audio_tokenizer, audio_fn, offset=0, num_frames=prompt_end_frame)
+    original_audio = encoded_frames[0][0].transpose(2,1) # [1,T,K]
+    assert original_audio.ndim==3 and original_audio.shape[0] == 1 and original_audio.shape[2] == model_args.n_codebooks, original_audio.shape
+    logging.info(f"original audio length: {original_audio.shape[1]} codec frames, which is {original_audio.shape[1]/decode_config['codec_sr']:.2f} sec.")
+
+    # forward
+    stime = time.time()
+    if decode_config['sample_batch_size'] <= 1:
+        logging.info(f"running inference with batch size 1")
+        concat_frames, gen_frames = model.inference_tts(
+            text_tokens.to(device),
+            text_tokens_lens.to(device),
+            original_audio[...,:model_args.n_codebooks].to(device), # [1,T,8]
+            top_k=decode_config['top_k'],
+            top_p=decode_config['top_p'],
+            temperature=decode_config['temperature'],
+            stop_repetition=decode_config['stop_repetition'],
+            kvcache=decode_config['kvcache'],
+            silence_tokens=eval(decode_config['silence_tokens']) if type(decode_config['silence_tokens'])==str else decode_config['silence_tokens']
+        ) # output is [1,K,T]
+    else:
+        logging.info(f"running inference with batch size {decode_config['sample_batch_size']}, i.e. return the shortest among {decode_config['sample_batch_size']} generations.")
+        concat_frames, gen_frames = model.inference_tts_batch(
+            text_tokens.to(device),
+            text_tokens_lens.to(device),
+            original_audio[...,:model_args.n_codebooks].to(device), # [1,T,8]
+            top_k=decode_config['top_k'],
+            top_p=decode_config['top_p'],
+            temperature=decode_config['temperature'],
+            stop_repetition=decode_config['stop_repetition'],
+            kvcache=decode_config['kvcache'],
+            batch_size = decode_config['sample_batch_size'],
+            silence_tokens=eval(decode_config['silence_tokens']) if type(decode_config['silence_tokens'])==str else decode_config['silence_tokens']
+        ) # output is [1,K,T]
+    logging.info(f"inference on one sample take: {time.time() - stime:.4f} sec.")
+
+    logging.info(f"generated encoded_frames.shape: {gen_frames.shape}, which is {gen_frames.shape[-1]/decode_config['codec_sr']} sec.")
+    
+    # for timestamp, codes in enumerate(gen_frames[0].transpose(1,0)):
+    #     logging.info(f"{timestamp}: {codes.tolist()}")
+    # decode (both original and generated)
+    concat_sample = audio_tokenizer.decode(
+        [(concat_frames, None)] # [1,T,8] -> [1,8,T]
+    )
+    gen_sample = audio_tokenizer.decode(
+        [(gen_frames, None)]
+    )
+
+    # return
+    return concat_sample, gen_sample
+
+def get_model(exp_dir, device=None):
+    with open(os.path.join(exp_dir, "args.pkl"), "rb") as f:
+        model_args = pickle.load(f)
+
+    logging.info("load model weights...")
+    model = voicecraft.VoiceCraft(model_args)
+    ckpt_fn = os.path.join(exp_dir, "best_bundle.pth")
+    ckpt = torch.load(ckpt_fn, map_location='cpu')['model']
+    phn2num = torch.load(ckpt_fn, map_location='cpu')['phn2num']
+    model.load_state_dict(ckpt)
+    del ckpt
+    logging.info("done loading weights...")
+    if device == None:
+        device = torch.device("cpu")
+        if torch.cuda.is_available():
+            device = torch.device("cuda:0")
+    model.to(device)
+    model.eval()
+    return model, model_args, phn2num
+
+if __name__ == "__main__":
+    def seed_everything(seed):
+        os.environ['PYTHONHASHSEED'] = str(seed)
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+    formatter = (
+        "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d || %(message)s"
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO)
+    args = get_args()
+    # args.device='cpu'
+    seed_everything(args.seed)
+    
+    os.makedirs(args.output_dir, exist_ok=True)
+    # load model
+
+    with open(args.manifest_fn, "r") as rf:
+        manifest = [l.strip().split("\t") for l in rf.readlines()]
+    manifest = manifest[1:]
+    manifest = [[item[0], item[2], item[3], item[1], item[5]] for item in manifest]
+    
+    stime = time.time()
+    logging.info(f"loading model from {args.exp_dir}")
+    model, model_args, phn2num = get_model(args.exp_dir)
+    logging.info(f"loading model done, took {time.time() - stime:.4f} sec")
+
+    # setup text and audio tokenizer
+    text_tokenizer = TextTokenizer(backend="espeak")
+    audio_tokenizer = AudioTokenizer(signature=args.signature) # will also put the neural codec model on gpu
+    
+    audio_fns = []
+    texts = []
+    prompt_end_frames = []
+    new_audio_fns = []
+    text_to_syn = []
+
+    for item in manifest:
+        audio_fn = os.path.join(args.audio_root, item[0])
+        audio_fns.append(audio_fn)
+        temp = torchaudio.info(audio_fn)
+        prompt_end_frames.append(round(float(item[2])*temp.sample_rate))
+        texts.append(item[1])
+        new_audio_fns.append(item[-2])
+        all_text = item[1].split(" ")
+        start_ind = int(item[-1].split(",")[0])
+        text_to_syn.append(" ".join(all_text[start_ind:]))
+
+    for i, (audio_fn, text, prompt_end_frame, new_audio_fn, to_syn) in enumerate(tqdm.tqdm((zip(audio_fns, texts, prompt_end_frames, new_audio_fns, text_to_syn)))):
+        output_expected_sr = args.codec_audio_sr
+        concated_audio, gen_audio = inference_one_sample(model, model_args, phn2num, text_tokenizer, audio_tokenizer, audio_fn, text, args.device, vars(args), prompt_end_frame)
+    
+        # save segments for comparison
+        concated_audio, gen_audio = concated_audio[0].cpu(), gen_audio[0].cpu()
+        if output_expected_sr != args.codec_audio_sr:
+            gen_audio = torchaudio.transforms.Resample(output_expected_sr, args.codec_audio_sr)(gen_audio)
+            concated_audio = torchaudio.transforms.Resample(output_expected_sr, args.codec_audio_sr)(concated_audio)
+
+        seg_save_fn_gen = f"{args.output_dir}/gen_{new_audio_fn[:-4]}_{i}_seed{args.seed}.wav"
+        seg_save_fn_concat = f"{args.output_dir}/concat_{new_audio_fn[:-4]}_{i}_seed{args.seed}.wav"        
+
+        torchaudio.save(seg_save_fn_gen, gen_audio, args.codec_audio_sr)
+        torchaudio.save(seg_save_fn_concat, concated_audio, args.codec_audio_sr)

models/codebooks_patterns.py

@@ -0,0 +1,538 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections import namedtuple
+from dataclasses import dataclass
+from functools import lru_cache
+import logging
+import typing as tp
+
+from abc import ABC, abstractmethod
+import torch
+
+LayoutCoord = namedtuple('LayoutCoord', ['t', 'q'])  # (timestep, codebook index)
+PatternLayout = tp.List[tp.List[LayoutCoord]]  # Sequence of coordinates
+
+
+@dataclass
+class Pattern:
+    """Base implementation of a pattern over a sequence with multiple codebooks.
+
+    The codebook pattern consists in a layout, defining for each sequence step
+    the list of coordinates of each codebook timestep in the resulting interleaved sequence.
+    The first item of the pattern is always an empty list in order to properly insert a special token
+    to start with. For convenience, we also keep track of ``n_q`` the number of codebooks used for the pattern
+    and ``timesteps`` the number of timesteps corresponding to the original sequence.
+
+    The pattern provides convenient methods to build and revert interleaved sequences from it:
+    ``build_pattern_sequence`` maps a given a dense input tensor of multi-codebook sequence from [B, K, T]
+        to the interleaved sequence of shape [B, K, S] applying the pattern, with S being the batch size,
+        K being the number of codebooks, T the number of original timesteps and S the number of sequence steps
+        for the output sequence. The unfilled positions are replaced with a special token and the built sequence
+        is returned along with a mask indicating valid tokens.
+    ``revert_pattern_sequence`` maps back an interleaved sequence of shape [B, K, S] to the original alignment
+        of codebooks across timesteps to an output tensor of shape [B, K, T], using again a special token and a mask
+        to fill and specify invalid positions if needed.
+    See the dedicated methods for more details.
+    """
+    # Pattern layout, for each sequence step, we have a list of coordinates
+    # corresponding to the original codebook timestep and position.
+    # The first list is always an empty list in order to properly insert
+    # a special token to start with.
+    layout: PatternLayout
+    timesteps: int
+    n_q: int
+
+    def __post_init__(self):
+        assert len(self.layout) > 0
+        assert self.layout[0] == []
+        self._validate_layout()
+        self._build_reverted_sequence_scatter_indexes = lru_cache(100)(self._build_reverted_sequence_scatter_indexes)
+        self._build_pattern_sequence_scatter_indexes = lru_cache(100)(self._build_pattern_sequence_scatter_indexes)
+        # logging.info("New pattern, time steps: %d, sequence steps: %d", self.timesteps, len(self.layout))
+
+    def _validate_layout(self):
+        """Runs checks on the layout to ensure a valid pattern is defined.
+        A pattern is considered invalid if:
+            - Multiple timesteps for a same codebook are defined in the same sequence step
+            - The timesteps for a given codebook are not in ascending order as we advance in the sequence
+              (this would mean that we have future timesteps before past timesteps).
+        """
+        q_timesteps = {q: 0 for q in range(self.n_q)}
+        for s, seq_coords in enumerate(self.layout):
+            if len(seq_coords) > 0:
+                qs = set()
+                for coord in seq_coords:
+                    qs.add(coord.q)
+                    last_q_timestep = q_timesteps[coord.q]
+                    assert coord.t >= last_q_timestep, \
+                        f"Past timesteps are found in the sequence for codebook = {coord.q} at step {s}"
+                    q_timesteps[coord.q] = coord.t
+                # each sequence step contains at max 1 coordinate per codebook
+                assert len(qs) == len(seq_coords), \
+                    f"Multiple entries for a same codebook are found at step {s}"
+
+    @property
+    def num_sequence_steps(self):
+        return len(self.layout) - 1
+
+    @property
+    def max_delay(self):
+        max_t_in_seq_coords = 0
+        for seq_coords in self.layout[1:]:
+            for coords in seq_coords:
+                max_t_in_seq_coords = max(max_t_in_seq_coords, coords.t + 1)
+        return max_t_in_seq_coords - self.timesteps
+
+    @property
+    def valid_layout(self):
+        valid_step = len(self.layout) - self.max_delay
+        return self.layout[:valid_step]
+
+    def get_sequence_coords_with_timestep(self, t: int, q: tp.Optional[int] = None):
+        """Get codebook coordinates in the layout that corresponds to the specified timestep t
+        and optionally to the codebook q. Coordinates are returned as a tuple with the sequence step
+        and the actual codebook coordinates.
+        """
+        assert t <= self.timesteps, "provided timesteps is greater than the pattern's number of timesteps"
+        if q is not None:
+            assert q <= self.n_q, "provided number of codebooks is greater than the pattern's number of codebooks"
+        coords = []
+        for s, seq_codes in enumerate(self.layout):
+            for code in seq_codes:
+                if code.t == t and (q is None or code.q == q):
+                    coords.append((s, code))
+        return coords
+
+    def get_steps_with_timestep(self, t: int, q: tp.Optional[int] = None) -> tp.List[int]:
+        return [step for step, coords in self.get_sequence_coords_with_timestep(t, q)]
+
+    def get_first_step_with_timesteps(self, t: int, q: tp.Optional[int] = None) -> tp.Optional[int]:
+        steps_with_timesteps = self.get_steps_with_timestep(t, q)
+        return steps_with_timesteps[0] if len(steps_with_timesteps) > 0 else None
+
+    def _build_pattern_sequence_scatter_indexes(self, timesteps: int, n_q: int, keep_only_valid_steps: bool,
+                                                device: tp.Union[torch.device, str] = 'cpu'):
+        """Build scatter indexes corresponding to the pattern, up to the provided sequence_steps.
+
+        Args:
+            timesteps (int): Maximum number of timesteps steps to consider.
+            keep_only_valid_steps (bool): Restrict the pattern layout to match only valid steps.
+            device (Union[torch.device, str]): Device for created tensors.
+        Returns:
+            indexes (torch.Tensor): Indexes corresponding to the sequence, of shape [K, S].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes, of shape [K, S].
+        """
+        assert n_q == self.n_q, f"invalid number of codebooks for the sequence and the pattern: {n_q} != {self.n_q}"
+        assert timesteps <= self.timesteps, "invalid number of timesteps used to build the sequence from the pattern"
+        # use the proper layout based on whether we limit ourselves to valid steps only or not,
+        # note that using the valid_layout will result in a truncated sequence up to the valid steps
+        ref_layout = self.valid_layout if keep_only_valid_steps else self.layout
+        # single item indexing being super slow with pytorch vs. numpy, so we use numpy here
+        indexes = torch.zeros(n_q, len(ref_layout), dtype=torch.long).numpy()
+        mask = torch.zeros(n_q, len(ref_layout), dtype=torch.bool).numpy()
+        # fill indexes with last sequence step value that will correspond to our special token
+        # the last value is n_q * timesteps as we have flattened z and append special token as the last token
+        # which will correspond to the index: n_q * timesteps
+        indexes[:] = n_q * timesteps
+        # iterate over the pattern and fill scattered indexes and mask
+        for s, sequence_coords in enumerate(ref_layout):
+            for coords in sequence_coords:
+                if coords.t < timesteps:
+                    indexes[coords.q, s] = coords.t + coords.q * timesteps
+                    mask[coords.q, s] = 1
+        indexes = torch.from_numpy(indexes).to(device)
+        mask = torch.from_numpy(mask).to(device)
+        return indexes, mask
+
+    def build_pattern_sequence(self, z: torch.Tensor, special_token: int, keep_only_valid_steps: bool = False):
+        """Build sequence corresponding to the pattern from the input tensor z.
+        The sequence is built using up to sequence_steps if specified, and non-pattern
+        coordinates are filled with the special token.
+
+        Args:
+            z (torch.Tensor): Input tensor of multi-codebooks sequence, of shape [B, K, T].
+            special_token (int): Special token used to fill non-pattern coordinates in the new sequence.
+            keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.
+                Steps that are beyond valid steps will be replaced by the special_token in that case.
+        Returns:
+            values (torch.Tensor): Interleaved sequence matching the pattern, of shape [B, K, S] with S
+                corresponding either to the sequence_steps if provided, otherwise to the length of the pattern.
+            indexes (torch.Tensor): Indexes corresponding to the interleaved sequence, of shape [K, S].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, S].
+        """
+        B, K, T = z.shape
+        indexes, mask = self._build_pattern_sequence_scatter_indexes(
+            T, K, keep_only_valid_steps=keep_only_valid_steps, device=str(z.device)
+        )
+        z = z.view(B, -1)
+        # we append the special token as the last index of our flattened z tensor
+        z = torch.cat([z, torch.zeros_like(z[:, :1]) + special_token], dim=1)
+        values = z[:, indexes.view(-1)]
+        values = values.view(B, K, indexes.shape[-1])
+        return values, indexes, mask
+
+    def _build_reverted_sequence_scatter_indexes(self, sequence_steps: int, n_q: int,
+                                                 keep_only_valid_steps: bool = False,
+                                                 is_model_output: bool = False,
+                                                 device: tp.Union[torch.device, str] = 'cpu'):
+        """Builds scatter indexes required to retrieve the original multi-codebook sequence
+        from interleaving pattern.
+
+        Args:
+            sequence_steps (int): Sequence steps.
+            n_q (int): Number of codebooks.
+            keep_only_valid_steps (bool): Build a sequence from the pattern up to valid (= fully defined) steps.
+                Steps that are beyond valid steps will be replaced by the special_token in that case.
+            is_model_output (bool): Whether to keep the sequence item corresponding to initial special token or not.
+            device (Union[torch.device, str]): Device for created tensors.
+        Returns:
+            torch.Tensor: Indexes for reconstructing the output, of shape [K, T].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, T].
+        """
+        ref_layout = self.valid_layout if keep_only_valid_steps else self.layout
+        # TODO(jade): Do we want to further truncate to only valid timesteps here as well?
+        timesteps = self.timesteps
+        assert n_q == self.n_q, f"invalid number of codebooks for the sequence and the pattern: {n_q} != {self.n_q}"
+        assert sequence_steps <= len(ref_layout), \
+            f"sequence to revert is longer than the defined pattern: {sequence_steps} > {len(ref_layout)}"
+
+        # ensure we take the appropriate indexes to keep the model output from the first special token as well
+        if is_model_output:
+            ref_layout = ref_layout[1:]
+
+        # single item indexing being super slow with pytorch vs. numpy, so we use numpy here
+        indexes = torch.zeros(n_q, timesteps, dtype=torch.long).numpy()
+        mask = torch.zeros(n_q, timesteps, dtype=torch.bool).numpy()
+        # fill indexes with last sequence step value that will correspond to our special token
+        indexes[:] = n_q * sequence_steps
+        for s, sequence_codes in enumerate(ref_layout):
+            if s < sequence_steps:
+                for code in sequence_codes:
+                    if code.t < timesteps:
+                        indexes[code.q, code.t] = s + code.q * sequence_steps
+                        mask[code.q, code.t] = 1
+        indexes = torch.from_numpy(indexes).to(device)
+        mask = torch.from_numpy(mask).to(device)
+        return indexes, mask
+
+    def revert_pattern_sequence(self, s: torch.Tensor, special_token: int, keep_only_valid_steps: bool = False):
+        """Revert a sequence built from the pattern back to the original multi-codebook sequence without interleaving.
+        The sequence is reverted using up to timesteps if specified, and non-pattern coordinates
+        are filled with the special token.
+
+        Args:
+            s (torch.Tensor): Interleaved sequence tensor obtained from the pattern, of shape [B, K, S].
+            special_token (int or float): Special token used to fill non-pattern coordinates in the new sequence.
+        Returns:
+            values (torch.Tensor): Interleaved sequence matching the pattern, of shape [B, K, T] with T
+                corresponding either to the timesteps if provided, or the total timesteps in pattern otherwise.
+            indexes (torch.Tensor): Indexes corresponding to the interleaved sequence, of shape [K, T].
+            mask (torch.Tensor): Mask corresponding to indexes that matches valid indexes of shape [K, T].
+        """
+        B, K, S = s.shape
+        indexes, mask = self._build_reverted_sequence_scatter_indexes(
+            S, K, keep_only_valid_steps, is_model_output=False, device=str(s.device)
+        )
+        s = s.view(B, -1)
+        # we append the special token as the last index of our flattened z tensor
+        s = torch.cat([s, torch.zeros_like(s[:, :1]) + special_token], dim=1)
+        values = s[:, indexes.view(-1)]
+        values = values.view(B, K, indexes.shape[-1])
+        return values, indexes, mask
+
+    def revert_pattern_logits(self, logits: torch.Tensor, special_token: float, keep_only_valid_steps: bool = False):
+        """Revert model logits obtained on a sequence built from the pattern
+        back to a tensor matching the original sequence.
+
+        This method is similar to ``revert_pattern_sequence`` with the following specificities:
+        1. It is designed to work with the extra cardinality dimension
+        2. We return the logits for the first sequence item that matches the special_token and
+        which matching target in the original sequence is the first item of the sequence,
+        while we skip the last logits as there is no matching target
+        """
+        B, card, K, S = logits.shape
+        indexes, mask = self._build_reverted_sequence_scatter_indexes(
+            S, K, keep_only_valid_steps, is_model_output=True, device=logits.device
+        )
+        logits = logits.reshape(B, card, -1)
+        # we append the special token as the last index of our flattened z tensor
+        logits = torch.cat([logits, torch.zeros_like(logits[:, :, :1]) + special_token], dim=-1)  # [B, card, K x S]
+        values = logits[:, :, indexes.view(-1)]
+        values = values.view(B, card, K, indexes.shape[-1])
+        return values, indexes, mask
+
+
+class CodebooksPatternProvider(ABC):
+    """Abstraction around providing pattern for interleaving codebooks.
+
+    The CodebooksPatternProvider abstraction allows to implement various strategies to
+    define interleaving pattern of sequences composed of multiple codebooks. For a given
+    number of codebooks `n_q`, the pattern provider can generate a specified pattern
+    corresponding to a sequence of `T` timesteps with `n_q` parallel codebooks. This pattern
+    can be used to construct a new sequence from the original codes respecting the specified
+    pattern. The pattern is defined as a list of list of code coordinates, code coordinate
+    being a tuple with the original timestep and codebook to build the new sequence.
+    Note that all patterns must start with an empty list that is then used to insert a first
+    sequence step of special tokens in the newly generated sequence.
+
+    Args:
+        n_q (int): number of codebooks.
+        cached (bool): if True, patterns for a given length are cached. In general
+            that should be true for efficiency reason to avoid synchronization points.
+    """
+    def __init__(self, n_q: int, cached: bool = True):
+        assert n_q > 0
+        self.n_q = n_q
+        self.get_pattern = lru_cache(100)(self.get_pattern)  # type: ignore
+
+    @abstractmethod
+    def get_pattern(self, timesteps: int) -> Pattern:
+        """Builds pattern with specific interleaving between codebooks.
+
+        Args:
+            timesteps (int): Total numer of timesteps.
+        """
+        raise NotImplementedError()
+
+
+class DelayedPatternProvider(CodebooksPatternProvider):
+    """Provider for delayed pattern across delayed codebooks.
+    Codebooks are delayed in the sequence and sequence steps will contain codebooks
+    from different timesteps.
+
+    Example:
+        Taking timesteps=4 and n_q=3, delays=None, the multi-codebook sequence:
+        [[1, 2, 3, 4],
+        [1, 2, 3, 4],
+        [1, 2, 3, 4]]
+        The resulting sequence obtained from the returned pattern is:
+        [[S, 1, 2, 3, 4],
+        [S, S, 1, 2, 3],
+        [S, S, S, 1, 2]]
+        (with S being a special token)
+
+    Args:
+        n_q (int): Number of codebooks.
+        delays (Optional[List[int]]): Delay for each of the codebooks.
+            If delays not defined, each codebook is delayed by 1 compared to the previous one.
+        flatten_first (int): Flatten the first N timesteps.
+        empty_initial (int): Prepend with N empty list of coordinates.
+    """
+    def __init__(self, n_q: int, delays: tp.Optional[tp.List[int]] = None,
+                 flatten_first: int = 0, empty_initial: int = 0):
+        super().__init__(n_q)
+        if delays is None:
+            delays = list(range(n_q))
+        self.delays = delays
+        self.flatten_first = flatten_first
+        self.empty_initial = empty_initial
+        assert len(self.delays) == self.n_q
+        assert sorted(self.delays) == self.delays
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        out: PatternLayout = [[]]
+        max_delay = max(self.delays)
+        if self.empty_initial:
+            out += [[] for _ in range(self.empty_initial)]
+        if self.flatten_first:
+            for t in range(min(timesteps, self.flatten_first)):
+                for q in range(self.n_q):
+                    out.append([LayoutCoord(t, q)])
+        for t in range(self.flatten_first, timesteps + max_delay):
+            v = []
+            for q, delay in enumerate(self.delays):
+                t_for_q = t - delay
+                if t_for_q >= self.flatten_first:
+                    v.append(LayoutCoord(t_for_q, q))
+            out.append(v)
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class ParallelPatternProvider(DelayedPatternProvider):
+    """Provider for parallel pattern across codebooks.
+    This pattern provider is a special case of the delayed pattern with actually no delay,
+    hence delays=repeat(0, n_q).
+
+    Args:
+        n_q (int): Number of codebooks.
+    """
+    def __init__(self, n_q: int):
+        super().__init__(n_q, [0] * n_q)
+
+
+class UnrolledPatternProvider(CodebooksPatternProvider):
+    """Provider for unrolling codebooks pattern.
+    This pattern provider enables to represent the codebook flattened completely or only to some extend
+    while also specifying a given delay between the flattened codebooks representation, allowing to
+    unroll the codebooks in the sequence.
+
+    Example:
+        1. Flattening of the codebooks.
+        By default, the pattern provider will fully flatten the codebooks such as flattening=range(n_q),
+        taking n_q = 3 and timesteps = 4:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, S, 1, S, S, 2, S, S, 3, S, S, 4],
+         [S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [1, S, S, 2, S, S, 3, S, S, 4, S, S]]
+        2. Partial flattening of the codebooks. The ``flattening`` parameter allows to specify the inner step
+        for each of the codebook, allowing to define which codebook to flatten (or keep in parallel), for example
+        taking n_q = 3, timesteps = 4 and flattening = [0, 1, 1]:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [S, 1, S, S, 2, S, S, 3, S, S, 4, S],
+         [1, S, S, 2, S, S, 3, S, S, 4, S, S]]
+        3. Flattening with delay. The ``delay`` parameter allows to further unroll the sequence of codebooks
+        allowing to specify the delay per codebook. Note that the delay between codebooks flattened to the
+        same inner timestep should be coherent. For example, taking n_q = 3, timesteps = 4, flattening = [0, 1, 1]
+        and delays = [0, 3, 3]:
+        [[1, 2, 3, 4],
+         [1, 2, 3, 4],
+         [1, 2, 3, 4]]
+        will result into:
+        [[S, S, S, 1, S, 2, S, 3, S, 4],
+         [S, S, S, 1, S, 2, S, 3, S, 4],
+         [1, 2, 3, S, 4, S, 5, S, 6, S]]
+
+    Args:
+        n_q (int): Number of codebooks.
+        flattening (Optional[List[int]]): Flattening schema over the codebooks. If not defined,
+            the codebooks will be flattened to 1 codebook per step, meaning that the sequence will
+            have n_q extra steps for each timestep.
+        delays (Optional[List[int]]): Delay for each of the codebooks. If not defined,
+            no delay is added and therefore will default to [0] * ``n_q``.
+            Note that two codebooks that will be flattened to the same inner step
+            should have the same delay, otherwise the pattern is considered as invalid.
+    """
+    FlattenedCodebook = namedtuple('FlattenedCodebook', ['codebooks', 'delay'])
+
+    def __init__(self, n_q: int, flattening: tp.Optional[tp.List[int]] = None,
+                 delays: tp.Optional[tp.List[int]] = None):
+        super().__init__(n_q)
+        if flattening is None:
+            flattening = list(range(n_q))
+        if delays is None:
+            delays = [0] * n_q
+        assert len(flattening) == n_q
+        assert len(delays) == n_q
+        assert sorted(flattening) == flattening
+        assert sorted(delays) == delays
+        self._flattened_codebooks = self._build_flattened_codebooks(delays, flattening)
+        self.max_delay = max(delays)
+
+    def _build_flattened_codebooks(self, delays: tp.List[int], flattening: tp.List[int]):
+        """Build a flattened codebooks representation as a dictionary of inner step
+        and the actual codebook indices corresponding to the flattened codebook. For convenience, we
+        also store the delay associated to the flattened codebook to avoid maintaining an extra mapping.
+        """
+        flattened_codebooks: dict = {}
+        for q, (inner_step, delay) in enumerate(zip(flattening, delays)):
+            if inner_step not in flattened_codebooks:
+                flat_codebook = UnrolledPatternProvider.FlattenedCodebook(codebooks=[q], delay=delay)
+            else:
+                flat_codebook = flattened_codebooks[inner_step]
+                assert flat_codebook.delay == delay, (
+                    "Delay and flattening between codebooks is inconsistent: ",
+                    "two codebooks flattened to the same position should have the same delay."
+                )
+                flat_codebook.codebooks.append(q)
+            flattened_codebooks[inner_step] = flat_codebook
+        return flattened_codebooks
+
+    @property
+    def _num_inner_steps(self):
+        """Number of inner steps to unroll between timesteps in order to flatten the codebooks.
+        """
+        return max([inner_step for inner_step in self._flattened_codebooks.keys()]) + 1
+
+    def num_virtual_steps(self, timesteps: int) -> int:
+        return timesteps * self._num_inner_steps + 1
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        """Builds pattern for delay across codebooks.
+
+        Args:
+            timesteps (int): Total numer of timesteps.
+        """
+        # the PatternLayout is built as a tuple of sequence position and list of coordinates
+        # so that it can be reordered properly given the required delay between codebooks of given timesteps
+        indexed_out: list = [(-1, [])]
+        max_timesteps = timesteps + self.max_delay
+        for t in range(max_timesteps):
+            # for each timestep, we unroll the flattened codebooks,
+            # emitting the sequence step with the corresponding delay
+            for step in range(self._num_inner_steps):
+                if step in self._flattened_codebooks:
+                    # we have codebooks at this virtual step to emit
+                    step_codebooks = self._flattened_codebooks[step]
+                    t_for_q = t + step_codebooks.delay
+                    coords = [LayoutCoord(t, q) for q in step_codebooks.codebooks]
+                    if t_for_q < max_timesteps and t < max_timesteps:
+                        indexed_out.append((t_for_q, coords))
+                else:
+                    # there is no codebook in this virtual step so we emit an empty list
+                    indexed_out.append((t, []))
+        out = [coords for _, coords in sorted(indexed_out)]
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class VALLEPattern(CodebooksPatternProvider):
+    """Almost VALL-E style pattern. We futher allow some delays for the
+    codebooks other than the first one.
+
+    Args:
+        n_q (int): Number of codebooks.
+        delays (Optional[List[int]]): Delay for each of the codebooks.
+            If delays not defined, each codebook is delayed by 1 compared to the previous one.
+    """
+    def __init__(self, n_q: int, delays: tp.Optional[tp.List[int]] = None):
+        super().__init__(n_q)
+        if delays is None:
+            delays = [0] * (n_q - 1)
+        self.delays = delays
+        assert len(self.delays) == self.n_q - 1
+        assert sorted(self.delays) == self.delays
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        out: PatternLayout = [[]]
+        for t in range(timesteps):
+            out.append([LayoutCoord(t, 0)])
+        max_delay = max(self.delays)
+        for t in range(timesteps + max_delay):
+            v = []
+            for q, delay in enumerate(self.delays):
+                t_for_q = t - delay
+                if t_for_q >= 0:
+                    v.append(LayoutCoord(t_for_q, q + 1))
+            out.append(v)
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)
+
+
+class MusicLMPattern(CodebooksPatternProvider):
+    """Almost MusicLM style pattern. This is equivalent to full flattening
+    but in a different order.
+
+    Args:
+        n_q (int): Number of codebooks.
+        group_by (int): Number of codebooks to group together.
+    """
+    def __init__(self, n_q: int, group_by: int = 2):
+        super().__init__(n_q)
+        self.group_by = group_by
+
+    def get_pattern(self, timesteps: int) -> Pattern:
+        out: PatternLayout = [[]]
+        for offset in range(0, self.n_q, self.group_by):
+            for t in range(timesteps):
+                for q in range(offset, offset + self.group_by):
+                    out.append([LayoutCoord(t, q)])
+        return Pattern(out, n_q=self.n_q, timesteps=timesteps)

models/modules/activation.py

@@ -0,0 +1,653 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/activation.py, modified by Puyuan Peng, 2024
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+from torch.nn import Linear, Module
+from torch.nn import functional as F
+from torch.nn.init import constant_, xavier_normal_, xavier_uniform_
+from torch.nn.modules.linear import NonDynamicallyQuantizableLinear
+from torch.nn.parameter import Parameter
+import logging
+from typing import Callable, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from torch.types import _dtype as DType
+else:
+    # The JIT doesn't understand Union, nor torch.dtype here
+    DType = int
+
+def _canonical_mask(
+        mask: Optional[Tensor],
+        mask_name: str,
+        other_type: Optional[DType],
+        other_name: str,
+        target_type: DType,
+        check_other: bool = True,
+) -> Optional[Tensor]:
+
+    if mask is not None:
+        _mask_dtype = mask.dtype
+        _mask_is_float = torch.is_floating_point(mask)
+        if _mask_dtype != torch.bool and not _mask_is_float:
+            raise AssertionError(
+                f"only bool and floating types of {mask_name} are supported")
+        if check_other and other_type is not None:
+            if _mask_dtype != other_type:
+                warnings.warn(
+                    f"Support for mismatched {mask_name} and {other_name} "
+                    "is deprecated. Use same type for both instead."
+                )
+        if not _mask_is_float:
+            mask = (
+                torch.zeros_like(mask, dtype=target_type)
+                .masked_fill_(mask, float("-inf"))
+            )
+    return mask
+
+def _in_projection_packed(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    w: Tensor,
+    b: Optional[Tensor] = None,
+) -> List[Tensor]:
+    r"""
+    Performs the in-projection step of the attention operation, using packed weights.
+    Output is a triple containing projection tensors for query, key and value.
+
+    Args:
+        q, k, v: query, key and value tensors to be projected. For self-attention,
+            these are typically the same tensor; for encoder-decoder attention,
+            k and v are typically the same tensor. (We take advantage of these
+            identities for performance if they are present.) Regardless, q, k and v
+            must share a common embedding dimension; otherwise their shapes may vary.
+        w: projection weights for q, k and v, packed into a single tensor. Weights
+            are packed along dimension 0, in q, k, v order.
+        b: optional projection biases for q, k and v, packed into a single tensor
+            in q, k, v order.
+
+    Shape:
+        Inputs:
+        - q: :math:`(..., E)` where E is the embedding dimension
+        - k: :math:`(..., E)` where E is the embedding dimension
+        - v: :math:`(..., E)` where E is the embedding dimension
+        - w: :math:`(E * 3, E)` where E is the embedding dimension
+        - b: :math:`E * 3` where E is the embedding dimension
+
+        Output:
+        - in output list :math:`[q', k', v']`, each output tensor will have the
+            same shape as the corresponding input tensor.
+    """
+    E = q.size(-1)
+    if k is v:
+        if q is k:
+            # self-attention
+            proj = F.linear(q, w, b)
+            # reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()
+            proj = proj.unflatten(-1, (3, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+            return proj[0], proj[1], proj[2]
+        else:
+            # encoder-decoder attention
+            w_q, w_kv = w.split([E, E * 2])
+            if b is None:
+                b_q = b_kv = None
+            else:
+                b_q, b_kv = b.split([E, E * 2])
+            q_proj = F.linear(q, w_q, b_q)
+            kv_proj = F.linear(k, w_kv, b_kv)
+            # reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()
+            kv_proj = kv_proj.unflatten(-1, (2, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+            return (q_proj, kv_proj[0], kv_proj[1])
+    else:
+        w_q, w_k, w_v = w.chunk(3)
+        if b is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = b.chunk(3)
+        return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)
+
+def _none_or_dtype(input: Optional[Tensor]) -> Optional[DType]:
+    if input is None:
+        return None
+    elif isinstance(input, torch.Tensor):
+        return input.dtype
+    raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor")
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces as described in the paper:
+    `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.
+
+    Multi-Head Attention is defined as:
+
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+
+    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
+
+    ``forward()`` will use a special optimized implementation if all of the following
+    conditions are met:
+
+    - self attention is being computed (i.e., ``query``, ``key``, and ``value`` are the same tensor. This
+      restriction will be loosened in the future.)
+    - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor argument ``requires_grad``
+    - training is disabled (using ``.eval()``)
+    - dropout is 0
+    - ``add_bias_kv`` is ``False``
+    - ``add_zero_attn`` is ``False``
+    - ``batch_first`` is ``True`` and the input is batched
+    - ``kdim`` and ``vdim`` are equal to ``embed_dim``
+    - at most one of ``key_padding_mask`` or ``attn_mask`` is passed
+    - if a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ is passed, neither ``key_padding_mask``
+      nor ``attn_mask`` is passed
+
+    If the optimized implementation is in use, a
+    `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be passed for
+    ``query``/``key``/``value`` to represent padding more efficiently than using a
+    padding mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_
+    will be returned, and an additional speedup proportional to the fraction of the input
+    that is padding can be expected.
+
+    Args:
+        embed_dim: Total dimension of the model.
+        num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split
+            across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``).
+        dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout).
+        bias: If specified, adds bias to input / output projection layers. Default: ``True``.
+        add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``.
+        add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1.
+            Default: ``False``.
+        kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).
+        vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+
+    Examples::
+
+        >>> # xdoctest: +SKIP
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+
+    """
+    __constants__ = ["batch_first"]
+    bias_k: Optional[torch.Tensor]
+    bias_v: Optional[torch.Tensor]
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_attn=False,
+        kdim=None,
+        vdim=None,
+        batch_first=False,
+        linear1_cls=Linear,
+        linear2_cls=Linear,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(MultiheadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = (
+            self.kdim == embed_dim and self.vdim == embed_dim
+        )
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.batch_first = batch_first
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+
+        if add_bias_kv:
+            self.bias_k = Parameter(
+                torch.empty((1, 1, embed_dim), **factory_kwargs)
+            )
+            self.bias_v = Parameter(
+                torch.empty((1, 1, embed_dim), **factory_kwargs)
+            )
+        else:
+            self.bias_k = self.bias_v = None
+
+        if linear1_cls == Linear:
+            if not self._qkv_same_embed_dim:
+                self.q_proj_weight = Parameter(
+                    torch.empty((embed_dim, embed_dim), **factory_kwargs)
+                )
+                self.k_proj_weight = Parameter(
+                    torch.empty((embed_dim, self.kdim), **factory_kwargs)
+                )
+                self.v_proj_weight = Parameter(
+                    torch.empty((embed_dim, self.vdim), **factory_kwargs)
+                )
+                self.register_parameter("in_proj_weight", None)
+            else:
+                # go down this route with voicecraft
+                self.in_proj_weight = Parameter(
+                    torch.empty((3 * embed_dim, embed_dim), **factory_kwargs)
+                )
+                self.register_parameter("q_proj_weight", None)
+                self.register_parameter("k_proj_weight", None)
+                self.register_parameter("v_proj_weight", None)
+
+            if bias: # True by default
+                self.in_proj_bias = Parameter(
+                    torch.empty(3 * embed_dim, **factory_kwargs)
+                )
+            else:
+                self.register_parameter("in_proj_bias", None)
+            self.out_proj = NonDynamicallyQuantizableLinear(
+                embed_dim, embed_dim, bias=bias, **factory_kwargs
+            )
+
+            self._reset_parameters()
+        else:
+            if not self._qkv_same_embed_dim:
+                raise NotImplementedError
+            else:
+                self.in_proj_linear = linear1_cls(
+                    embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs
+                )
+                self.in_proj_weight = self.in_proj_linear.weight
+
+                self.register_parameter("q_proj_weight", None)
+                self.register_parameter("k_proj_weight", None)
+                self.register_parameter("v_proj_weight", None)
+
+                if bias:
+                    self.in_proj_bias = self.in_proj_linear.bias
+                else:
+                    self.register_parameter("in_proj_bias", None)
+
+            self.out_proj = linear2_cls(
+                embed_dim, embed_dim, bias=bias, **factory_kwargs
+            )
+
+            if self.bias_k is not None:
+                xavier_normal_(self.bias_k)
+            if self.bias_v is not None:
+                xavier_normal_(self.bias_v)
+
+        self.add_zero_attn = add_zero_attn
+
+    def _reset_parameters(self):
+        if self._qkv_same_embed_dim:
+            xavier_uniform_(self.in_proj_weight)
+        else:
+            xavier_uniform_(self.q_proj_weight)
+            xavier_uniform_(self.k_proj_weight)
+            xavier_uniform_(self.v_proj_weight)
+
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.0)
+            constant_(self.out_proj.bias, 0.0)
+
+        if self.bias_k is not None:
+            xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            xavier_normal_(self.bias_v)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if "_qkv_same_embed_dim" not in state:
+            state["_qkv_same_embed_dim"] = True
+
+        super(MultiheadAttention, self).__setstate__(state)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        key_padding_mask: Optional[Tensor] = None,
+        need_weights: bool = True,
+        attn_mask: Optional[Tensor] = None,
+        average_attn_weights: bool = True,
+        past: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        r"""
+        Args:
+            query: Query embeddings of shape :math:`(L, E_q)` for unbatched input, :math:`(L, N, E_q)` when ``batch_first=False``
+                or :math:`(N, L, E_q)` when ``batch_first=True``, where :math:`L` is the target sequence length,
+                :math:`N` is the batch size, and :math:`E_q` is the query embedding dimension ``embed_dim``.
+                Queries are compared against key-value pairs to produce the output.
+                See "Attention Is All You Need" for more details.
+            key: Key embeddings of shape :math:`(S, E_k)` for unbatched input, :math:`(S, N, E_k)` when ``batch_first=False``
+                or :math:`(N, S, E_k)` when ``batch_first=True``, where :math:`S` is the source sequence length,
+                :math:`N` is the batch size, and :math:`E_k` is the key embedding dimension ``kdim``.
+                See "Attention Is All You Need" for more details.
+            value: Value embeddings of shape :math:`(S, E_v)` for unbatched input, :math:`(S, N, E_v)` when
+                ``batch_first=False`` or :math:`(N, S, E_v)` when ``batch_first=True``, where :math:`S` is the source
+                sequence length, :math:`N` is the batch size, and :math:`E_v` is the value embedding dimension ``vdim``.
+                See "Attention Is All You Need" for more details.
+            key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key``
+                to ignore for the purpose of attention (i.e. treat as "padding"). For unbatched `query`, shape should be :math:`(S)`.
+                Binary and byte masks are supported.
+                For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for
+                the purpose of attention. For a float mask, it will be directly added to the corresponding ``key`` value.
+            need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``.
+                Default: ``True``.
+            attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape
+                :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size,
+                :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be
+                broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.
+                Binary, byte, and float masks are supported. For a binary mask, a ``True`` value indicates that the
+                corresponding position is not allowed to attend. For a byte mask, a non-zero value indicates that the
+                corresponding position is not allowed to attend. For a float mask, the mask values will be added to
+                the attention weight.
+            average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
+                heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
+                effect when ``need_weights=True``. Default: ``True`` (i.e. average weights across heads)
+
+        Outputs:
+            - **attn_output** - Attention outputs of shape :math:`(L, E)` when input is unbatched,
+              :math:`(L, N, E)` when ``batch_first=False`` or :math:`(N, L, E)` when ``batch_first=True``,
+              where :math:`L` is the target sequence length, :math:`N` is the batch size, and :math:`E` is the
+              embedding dimension ``embed_dim``.
+            - **attn_output_weights** - Only returned when ``need_weights=True``. If ``average_attn_weights=True``,
+              returns attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
+              :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
+              :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
+              head of shape :math:`(\text{num\_heads}, L, S)` when input is unbatched or :math:`(N, \text{num\_heads}, L, S)`.
+
+            .. note::
+                `batch_first` argument is ignored for unbatched inputs.
+        """
+        is_batched = query.dim() == 3
+        if key_padding_mask is not None:
+            _kpm_dtype = key_padding_mask.dtype
+            if _kpm_dtype != torch.bool and not torch.is_floating_point(
+                key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+        why_not_fast_path = ""
+        if not is_batched:
+            why_not_fast_path = f"input not batched; expected query.dim() of 3 but got {query.dim()}"
+        elif query is not key or key is not value:
+            # When lifting this restriction, don't forget to either
+            # enforce that the dtypes all match or test cases where
+            # they don't!
+            why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
+        elif (
+            self.in_proj_bias is not None
+            and query.dtype != self.in_proj_bias.dtype
+        ):
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
+        elif (
+            self.in_proj_weight is not None
+            and query.dtype != self.in_proj_weight.dtype
+        ):
+            # this case will fail anyway, but at least they'll get a useful error message.
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
+        elif self.training:
+            why_not_fast_path = "training is enabled"
+        elif not self.batch_first:
+            why_not_fast_path = "batch_first was not True"
+        elif self.bias_k is not None:
+            why_not_fast_path = "self.bias_k was not None"
+        elif self.bias_v is not None:
+            why_not_fast_path = "self.bias_v was not None"
+        elif self.dropout:
+            why_not_fast_path = f"dropout was {self.dropout}, required zero"
+        elif self.add_zero_attn:
+            why_not_fast_path = "add_zero_attn was enabled"
+        elif not self._qkv_same_embed_dim:
+            why_not_fast_path = "_qkv_same_embed_dim was not True"
+        elif attn_mask is not None:
+            why_not_fast_path = "attn_mask was not None"
+        elif query.is_nested and key_padding_mask is not None:
+            why_not_fast_path = (
+                "key_padding_mask is not supported with NestedTensor input"
+            )
+        elif self.num_heads % 2 == 1:
+            why_not_fast_path = "num_heads is odd"
+        elif torch.is_autocast_enabled():
+            why_not_fast_path = "autocast is enabled"
+
+        if not why_not_fast_path:
+            tensor_args = (
+                query,
+                key,
+                value,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.out_proj.weight,
+                self.out_proj.bias,
+            )
+            # We have to use list comprehensions below because TorchScript does not support
+            # generator expressions.
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_fast_path = "some Tensor argument has_torch_function"
+            elif not all(
+                [
+                    (x is None or x.is_cuda or "cpu" in str(x.device))
+                    for x in tensor_args
+                ]
+            ):
+                why_not_fast_path = (
+                    "some Tensor argument is neither CUDA nor CPU"
+                )
+            elif torch.is_grad_enabled() and any(
+                [x is not None and x.requires_grad for x in tensor_args]
+            ):
+                why_not_fast_path = (
+                    "grad is enabled and at least one of query or the "
+                    "input/output projection weights or biases requires_grad"
+                )
+            if not why_not_fast_path:
+                return torch._native_multi_head_attention(
+                    query,
+                    key,
+                    value,
+                    self.embed_dim,
+                    self.num_heads,
+                    self.in_proj_weight,
+                    self.in_proj_bias,
+                    self.out_proj.weight,
+                    self.out_proj.bias,
+                    key_padding_mask
+                    if key_padding_mask is not None
+                    else attn_mask,
+                    need_weights,
+                    average_attn_weights,
+                    1
+                    if key_padding_mask is not None
+                    else 0
+                    if attn_mask is not None
+                    else None,
+                )
+
+        any_nested = query.is_nested or key.is_nested or value.is_nested
+        assert not any_nested, (
+            "MultiheadAttention does not support NestedTensor outside of its fast path. "
+            + f"The fast path was not hit because {why_not_fast_path}"
+        )
+
+        if self.batch_first and is_batched:
+            # make sure that the transpose op does not affect the "is" property
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = [x.transpose(1, 0) for x in (query, key)]
+                    value = key
+            else:
+                query, key, value = [
+                    x.transpose(1, 0) for x in (query, key, value)
+                ]
+
+        if not self._qkv_same_embed_dim:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight,
+                k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight,
+                average_attn_weights=average_attn_weights,
+            )
+        else:
+            # re-write the self.attention here, to get k, v cache
+            tgt_len, bsz, embed_dim = query.shape
+            src_len, _, _ = key.shape
+            num_heads = self.num_heads
+            key_padding_mask = _canonical_mask(
+                mask=key_padding_mask,
+                mask_name="key_padding_mask",
+                other_type=_none_or_dtype(attn_mask),
+                other_name="attn_mask",
+                target_type=query.dtype
+            )
+            attn_mask = _canonical_mask(
+                            mask=attn_mask,
+                            mask_name="attn_mask",
+                            other_type=None,
+                            other_name="",
+                            target_type=query.dtype,
+                            check_other=False,
+                            )
+            head_dim = self.embed_dim // self.num_heads
+            assert head_dim * self.num_heads == self.embed_dim, f"embed_dim {self.embed_dim} not divisible by num_heads {self.num_heads}"
+            assert key.shape == value.shape, f"key shape {key.shape} does not match value shape {value.shape}"
+            q, k, v = _in_projection_packed(query, key, value, self.in_proj_weight, self.in_proj_bias)
+            # k_present, v_present = k, v
+            
+            #
+            # reshape q, k, v for multihead attention and make em batch first
+            #
+            
+            q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+            k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+            v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1) # (bsz * num_heads, src_len, head_dim)
+            src_len = k.size(1)
+            if past is not None and past.ndim > 2:
+                expected_src_len = src_len + past[0].shape[-2]
+            else:
+                expected_src_len = src_len
+
+
+            # ensure attn_mask's dim is 3
+            if attn_mask.dim() == 2:
+                correct_2d_size = (tgt_len, expected_src_len)
+                if attn_mask.shape != correct_2d_size:
+                    raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
+                attn_mask = attn_mask.unsqueeze(0)
+            elif attn_mask.dim() == 3:
+                correct_3d_size = (bsz * num_heads, tgt_len, expected_src_len)
+                if attn_mask.shape != correct_3d_size:
+                    raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
+            else:
+                raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
+            
+            if key_padding_mask is not None:
+                assert key_padding_mask.shape == (bsz, expected_src_len), \
+                    f"expecting key_padding_mask shape of {(bsz, expected_src_len)}, but got {key_padding_mask.shape}"
+                key_padding_mask = key_padding_mask.view(bsz, 1, 1, expected_src_len).   \
+                    expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, expected_src_len)
+                if attn_mask is None:
+                    attn_mask = key_padding_mask
+                else:
+                    attn_mask = attn_mask + key_padding_mask
+            
+            if not self.training:
+                dropout_p = 0.0
+            else:
+                dropout_p = self.dropout
+
+            if need_weights:
+                raise NotImplementedError("need_weights not implemented for voicecraft")
+                # B, Nt, E = q.shape
+                # q_scaled = q / math.sqrt(E)
+
+                # assert not (is_causal and attn_mask is None), "FIXME: is_causal not implemented for need_weights"
+
+                # if attn_mask is not None:
+                #     attn_output_weights = torch.baddbmm(attn_mask, q_scaled, k.transpose(-2, -1))
+                # else:
+                #     attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
+                # attn_output_weights = softmax(attn_output_weights, dim=-1)
+                # if dropout_p > 0.0:
+                #     attn_output_weights = dropout(attn_output_weights, p=dropout_p)
+
+                # attn_output = torch.bmm(attn_output_weights, v)
+
+                # attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
+                # attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+                # attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+
+                # # optionally average attention weights over heads
+                # attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+                # if average_attn_weights:
+                #     attn_output_weights = attn_output_weights.mean(dim=1)
+
+                # if not is_batched:
+                #     # squeeze the output if input was unbatched
+                #     attn_output = attn_output.squeeze(1)
+                #     attn_output_weights = attn_output_weights.squeeze(0)
+                # return attn_output, attn_output_weights
+            else:
+                # attn_mask can be either (L,S) or (N*num_heads, L, S)
+                # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)
+                # in order to match the input for SDPA of (N, num_heads, L, S)
+                if attn_mask is not None:
+                    if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
+                        attn_mask = attn_mask.unsqueeze(0)
+                    else:
+                        attn_mask = attn_mask.view(bsz, num_heads, -1, expected_src_len)
+
+                q = q.view(bsz, num_heads, tgt_len, head_dim)
+                k = k.view(bsz, num_heads, src_len, head_dim)
+                v = v.view(bsz, num_heads, src_len, head_dim)
+                # logging.info(f"shape of past: {past.shape}")
+                if past is not None:
+                    present = torch.stack([k, v], dim=0) # (2, bsz, num_heads, src_len, head_dim)
+                    if past.ndim > 2: # this means we use kvcache, otherwise we just pass in a placeholder, but not actually using kvcache
+                        pk, pv = past
+                        k = torch.cat([pk, k], dim=-2)
+                        v = torch.cat([pv, v], dim=-2)
+                else:
+                    present = None
+                attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal=False)
+                attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
+
+                attn_output = F.linear(attn_output, self.out_proj.weight, self.out_proj.bias)
+                attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+                if not is_batched:
+                    # squeeze the output if input was unbatched
+                    attn_output = attn_output.squeeze(1)
+                # if self.training:
+                #     return attn_output, None
+                # else:
+                #     return (attn_output, present), None
+
+        # harded coded, the code do not support returning attn weigths yet
+        attn_output_weights=None
+        if self.batch_first and is_batched:
+            return attn_output.transpose(1, 0), present
+        else:
+            return attn_output, present
+

models/modules/embedding.py

@@ -0,0 +1,98 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/embedding.py
+# Copyright    2023                             (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        vocab_size: int,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.dim_model = dim_model
+
+        self.dropout = torch.nn.Dropout(p=dropout)
+        self.word_embeddings = nn.Embedding(self.vocab_size, self.dim_model)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.word_embeddings.weight
+
+    def embedding(self, index: int) -> torch.Tensor:
+        return self.word_embeddings.weight[index : index + 1]
+
+    def forward(self, x: torch.Tensor):
+        X = self.word_embeddings(x)
+        X = self.dropout(X)
+
+        return X
+
+
+class SinePositionalEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        dropout: float = 0.0,
+        scale: bool = False,
+        alpha: bool = False,
+    ):
+        super().__init__()
+        self.dim_model = dim_model
+        self.x_scale = math.sqrt(dim_model) if scale else 1.0
+        self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)
+        self.dropout = torch.nn.Dropout(p=dropout)
+
+        self.reverse = False
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, 4000))
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            if self.pe.size(1) >= x.size(1):
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        pe = torch.zeros(x.size(1), self.dim_model)
+        if self.reverse:
+            position = torch.arange(
+                x.size(1) - 1, -1, -1.0, dtype=torch.float32
+            ).unsqueeze(1)
+        else:
+            position = torch.arange(
+                0, x.size(1), dtype=torch.float32
+            ).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.dim_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.dim_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.pe = pe.to(device=x.device, dtype=x.dtype).detach()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.extend_pe(x)
+        output = x.unsqueeze(-1) if x.ndim == 2 else x
+        output = output * self.x_scale + self.alpha * self.pe[:, : x.size(1)]
+        return self.dropout(output)

models/modules/sampling.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn.functional as F
+
+def top_k_top_p_filtering(
+    logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(
+            max(top_k, min_tokens_to_keep), logits.size(-1)
+        )  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(
+            F.softmax(sorted_logits, dim=-1), dim=-1
+        )
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+            ..., :-1
+        ].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = filter_value
+    return logits
+    
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+    # temperature: (`optional`) float
+    #     The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+    # top_k: (`optional`) int
+    #     The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+    # top_p: (`optional`) float
+    #     The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+    # Temperature (higher temperature => more likely to sample low probability tokens)
+    if temperature != 1.0:
+        logits = logits / temperature
+    # Top-p/top-k filtering
+    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+    # Sample
+    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+    return token

models/modules/scaling.py

@@ -0,0 +1,1406 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/scaling.py
+# Copyright    2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import collections
+import logging
+import random
+import math
+from functools import reduce
+from itertools import repeat
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Embedding as ScaledEmbedding
+
+# from valle.utils import Transpose
+
+class Transpose(nn.Identity):
+    """(N, T, D) -> (N, D, T)"""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input.transpose(1, 2)
+    
+class ActivationBalancerFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        scale_factor: Tensor,
+        sign_factor: Optional[Tensor],
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        if sign_factor is None:
+            ctx.save_for_backward(xgt0, scale_factor)
+        else:
+            ctx.save_for_backward(xgt0, scale_factor, sign_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        if len(ctx.saved_tensors) == 3:
+            xgt0, scale_factor, sign_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+                sign_factor = sign_factor.unsqueeze(-1)
+            factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        else:
+            xgt0, scale_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+            factor = scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+def _compute_scale_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_abs: float,
+    max_abs: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    x_abs_mean = torch.mean(x.abs(), dim=sum_dims).to(torch.float32)
+
+    if min_abs == 0.0:
+        below_threshold = 0.0
+    else:
+        # below_threshold is 0 if x_abs_mean > min_abs, can be at most max_factor if
+        # x_abs)_mean , min_abs.
+        below_threshold = (
+            (min_abs - x_abs_mean) * (gain_factor / min_abs)
+        ).clamp(min=0, max=max_factor)
+
+    above_threshold = ((x_abs_mean - max_abs) * (gain_factor / max_abs)).clamp(
+        min=0, max=max_factor
+    )
+
+    return below_threshold - above_threshold
+
+
+def _compute_sign_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_positive: float,
+    max_positive: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    proportion_positive = torch.mean((x > 0).to(torch.float32), dim=sum_dims)
+    if min_positive == 0.0:
+        factor1 = 0.0
+    else:
+        # 0 if proportion_positive >= min_positive, else can be
+        # as large as max_factor.
+        factor1 = (
+            (min_positive - proportion_positive) * (gain_factor / min_positive)
+        ).clamp_(min=0, max=max_factor)
+
+    if max_positive == 1.0:
+        factor2 = 0.0
+    else:
+        # 0 if self.proportion_positive <= max_positive, else can be
+        # as large as -max_factor.
+        factor2 = (
+            (proportion_positive - max_positive)
+            * (gain_factor / (1.0 - max_positive))
+        ).clamp_(min=0, max=max_factor)
+    sign_factor = factor1 - factor2
+    # require min_positive != 0 or max_positive != 1:
+    assert not isinstance(sign_factor, float)
+    return sign_factor
+
+
+class ActivationScaleBalancerFunction(torch.autograd.Function):
+    """
+    This object is used in class ActivationBalancer when the user specified
+    min_positive=0, max_positive=1, so there are no constraints on the signs
+    of the activations and only the absolute value has a constraint.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        sign_factor: Tensor,
+        scale_factor: Tensor,
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        ctx.save_for_backward(xgt0, sign_factor, scale_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        xgt0, sign_factor, scale_factor = ctx.saved_tensors
+        for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+            sign_factor = sign_factor.unsqueeze(-1)
+            scale_factor = scale_factor.unsqueeze(-1)
+
+        factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+class RandomClampFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        min: Optional[float],
+        max: Optional[float],
+        prob: float,
+        reflect: float,
+    ) -> Tensor:
+        x_clamped = torch.clamp(x, min=min, max=max)
+        mask = torch.rand_like(x) < prob
+        ans = torch.where(mask, x_clamped, x)
+        if x.requires_grad:
+            ctx.save_for_backward(ans == x)
+            ctx.reflect = reflect
+        if reflect != 0.0:
+            ans = ans * (1.0 + reflect) - (x * reflect)
+        return ans
+
+    @staticmethod
+    def backward(
+        ctx, ans_grad: Tensor
+    ) -> Tuple[Tensor, None, None, None, None]:
+        (is_same,) = ctx.saved_tensors
+        x_grad = ans_grad * is_same.to(ans_grad.dtype)
+        reflect = ctx.reflect
+        if reflect != 0.0:
+            x_grad = x_grad * (1.0 + reflect) - (ans_grad * reflect)
+        return x_grad, None, None, None, None
+
+
+def random_clamp(
+    x: Tensor,
+    min: Optional[float] = None,
+    max: Optional[float] = None,
+    prob: float = 0.5,
+    reflect: float = 0.0,
+):
+    return RandomClampFunction.apply(x, min, max, prob, reflect)
+
+
+def random_cast_to_half(x: Tensor, min_abs: float = 5.0e-06) -> Tensor:
+    """
+    A randomized way of casting a floating point value to half precision.
+    """
+    if x.dtype == torch.float16:
+        return x
+    x_abs = x.abs()
+    is_too_small = x_abs < min_abs
+    # for elements where is_too_small is true, random_val will contain +-min_abs with
+    # probability (x.abs() / min_abs), and 0.0 otherwise.  [so this preserves expectations,
+    # for those elements].
+    random_val = min_abs * x.sign() * (torch.rand_like(x) * min_abs < x_abs)
+    return torch.where(is_too_small, random_val, x).to(torch.float16)
+
+
+class RandomGradFunction(torch.autograd.Function):
+    """
+    Does nothing in forward pass; in backward pass, gets rid of very small grads using
+    randomized approach that preserves expectations (intended to reduce roundoff).
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor, min_abs: float) -> Tensor:
+        ctx.min_abs = min_abs
+        return x
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor) -> Tuple[Tensor, None]:
+        if ans_grad.dtype == torch.float16:
+            return (
+                random_cast_to_half(
+                    ans_grad.to(torch.float32), min_abs=ctx.min_abs
+                ),
+                None,
+            )
+        else:
+            return ans_grad, None
+
+
+class RandomGrad(torch.nn.Module):
+    """
+    Gets rid of very small gradients using an expectation-preserving method, intended to increase
+    accuracy of training when using amp (automatic mixed precision)
+    """
+
+    def __init__(self, min_abs: float = 5.0e-06):
+        super(RandomGrad, self).__init__()
+        self.min_abs = min_abs
+
+    def forward(self, x: Tensor):
+        if (
+            torch.jit.is_scripting()
+            or not self.training
+            or torch.jit.is_tracing()
+        ):
+            return x
+        else:
+            return RandomGradFunction.apply(x, self.min_abs)
+
+
+class SoftmaxFunction(torch.autograd.Function):
+    """
+    Tries to handle half-precision derivatives in a randomized way that should
+    be more accurate for training than the default behavior.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor, dim: int):
+        ans = x.softmax(dim=dim)
+        # if x dtype is float16, x.softmax() returns a float32 because
+        # (presumably) that op does not support float16, and autocast
+        # is enabled.
+        if torch.is_autocast_enabled():
+            ans = ans.to(torch.float16)
+        ctx.save_for_backward(ans)
+        ctx.x_dtype = x.dtype
+        ctx.dim = dim
+        return ans
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor):
+        (ans,) = ctx.saved_tensors
+        with torch.cuda.amp.autocast(enabled=False):
+            ans_grad = ans_grad.to(torch.float32)
+            ans = ans.to(torch.float32)
+            x_grad = ans_grad * ans
+            x_grad = x_grad - ans * x_grad.sum(dim=ctx.dim, keepdim=True)
+            return x_grad, None
+
+
+def softmax(x: Tensor, dim: int):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x.softmax(dim)
+
+    return SoftmaxFunction.apply(x, dim)
+
+
+class MaxEigLimiterFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        coeffs: Tensor,
+        direction: Tensor,
+        channel_dim: int,
+        grad_scale: float,
+    ) -> Tensor:
+        ctx.channel_dim = channel_dim
+        ctx.grad_scale = grad_scale
+        ctx.save_for_backward(x.detach(), coeffs.detach(), direction.detach())
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad, *args):
+        with torch.enable_grad():
+            (x_orig, coeffs, new_direction) = ctx.saved_tensors
+            x_orig.requires_grad = True
+            num_channels = x_orig.shape[ctx.channel_dim]
+            x = x_orig.transpose(ctx.channel_dim, -1).reshape(-1, num_channels)
+            new_direction.requires_grad = False
+            x = x - x.mean(dim=0)
+            x_var = (x ** 2).mean()
+            x_residual = x - coeffs * new_direction
+            x_residual_var = (x_residual ** 2).mean()
+            # `variance_proportion` is the proportion of the variance accounted for
+            # by the top eigen-direction.  This is to be minimized.
+            variance_proportion = (x_var - x_residual_var) / (x_var + 1.0e-20)
+            variance_proportion.backward()
+        x_orig_grad = x_orig.grad
+        x_extra_grad = (
+            x_orig.grad
+            * ctx.grad_scale
+            * x_grad.norm()
+            / (x_orig_grad.norm() + 1.0e-20)
+        )
+        return x_grad + x_extra_grad.detach(), None, None, None, None
+
+
+class BasicNorm(torch.nn.Module):
+    """
+    This is intended to be a simpler, and hopefully cheaper, replacement for
+    LayerNorm.  The observation this is based on, is that Transformer-type
+    networks, especially with pre-norm, sometimes seem to set one of the
+    feature dimensions to a large constant value (e.g. 50), which "defeats"
+    the LayerNorm because the output magnitude is then not strongly dependent
+    on the other (useful) features.  Presumably the weight and bias of the
+    LayerNorm are required to allow it to do this.
+
+    So the idea is to introduce this large constant value as an explicit
+    parameter, that takes the role of the "eps" in LayerNorm, so the network
+    doesn't have to do this trick.  We make the "eps" learnable.
+
+    Args:
+       num_channels: the number of channels, e.g. 512.
+      channel_dim: the axis/dimension corresponding to the channel,
+        interprted as an offset from the input's ndim if negative.
+        shis is NOT the num_channels; it should typically be one of
+        {-2, -1, 0, 1, 2, 3}.
+       eps: the initial "epsilon" that we add as ballast in:
+             scale = ((input_vec**2).mean() + epsilon)**-0.5
+          Note: our epsilon is actually large, but we keep the name
+          to indicate the connection with conventional LayerNorm.
+       learn_eps: if true, we learn epsilon; if false, we keep it
+         at the initial value.
+    eps_min: float
+    eps_max: float
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int = -1,  # CAUTION: see documentation.
+        eps: float = 0.25,
+        learn_eps: bool = True,
+        eps_min: float = -3.0,
+        eps_max: float = 3.0,
+    ) -> None:
+        super(BasicNorm, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        if learn_eps:
+            self.eps = nn.Parameter(torch.tensor(eps).log().detach())
+        else:
+            self.register_buffer("eps", torch.tensor(eps).log().detach())
+        self.eps_min = eps_min
+        self.eps_max = eps_max
+
+    def forward(self, x: Tensor) -> Tensor:
+        assert x.shape[self.channel_dim] == self.num_channels
+        eps = self.eps
+        if self.training and random.random() < 0.25:
+            # with probability 0.25, in training mode, clamp eps between the min
+            # and max; this will encourage it to learn parameters within the
+            # allowed range by making parameters that are outside the allowed
+            # range noisy.
+
+            # gradients to allow the parameter to get back into the allowed region if it happens to exit it.
+            eps = eps.clamp(min=self.eps_min, max=self.eps_max)
+        scales = (
+            torch.mean(x ** 2, dim=self.channel_dim, keepdim=True) + eps.exp()
+        ) ** -0.5
+        return x * scales
+
+
+def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs) -> nn.Linear:
+    """
+    Behaves like a constructor of a modified version of nn.Linear
+    that gives an easy way to set the default initial parameter scale.
+
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Linear(*args, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(
+                ans.bias, -0.1 * initial_scale, 0.1 * initial_scale
+            )
+    return ans
+
+
+def ScaledConv1d(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Conv1d:
+    """
+    Behaves like a constructor of a modified version of nn.Conv1d
+    that gives an easy way to set the default initial parameter scale.
+
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(
+                ans.bias, -0.1 * initial_scale, 0.1 * initial_scale
+            )
+    return ans
+
+
+def TransposeScaledConv1d(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Sequential:
+    """
+    Transpose -> ScaledConv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        ScaledConv1d(
+            *args,
+            initial_scale=initial_scale,
+            kernel_size=kernel_size,
+            padding=padding,
+            **kwargs,
+        ),
+    )
+
+
+def ScaledConv1dTranspose(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Sequential:
+    """
+    Transpose -> ScaledConv1d
+    """
+    return nn.Sequential(
+        ScaledConv1d(
+            *args,
+            initial_scale=initial_scale,
+            kernel_size=kernel_size,
+            padding=padding,
+            **kwargs,
+        ),
+        Transpose(),
+    )
+
+
+def TransposeConv1d(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    Transpose -> Conv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+    )
+
+
+def Conv1dTranspose(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    ScaledConv1d -> Transpose
+    """
+    return nn.Sequential(
+        nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+        Transpose(),
+    )
+
+
+class SRLinear(nn.Linear):
+    """https://arxiv.org/abs/2303.06296
+    Stabilizing Transformer Training by Preventing Attention Entropy Collapse
+    """
+
+    def __init__(self, in_features, out_features, bias=True, **kwargs):
+        super().__init__(in_features, out_features, bias=bias, **kwargs)
+        self.register_buffer(
+            "u", nn.functional.normalize(torch.randn(in_features), dim=0)
+        )
+        with torch.no_grad():
+            sigma = self.get_sigma()
+        self.register_buffer("spectral_norm", sigma)
+        self.sigma = nn.Parameter(torch.ones(1))
+
+    def get_sigma(self):
+        with torch.no_grad():
+            u = self.u
+            v = self.weight.mv(u)
+            v = nn.functional.normalize(v, dim=0)
+            u = self.weight.T.mv(v)
+            u = nn.functional.normalize(u, dim=0)
+            self.u.data.copy_(u)
+        return torch.einsum("c,cd,d->", v, self.weight, u)
+
+    def get_weight(self):
+        sigma = self.get_sigma()
+        if self.training:
+            self.spectral_norm.data.copy_(sigma)
+        weight = (self.sigma / sigma) * self.weight
+        return weight
+
+    def forward(self, x):
+        return nn.functional.linear(x, self.get_weight(), self.bias)
+
+
+class SRConv1d(SRLinear):
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        kernel_size,
+        stride: int = 1,
+        padding: str = "same",
+        bias: bool = True,
+        **kwargs,
+    ):
+        in_features = in_features * kernel_size
+        super().__init__(in_features, out_features, bias=bias, **kwargs)
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+
+    def forward(self, x):
+        in_features = self.in_features // self.kernel_size
+        weight = self.get_weight().view(
+            self.out_features, in_features, self.kernel_size
+        )
+        return nn.functional.conv1d(
+            x, weight, bias=self.bias, stride=self.stride, padding=self.padding
+        )
+
+
+def TransposeSRConv1d(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    Transpose -> SRConv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        SRConv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+    )
+
+
+def SRConv1dTranspose(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    SRConv1d -> Transpose
+    """
+    return nn.Sequential(
+        SRConv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+        Transpose(),
+    )
+
+
+class ActivationBalancer(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to encourage, for
+    each channel, that it is positive at least a proportion `threshold` of the
+    time.  It does this by multiplying negative derivative values by up to
+    (1+max_factor), and positive derivative values by up to (1-max_factor),
+    interpolated from 1 at the threshold to those extremal values when none
+    of the inputs are positive.
+
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           min_positive: the minimum, per channel, of the proportion of the time
+               that (x > 0), below which we start to modify the derivatives.
+           max_positive: the maximum, per channel, of the proportion of the time
+               that (x > 0), above which we start to modify the derivatives.
+           max_factor: the maximum factor by which we modify the derivatives for
+              either the sign constraint or the magnitude constraint;
+              e.g. with max_factor=0.02, the the derivatives would be multiplied by
+              values in the range [0.98..1.02].
+           sign_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_positive and max_positive
+              are violated.
+           scale_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_abs and max_abs
+              are violated.
+           min_abs:  the minimum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+           max_abs:  the maximum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+          min_prob: determines the minimum probability with which we modify the
+             gradients for the {min,max}_positive and {min,max}_abs constraints,
+             on each forward().  This is done randomly to prevent all layers
+             from doing it at the same time.  Early in training we may use
+             higher probabilities than this; it will decay to this value.
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        min_positive: float = 0.05,
+        max_positive: float = 0.95,
+        max_factor: float = 0.04,
+        sign_gain_factor: float = 0.01,
+        scale_gain_factor: float = 0.02,
+        min_abs: float = 0.2,
+        max_abs: float = 100.0,
+        min_prob: float = 0.1,
+    ):
+        super(ActivationBalancer, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.min_positive = min_positive
+        self.max_positive = max_positive
+        self.max_factor = max_factor
+        self.min_abs = min_abs
+        self.max_abs = max_abs
+        self.min_prob = min_prob
+        self.sign_gain_factor = sign_gain_factor
+        self.scale_gain_factor = scale_gain_factor
+
+        # count measures how many times the forward() function has been called.
+        # We occasionally sync this to a tensor called `count`, that exists to
+        # make sure it is synced to disk when we load and save the model.
+        self.cpu_count = 0
+        self.register_buffer("count", torch.tensor(0, dtype=torch.int64))
+
+    def forward(self, x: Tensor) -> Tensor:
+        if (
+            torch.jit.is_scripting()
+            or not x.requires_grad
+            or torch.jit.is_tracing()
+        ):
+            return _no_op(x)
+
+        count = self.cpu_count
+        self.cpu_count += 1
+
+        if random.random() < 0.01:
+            # Occasionally sync self.cpu_count with self.count.
+            # count affects the decay of 'prob'.  don't do this on every iter,
+            # because syncing with the GPU is slow.
+            self.cpu_count = max(self.cpu_count, self.count.item())
+            self.count.fill_(self.cpu_count)
+
+        # the prob of doing some work exponentially decreases from 0.5 till it hits
+        # a floor at min_prob (==0.1, by default)
+        prob = max(self.min_prob, 0.5 ** (1 + (count / 4000.0)))
+
+        if random.random() < prob:
+            sign_gain_factor = 0.5
+            if self.min_positive != 0.0 or self.max_positive != 1.0:
+                sign_factor = _compute_sign_factor(
+                    x,
+                    self.channel_dim,
+                    self.min_positive,
+                    self.max_positive,
+                    gain_factor=self.sign_gain_factor / prob,
+                    max_factor=self.max_factor,
+                )
+            else:
+                sign_factor = None
+
+            scale_factor = _compute_scale_factor(
+                x.detach(),
+                self.channel_dim,
+                min_abs=self.min_abs,
+                max_abs=self.max_abs,
+                gain_factor=self.scale_gain_factor / prob,
+                max_factor=self.max_factor,
+            )
+            return ActivationBalancerFunction.apply(
+                x,
+                scale_factor,
+                sign_factor,
+                self.channel_dim,
+            )
+        else:
+            return _no_op(x)
+
+
+def penalize_abs_values_gt(x: Tensor, limit: float, penalty: float) -> Tensor:
+    """
+    Returns x unmodified, but in backprop will put a penalty for the excess of
+    the absolute values of elements of x over the limit "limit".  E.g. if
+    limit == 10.0, then if x has any values over 10 it will get a penalty.
+
+    Caution: the value of this penalty will be affected by grad scaling used
+    in automatic mixed precision training.  For this reasons we use this,
+    it shouldn't really matter, or may even be helpful; we just use this
+    to disallow really implausible values of scores to be given to softmax.
+    """
+    x_sign = x.sign()
+    over_limit = (x.abs() - limit) > 0
+    # The following is a memory efficient way to penalize the absolute values of
+    # x that's over the limit.  (The memory efficiency comes when you think
+    # about which items torch needs to cache for the autograd, and which ones it
+    # can throw away).  The numerical value of aux_loss as computed here will
+    # actually be larger than it should be, by limit * over_limit.sum(), but it
+    # has the same derivative as the real aux_loss which is penalty * (x.abs() -
+    # limit).relu().
+    aux_loss = penalty * ((x_sign * over_limit).to(torch.int8) * x)
+    # note: we don't do sum() here on aux)_loss, but it's as if we had done
+    # sum() due to how with_loss() works.
+    x = with_loss(x, aux_loss)
+    # you must use x for something, or this will be ineffective.
+    return x
+
+
+def _diag(x: Tensor):  # like .diag(), but works for tensors with 3 dims.
+    if x.ndim == 2:
+        return x.diag()
+    else:
+        (batch, dim, dim) = x.shape
+        x = x.reshape(batch, dim * dim)
+        x = x[:, :: dim + 1]
+        assert x.shape == (batch, dim)
+        return x
+
+
+def _whitening_metric(x: Tensor, num_groups: int):
+    """
+    Computes the "whitening metric", a value which will be 1.0 if all the eigenvalues of
+    of the centered feature covariance are the same within each group's covariance matrix
+    and also between groups.
+    Args:
+        x: a Tensor of shape (*, num_channels)
+     num_groups:  the number of groups of channels, a number >=1 that divides num_channels
+    Returns:
+        Returns a scalar Tensor that will be 1.0 if the data is "perfectly white" and
+    greater than 1.0 otherwise.
+    """
+    assert x.dtype != torch.float16
+    x = x.reshape(-1, x.shape[-1])
+    (num_frames, num_channels) = x.shape
+    assert num_channels % num_groups == 0
+    channels_per_group = num_channels // num_groups
+    x = x.reshape(num_frames, num_groups, channels_per_group).transpose(0, 1)
+    # x now has shape (num_groups, num_frames, channels_per_group)
+    # subtract the mean so we use the centered, not uncentered, covariance.
+    # My experience has been that when we "mess with the gradients" like this,
+    # it's better not do anything that tries to move the mean around, because
+    # that can easily cause instability.
+    x = x - x.mean(dim=1, keepdim=True)
+    # x_covar: (num_groups, channels_per_group, channels_per_group)
+    x_covar = torch.matmul(x.transpose(1, 2), x)
+    x_covar_mean_diag = _diag(x_covar).mean()
+    # the following expression is what we'd get if we took the matrix product
+    # of each covariance and measured the mean of its trace, i.e.
+    # the same as _diag(torch.matmul(x_covar, x_covar)).mean().
+    x_covarsq_mean_diag = (x_covar ** 2).sum() / (
+        num_groups * channels_per_group
+    )
+    # this metric will be >= 1.0; the larger it is, the less 'white' the data was.
+    metric = x_covarsq_mean_diag / (x_covar_mean_diag ** 2 + 1.0e-20)
+    return metric
+
+
+class WhiteningPenaltyFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        num_groups: int,
+        whitening_limit: float,
+        grad_scale: float,
+    ) -> Tensor:
+        ctx.save_for_backward(x)
+        ctx.num_groups = num_groups
+        ctx.whitening_limit = whitening_limit
+        ctx.grad_scale = grad_scale
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor):
+        (x_orig,) = ctx.saved_tensors
+        with torch.enable_grad():
+            with torch.cuda.amp.autocast(enabled=False):
+                x_detached = x_orig.to(torch.float32).detach()
+                x_detached.requires_grad = True
+
+                metric = _whitening_metric(x_detached, ctx.num_groups)
+
+                if random.random() < 0.005 or __name__ == "__main__":
+                    logging.info(
+                        f"Whitening: num_groups={ctx.num_groups}, num_channels={x_orig.shape[-1]}, "
+                        f"metric={metric.item():.2f} vs. limit={ctx.whitening_limit}"
+                    )
+
+                (metric - ctx.whitening_limit).relu().backward()
+                penalty_grad = x_detached.grad
+                scale = ctx.grad_scale * (
+                    x_grad.to(torch.float32).norm()
+                    / (penalty_grad.norm() + 1.0e-20)
+                )
+                penalty_grad = penalty_grad * scale
+        return x_grad + penalty_grad.to(x_grad.dtype), None, None, None
+
+
+class Whiten(nn.Module):
+    def __init__(
+        self,
+        num_groups: int,
+        whitening_limit: float,
+        prob: Union[float, Tuple[float, float]],
+        grad_scale: float,
+    ):
+        """
+        Args:
+          num_groups: the number of groups to divide the channel dim into before
+            whitening.  We will attempt to make the feature covariance
+            within each group, after mean subtraction, as "white" as possible,
+            while having the same trace across all groups.
+         whitening_limit: a value greater than 1.0, that dictates how much
+           freedom we have to violate the constraints.  1.0 would mean perfectly
+           white, with exactly the same trace across groups; larger values
+           give more freedom.  E.g. 2.0.
+         prob: the probability with which we apply the gradient modification
+           (also affects the grad scale).  May be supplied as a float,
+           or as a pair (min_prob, max_prob)
+
+          grad_scale: determines the scale on the gradient term from this object,
+            relative to the rest of the gradient on the attention weights.
+            E.g. 0.02 (you may want to use smaller values than this if prob is large)
+        """
+        super(Whiten, self).__init__()
+        assert num_groups >= 1
+        assert whitening_limit >= 1
+        assert grad_scale >= 0
+        self.num_groups = num_groups
+        self.whitening_limit = whitening_limit
+        if isinstance(prob, float):
+            assert 0 < prob <= 1
+            self.prob = prob
+        else:
+            (self.min_prob, self.max_prob) = prob
+            assert 0 < self.min_prob < self.max_prob <= 1
+            self.prob = self.max_prob
+
+        self.grad_scale = grad_scale
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        In the forward pass, this function just returns the input unmodified.
+        In the backward pass, it will modify the gradients to ensure that the
+        distribution in each group has close to (lambda times I) as the covariance
+        after mean subtraction, with the same lambda across groups.
+        For whitening_limit > 1, there will be more freedom to violate this
+        constraint.
+
+        Args:
+           x: the input of shape (*, num_channels)
+
+        Returns:
+            x, unmodified.   You should make sure
+        you use the returned value, or the graph will be freed
+        and nothing will happen in backprop.
+        """
+        if (
+            not x.requires_grad
+            or random.random() > self.prob
+            or self.grad_scale == 0
+        ):
+            return _no_op(x)
+        else:
+            if hasattr(self, "min_prob") and random.random() < 0.25:
+                # occasionally switch between min_prob and max_prob, based on whether
+                # we are above or below the threshold.
+                if (
+                    _whitening_metric(x.to(torch.float32), self.num_groups)
+                    > self.whitening_limit
+                ):
+                    # there would be a change to the grad.
+                    self.prob = self.max_prob
+                else:
+                    self.prob = self.min_prob
+
+            return WhiteningPenaltyFunction.apply(
+                x, self.num_groups, self.whitening_limit, self.grad_scale
+            )
+
+
+class WithLoss(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: Tensor, y: Tensor):
+        ctx.y_shape = y.shape
+        return x
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor):
+        return ans_grad, torch.ones(
+            ctx.y_shape, dtype=ans_grad.dtype, device=ans_grad.device
+        )
+
+
+def with_loss(x, y):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x
+    # returns x but adds y.sum() to the loss function.
+    return WithLoss.apply(x, y)
+
+
+def _no_op(x: Tensor) -> Tensor:
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x
+    else:
+        # a no-op function that will have a node in the autograd graph,
+        # to avoid certain bugs relating to backward hooks
+        return x.chunk(1, dim=-1)[0]
+
+
+class Identity(torch.nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return _no_op(x)
+
+
+class MaxEig(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to discourage
+    that any given direction in activation space accounts for more than
+    a specified proportion of the covariance (e.g. 0.2).
+
+
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           max_var_per_eig:  the maximum proportion of the variance of the
+               features/channels, after mean subtraction, that can come from
+               any given eigenvalue.
+           min_prob: the minimum probability with which we apply this during any invocation
+               of forward(), assuming last time we applied the constraint it was
+               not active; supplied for speed.
+           scale: determines the scale with which we modify the gradients, relative
+               to the existing / unmodified gradients
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        max_var_per_eig: float = 0.2,
+        min_prob: float = 0.01,
+        scale: float = 0.01,
+    ):
+        super(MaxEig, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.scale = scale
+        assert max_var_per_eig == 0.0 or max_var_per_eig > 1.0 / num_channels
+        self.max_var_per_eig = max_var_per_eig
+
+        # we figure out the dominant direction using the power method: starting with
+        # a random vector, keep multiplying by the covariance and renormalizing.
+        with torch.no_grad():
+            # arbitrary.. would use randn() but want to leave the rest of the model's
+            # random parameters unchanged for comparison
+            direction = torch.arange(num_channels).to(torch.float)
+            direction = direction / direction.norm()
+            self.register_buffer("max_eig_direction", direction)
+
+        self.min_prob = min_prob
+        # cur_prob is the current probability we'll use to apply the ActivationBalancer.
+        # We'll regress this towards prob, each tiem we try to apply it and it is not
+        # active.
+        self.cur_prob = 1.0
+
+    def forward(self, x: Tensor) -> Tensor:
+        if (
+            torch.jit.is_scripting()
+            or self.max_var_per_eig <= 0
+            or random.random() > self.cur_prob
+            or torch.jit.is_tracing()
+        ):
+            return _no_op(x)
+
+        with torch.cuda.amp.autocast(enabled=False):
+            eps = 1.0e-20
+            orig_x = x
+            x = x.to(torch.float32)
+            with torch.no_grad():
+                x = x.transpose(self.channel_dim, -1).reshape(
+                    -1, self.num_channels
+                )
+                x = x - x.mean(dim=0)
+                new_direction, coeffs = self._find_direction_coeffs(
+                    x, self.max_eig_direction
+                )
+                x_var = (x ** 2).mean()
+                x_residual = x - coeffs * new_direction
+                x_residual_var = (x_residual ** 2).mean()
+
+                # `variance_proportion` is the proportion of the variance accounted for
+                # by the top eigen-direction.
+                variance_proportion = (x_var - x_residual_var) / (
+                    x_var + 1.0e-20
+                )
+
+                # ensure new direction is nonzero even if x == 0, by including `direction`.
+                self._set_direction(
+                    0.1 * self.max_eig_direction + new_direction
+                )
+
+            if random.random() < 0.01 or __name__ == "__main__":
+                logging.info(
+                    f"variance_proportion = {variance_proportion.item()}, shape={tuple(orig_x.shape)}, cur_prob={self.cur_prob}"
+                )
+
+            if variance_proportion >= self.max_var_per_eig:
+                # The constraint is active.  Note, we should quite rarely
+                # reach here, only near the beginning of training if we are
+                # starting to diverge, should this constraint be active.
+                cur_prob = self.cur_prob
+                self.cur_prob = (
+                    1.0  # next time, do the update with probability 1.0.
+                )
+                return MaxEigLimiterFunction.apply(
+                    orig_x, coeffs, new_direction, self.channel_dim, self.scale
+                )
+            else:
+                # let self.cur_prob exponentially approach self.min_prob, as
+                # long as the constraint is inactive.
+                self.cur_prob = 0.75 * self.cur_prob + 0.25 * self.min_prob
+                return orig_x
+
+    def _set_direction(self, direction: Tensor):
+        """
+        Sets self.max_eig_direction to a normalized version of `direction`
+        """
+        direction = direction.detach()
+        direction = direction / direction.norm()
+        direction_sum = direction.sum().item()
+        if direction_sum - direction_sum == 0:  # no inf/nan
+            self.max_eig_direction[:] = direction
+        else:
+            logging.info(
+                f"Warning: sum of direction in MaxEig is {direction_sum}, "
+                "num_channels={self.num_channels}, channel_dim={self.channel_dim}"
+            )
+
+    def _find_direction_coeffs(
+        self, x: Tensor, prev_direction: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+            Figure out (an approximation to) the proportion of the variance of a set of
+            feature vectors that can be attributed to the top eigen-direction.
+            Args:
+             x: a Tensor of shape (num_frames, num_channels), with num_frames > 1.
+          prev_direction:  a Tensor of shape (num_channels,), that is our previous estimate
+                   of the top eigen-direction, or a random direction if this is the first
+                   iteration.  Does not have to be normalized, but should be nonzero.
+
+        Returns: (cur_direction, coeffs), where:
+             cur_direction: a Tensor of shape (num_channels,) that is the current
+                estimate of the top eigen-direction.
+             coeffs: a Tensor of shape (num_frames, 1) that minimizes, or
+                approximately minimizes, (x - coeffs * cur_direction).norm()
+        """
+        (num_frames, num_channels) = x.shape
+        assert num_channels > 1 and num_frames > 1
+        assert prev_direction.shape == (num_channels,)
+        # `coeffs` are the coefficients of `prev_direction` in x.
+        # actually represent the coeffs up to a constant positive factor.
+        coeffs = (x * prev_direction).sum(dim=1, keepdim=True) + 1.0e-10
+        cur_direction = (x * coeffs).sum(dim=0) / (
+            (coeffs ** 2).sum() + 1.0e-20
+        )
+        return cur_direction, coeffs
+
+
+class DoubleSwishFunction(torch.autograd.Function):
+    """
+      double_swish(x) = x * torch.sigmoid(x-1)
+    This is a definition, originally motivated by its close numerical
+    similarity to swish(swish(x)), where swish(x) =  x * sigmoid(x).
+
+    Memory-efficient derivative computation:
+     double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1)
+     double_swish'(x) = d/dx double_swish(x) =  x * s'(x) + x' * s(x) = x * s'(x) + s(x).
+     Now, s'(x) = s(x) * (1-s(x)).
+     double_swish'(x) =  x * s'(x) + s(x).
+                      =  x * s(x) * (1-s(x)) + s(x).
+                     = double_swish(x) * (1-s(x)) + s(x)
+     ... so we just need to remember s(x) but not x itself.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor) -> Tensor:
+        requires_grad = x.requires_grad
+        x_dtype = x.dtype
+        if x.dtype == torch.float16:
+            x = x.to(torch.float32)
+
+        s = torch.sigmoid(x - 1.0)
+        y = x * s
+
+        if requires_grad:
+            deriv = y * (1 - s) + s
+            # notes on derivative of x * sigmoid(x - 1):
+            # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29
+            # min \simeq -0.043638.  Take floor as -0.043637 so it's a lower bund
+            # max \simeq 1.1990.   Take ceil to be 1.2 so it's an upper bound.
+            # the combination of "+ torch.rand_like(deriv)" and casting to torch.uint8 (which
+            # floors), should be expectation-preserving.
+            floor = -0.043637
+            ceil = 1.2
+            d_scaled = (deriv - floor) * (
+                255.0 / (ceil - floor)
+            ) + torch.rand_like(deriv)
+            if __name__ == "__main__":
+                # for self-testing only.
+                assert d_scaled.min() >= 0.0
+                assert d_scaled.max() < 256.0
+            d_int = d_scaled.to(torch.uint8)
+            ctx.save_for_backward(d_int)
+        if x.dtype == torch.float16 or torch.is_autocast_enabled():
+            y = y.to(torch.float16)
+        return y
+
+    @staticmethod
+    def backward(ctx, y_grad: Tensor) -> Tensor:
+        (d,) = ctx.saved_tensors
+        # the same constants as used in forward pass.
+        floor = -0.043637
+        ceil = 1.2
+        d = d * ((ceil - floor) / 255.0) + floor
+        return y_grad * d
+
+
+class DoubleSwish(torch.nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        """Return double-swish activation function which is an approximation to Swish(Swish(x)),
+        that we approximate closely with x * sigmoid(x-1).
+        """
+        if torch.jit.is_scripting() or torch.jit.is_tracing():
+            return x * torch.sigmoid(x - 1.0)
+        return DoubleSwishFunction.apply(x)
+
+
+def BalancedDoubleSwish(
+    d_model, channel_dim=-1, max_abs=10.0, min_prob=0.25
+) -> nn.Sequential:
+    """
+    ActivationBalancer -> DoubleSwish
+    """
+    balancer = ActivationBalancer(
+        d_model, channel_dim=channel_dim, max_abs=max_abs, min_prob=min_prob
+    )
+    return nn.Sequential(
+        balancer,
+        DoubleSwish(),
+    )
+
+
+def _test_max_eig():
+    for proportion in [0.1, 0.5, 10.0]:
+        logging.info(f"proportion = {proportion}")
+        x = torch.randn(100, 128)
+        direction = torch.randn(128)
+        coeffs = torch.randn(100, 1)
+        x += proportion * direction * coeffs
+
+        x.requires_grad = True
+
+        num_channels = 128
+        m = MaxEig(
+            num_channels, 1, 0.5, scale=0.1  # channel_dim  # max_var_per_eig
+        )  # grad_scale
+
+        for _ in range(4):
+            y = m(x)
+
+        y_grad = torch.randn_like(x)
+        y.backward(gradient=y_grad)
+
+        if proportion < 0.2:
+            assert torch.allclose(x.grad, y_grad, atol=1.0e-02)
+        elif proportion > 1.0:
+            assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_whiten():
+    for proportion in [0.1, 0.5, 10.0]:
+        logging.info(f"_test_whiten(): proportion = {proportion}")
+        x = torch.randn(100, 128)
+        direction = torch.randn(128)
+        coeffs = torch.randn(100, 1)
+        x += proportion * direction * coeffs
+
+        x.requires_grad = True
+
+        num_channels = 128
+        m = Whiten(
+            1, 5.0, prob=1.0, grad_scale=0.1  # num_groups  # whitening_limit,
+        )  # grad_scale
+
+        for _ in range(4):
+            y = m(x)
+
+        y_grad = torch.randn_like(x)
+        y.backward(gradient=y_grad)
+
+        if proportion < 0.2:
+            assert torch.allclose(x.grad, y_grad)
+        elif proportion > 1.0:
+            assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_activation_balancer_sign():
+    probs = torch.arange(0, 1, 0.01)
+    N = 1000
+    x = 1.0 * (
+        (2.0 * (torch.rand(probs.numel(), N) < probs.unsqueeze(-1))) - 1.0
+    )
+    x = x.detach()
+    x.requires_grad = True
+    m = ActivationBalancer(
+        probs.numel(),
+        channel_dim=0,
+        min_positive=0.05,
+        max_positive=0.95,
+        max_factor=0.2,
+        min_abs=0.0,
+    )
+
+    y_grad = torch.sign(torch.randn(probs.numel(), N))
+
+    y = m(x)
+    y.backward(gradient=y_grad)
+    print("_test_activation_balancer_sign: x = ", x)
+    print("_test_activation_balancer_sign: y grad = ", y_grad)
+    print("_test_activation_balancer_sign: x grad = ", x.grad)
+
+
+def _test_activation_balancer_magnitude():
+    magnitudes = torch.arange(0, 1, 0.01)
+    N = 1000
+    x = torch.sign(torch.randn(magnitudes.numel(), N)) * magnitudes.unsqueeze(
+        -1
+    )
+    x = x.detach()
+    x.requires_grad = True
+    m = ActivationBalancer(
+        magnitudes.numel(),
+        channel_dim=0,
+        min_positive=0.0,
+        max_positive=1.0,
+        max_factor=0.2,
+        min_abs=0.2,
+        max_abs=0.8,
+        min_prob=1.0,
+    )
+
+    y_grad = torch.sign(torch.randn(magnitudes.numel(), N))
+
+    y = m(x)
+    y.backward(gradient=y_grad)
+    print("_test_activation_balancer_magnitude: x = ", x)
+    print("_test_activation_balancer_magnitude: y grad = ", y_grad)
+    print("_test_activation_balancer_magnitude: x grad = ", x.grad)
+
+
+def _test_basic_norm():
+    num_channels = 128
+    m = BasicNorm(num_channels=num_channels, channel_dim=1)
+
+    x = torch.randn(500, num_channels)
+
+    y = m(x)
+
+    assert y.shape == x.shape
+    x_rms = (x ** 2).mean().sqrt()
+    y_rms = (y ** 2).mean().sqrt()
+    print("x rms = ", x_rms)
+    print("y rms = ", y_rms)
+    assert y_rms < x_rms
+    assert y_rms > 0.5 * x_rms
+
+
+def _test_double_swish_deriv():
+    x = torch.randn(10, 12, dtype=torch.double) * 3.0
+    x.requires_grad = True
+    m = DoubleSwish()
+
+    tol = (1.2 - (-0.043637)) / 255.0
+    torch.autograd.gradcheck(m, x, atol=tol)
+
+    # for self-test.
+    x = torch.randn(1000, 1000, dtype=torch.double) * 3.0
+    x.requires_grad = True
+    y = m(x)
+
+
+def _test_softmax():
+    a = torch.randn(2, 10, dtype=torch.float64)
+    b = a.clone()
+    a.requires_grad = True
+    b.requires_grad = True
+    a.softmax(dim=1)[:, 0].sum().backward()
+    print("a grad = ", a.grad)
+    softmax(b, dim=1)[:, 0].sum().backward()
+    print("b grad = ", b.grad)
+    assert torch.allclose(a.grad, b.grad)
+
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+    _test_softmax()
+    _test_whiten()
+    _test_max_eig()
+    _test_activation_balancer_sign()
+    _test_activation_balancer_magnitude()
+    _test_basic_norm()
+    _test_double_swish_deriv()

models/modules/transformer.py

@@ -0,0 +1,698 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/transformer.py, modified by Puyuan Peng 2024
+import copy
+import numbers
+from functools import partial
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .activation import MultiheadAttention
+from .scaling import ActivationBalancer, BalancedDoubleSwish
+from .scaling import BasicNorm as _BasicNorm
+
+_shape_t = Union[int, List[int], torch.Size]
+
+
+class LayerNorm(nn.Module):
+    __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
+    normalized_shape: Tuple[int, ...]
+    eps: float
+    elementwise_affine: bool
+
+    def __init__(
+        self,
+        normalized_shape: _shape_t,
+        eps: float = 1e-5,
+        elementwise_affine: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            # mypy error: incompatible types in assignment
+            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+            self.bias = nn.Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        if self.elementwise_affine:
+            nn.init.ones_(self.weight)
+            nn.init.zeros_(self.bias)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return (
+                F.layer_norm(
+                    input,
+                    self.normalized_shape,
+                    self.weight,
+                    self.bias,
+                    self.eps,
+                ),
+                embedding,
+            )
+
+        assert embedding is None
+        return F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps
+        )
+
+    def extra_repr(self) -> str:
+        return (
+            "{normalized_shape}, eps={eps}, "
+            "elementwise_affine={elementwise_affine}".format(**self.__dict__)
+        )
+
+
+class AdaptiveLayerNorm(nn.Module):
+    r"""Adaptive Layer Normalization"""
+
+    def __init__(self, d_model, norm) -> None:
+        super(AdaptiveLayerNorm, self).__init__()
+        self.project_layer = nn.Linear(d_model, 2 * d_model)
+        self.norm = norm
+        self.d_model = d_model
+        self.eps = self.norm.eps
+
+    def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            weight, bias = torch.split(
+                self.project_layer(embedding),
+                split_size_or_sections=self.d_model,
+                dim=-1,
+            )
+            return (weight * self.norm(input) + bias, embedding)
+
+        weight, bias = torch.split(
+            self.project_layer(embedding),
+            split_size_or_sections=self.d_model,
+            dim=-1,
+        )
+        return weight * self.norm(input) + bias
+
+
+class BasicNorm(_BasicNorm):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ):
+        super(BasicNorm, self).__init__(d_model, eps=eps)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return (
+                super(BasicNorm, self).forward(input),
+                embedding,
+            )
+
+        assert embedding is None
+        return super(BasicNorm, self).forward(input)
+
+
+class BalancedBasicNorm(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ):
+        super(BalancedBasicNorm, self).__init__()
+        self.balancer = ActivationBalancer(
+            d_model,
+            channel_dim=-1,
+            min_positive=0.45,
+            max_positive=0.55,
+            max_abs=6.0,
+        )
+        self.norm = BasicNorm(d_model, eps, device=device, dtype=dtype)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return self.norm((self.balancer(input), embedding))
+
+        assert embedding is None
+        return self.norm(self.balancer(input))
+
+
+class IdentityNorm(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ) -> None:
+        super(IdentityNorm, self).__init__()
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            return input
+
+        assert embedding is None
+        return input
+
+
+class TransformerEncoderLayer(nn.Module):
+    __constants__ = ["batch_first", "norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        device=None,
+        dtype=None,
+        linear1_self_attention_cls: nn.Module = nn.Linear,
+        linear2_self_attention_cls: nn.Module = nn.Linear,
+        linear1_feedforward_cls: nn.Module = nn.Linear,
+        linear2_feedforward_cls: nn.Module = nn.Linear,
+        layer_norm_cls: nn.Module = LayerNorm,
+        layer_norm_eps: float = 1e-5,
+        adaptive_layer_norm=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+
+        # Implementation of Feedforward model
+        self.linear1 = linear1_feedforward_cls(
+            d_model, dim_feedforward, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = linear2_feedforward_cls(
+            dim_feedforward, d_model, **factory_kwargs
+        )
+
+        self.norm_first = norm_first
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            activation = BalancedDoubleSwish(d_model)
+
+        # # We can't test self.activation in forward() in TorchScript,
+        # # so stash some information about it instead.
+        # if activation is F.relu or isinstance(activation, torch.nn.ReLU):
+        #     self.activation_relu_or_gelu = 1
+        # elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
+        #     self.activation_relu_or_gelu = 2
+        # else:
+        #     self.activation_relu_or_gelu = 0
+        self.activation = activation
+
+        norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)
+        if layer_norm_cls == IdentityNorm:
+            norm2 = BalancedBasicNorm(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+        else:
+            norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+
+        if adaptive_layer_norm:
+            self.norm1 = AdaptiveLayerNorm(d_model, norm1)
+            self.norm2 = AdaptiveLayerNorm(d_model, norm2)
+        else:
+            self.norm1 = norm1
+            self.norm2 = norm2
+
+    def __setstate__(self, state):
+        super(TransformerEncoderLayer, self).__setstate__(state)
+        if not hasattr(self, "activation"):
+            self.activation = F.relu
+
+    def forward(
+        self,
+        src: Tensor,
+        src_mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        need_weights: Optional[bool] = False,
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        x, stage_embedding = src, None
+        is_src_tuple = False
+        if isinstance(src, tuple):
+            x, stage_embedding = src
+            is_src_tuple = True
+
+        if src_key_padding_mask is not None:
+            _skpm_dtype = src_key_padding_mask.dtype
+            if _skpm_dtype != torch.bool and not torch.is_floating_point(
+                src_key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+        if need_weights:
+            if self.norm_first:
+                out, attn = self._sa_block_attn(
+                    self.norm1(x, stage_embedding),
+                    src_mask,
+                    src_key_padding_mask,
+                    past
+                )
+                out, present = out # present is the kvcache of the present timestep
+                x = x + out
+                x = x + self._ff_block(self.norm2(x, stage_embedding))
+            else:
+                out, attn = self._sa_block_attn(x, src_mask, src_key_padding_mask, past)
+                out, present = out # present is the kvcache of the present timestep
+                x = self.norm1(
+                    x + out,
+                    stage_embedding,
+                )
+                x = self.norm2(x + self._ff_block(x), stage_embedding)
+            assert not is_src_tuple
+                # return (x, stage_embedding)
+            return (x, attn)
+        else:
+            if self.norm_first:
+                out = self._sa_block(
+                    self.norm1(x, stage_embedding),
+                    src_mask,
+                    src_key_padding_mask, past
+                )
+                out, present = out # present is the kvcache of the present timestep
+                x = x + out
+                x = x + self._ff_block(self.norm2(x, stage_embedding))
+            else:
+                out = self._sa_block(x, src_mask, src_key_padding_mask)
+                out, present = out # present is the kvcache of the present timestep
+                x = self.norm1(
+                    x + out,
+                    stage_embedding, past
+                )
+                x = self.norm2(x + self._ff_block(x), stage_embedding)
+            
+            if is_src_tuple:
+                x = (x, stage_embedding)
+            if present != None:
+                x = [x, present]
+            return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+            past=past
+        )
+        x, present = x
+        return self.dropout1(x), present
+
+    # self-attention block, also return attention weights
+    def _sa_block_attn(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        x, attn = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=True,
+            past=past
+        )
+        x, present = x
+        return (self.dropout1(x), present), attn
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout2(x)
+
+
+class TransformerEncoder(nn.Module):
+    r"""TransformerEncoder is a stack of N encoder layers. Users can build the
+    BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.
+
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+        enable_nested_tensor: if True, input will automatically convert to nested tensor
+            (and convert back on output). This will improve the overall performance of
+            TransformerEncoder when padding rate is high. Default: ``True`` (enabled).
+
+    Examples::
+        >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = transformer_encoder(src)
+    """
+    __constants__ = ["norm"]
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super(TransformerEncoder, self).__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(
+        self,
+        src: Tensor,
+        mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        return_layer_states: bool = False,
+        need_weights:Optional[bool] = False,
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+            return_layer_states: return layers' state (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        if return_layer_states:
+            assert not need_weights
+            layer_states = []  # layers' output
+            output = src
+            for mod in self.layers:
+                output = mod(
+                    output,
+                    src_mask=mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                    past=past
+                )
+                layer_states.append(output[0])
+
+            if self.norm is not None:
+                output = self.norm(output)
+
+            return layer_states, output
+        if need_weights:
+            assert not return_layer_states
+            layer_attn = []  # layers' output
+            output = src
+            for mod in self.layers:
+                output = mod(
+                    output,
+                    src_mask=mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                    need_weights=True,
+                    past=past
+                )
+                layer_attn.append(output[1])
+
+            if self.norm is not None:
+                output = self.norm(output)
+
+            return layer_attn, output
+        
+        output = src
+        all_present = []
+        for n_layer, mod in enumerate(self.layers):
+            output = mod(
+                output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, past=None if past is None else past[n_layer]
+            )
+            if isinstance(output, list):
+                output, present = output
+                all_present.append(present)
+
+        if self.norm is not None:
+            output = self.norm(output)
+        if all_present != []:
+            all_present = torch.stack(all_present, dim=0) # (num_layers, 2, batch_size, num_heads, seq_len, head_dim)
+            output = [output, all_present]
+        return output
+
+
+class TransformerDecoderLayer(nn.Module):
+    __constants__ = ["batch_first", "norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+        linear1_self_attention_cls: nn.Module = nn.Linear,
+        linear2_self_attention_cls: nn.Module = nn.Linear,
+        linear1_feedforward_cls: nn.Module = nn.Linear,
+        linear2_feedforward_cls: nn.Module = nn.Linear,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        device=None,
+        dtype=None,
+        layer_norm_cls: nn.Module = LayerNorm,
+        layer_norm_eps: float = 1e-5,
+        adaptive_layer_norm=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(TransformerDecoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+        self.multihead_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+        # Implementation of Feedforward model
+        self.linear1 = linear1_feedforward_cls(
+            d_model, dim_feedforward, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = linear2_feedforward_cls(
+            dim_feedforward, d_model, **factory_kwargs
+        )
+
+        self.norm_first = norm_first
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            self.activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            self.activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            self.activation = BalancedDoubleSwish(d_model)
+        else:
+            self.activation = activation
+
+        if adaptive_layer_norm:
+            norm1 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            norm3 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+
+            self.norm1 = AdaptiveLayerNorm(d_model, norm1)
+            self.norm2 = AdaptiveLayerNorm(d_model, norm2)
+            self.norm3 = AdaptiveLayerNorm(d_model, norm3)
+        else:
+            self.norm1 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            self.norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            if layer_norm_cls == IdentityNorm:
+                self.norm3 = BalancedBasicNorm(
+                    d_model, eps=layer_norm_eps, **factory_kwargs
+                )
+            else:
+                self.norm3 = layer_norm_cls(
+                    d_model, eps=layer_norm_eps, **factory_kwargs
+                )
+
+    def forward(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the inputs (and mask) through the decoder layer.
+
+        Args:
+            tgt: the sequence to the decoder layer (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        tgt_is_tuple = False
+        if isinstance(tgt, tuple):
+            x, stage_embedding = tgt
+            tgt_is_tuple = True
+        else:
+            x, stage_embedding = tgt, None
+
+        if self.norm_first:
+            x = x + self._sa_block(
+                self.norm1(x, stage_embedding), tgt_mask, tgt_key_padding_mask
+            )
+            x = x + self._mha_block(
+                self.norm2(x, stage_embedding),
+                memory,
+                memory_mask,
+                memory_key_padding_mask,
+            )
+            x = x + self._ff_block(self.norm3(x, stage_embedding))
+        else:
+            x = self.norm1(
+                x + self._sa_block(x, tgt_mask, tgt_key_padding_mask),
+                stage_embedding,
+            )
+            x = self.norm2(
+                x
+                + self._mha_block(
+                    x, memory, memory_mask, memory_key_padding_mask
+                ),
+                stage_embedding,
+            )
+            x = self.norm3(x + self._ff_block(x), stage_embedding)
+
+        if tgt_is_tuple:
+            return (x, stage_embedding)
+        return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+        )[0]
+        return self.dropout1(x)
+
+    # multihead attention block
+    def _mha_block(
+        self,
+        x: Tensor,
+        mem: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+    ) -> Tensor:
+        x = self.multihead_attn(
+            x,
+            mem,
+            mem,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+        )[0]
+        return self.dropout2(x)
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout3(x)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError(
+        "activation should be relu/gelu, not {}".format(activation)
+    )

models/modules/utils.py

@@ -0,0 +1,37 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/transformer.py, modified by Puyuan Peng
+import torch
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """
+    Args:
+      lengths:
+        A 1-D tensor containing sentence lengths.
+      max_len:
+        The length of masks.
+    Returns:
+      Return a 2-D bool tensor, where masked positions
+      are filled with `True` and non-masked positions are
+      filled with `False`.
+
+    >>> lengths = torch.tensor([1, 3, 2, 5])
+    >>> make_pad_mask(lengths)
+    tensor([[False,  True,  True,  True,  True],
+            [False, False, False,  True,  True],
+            [False, False,  True,  True,  True],
+            [False, False, False, False, False]])
+    """
+    assert lengths.ndim == 1, lengths.ndim
+    max_len = max(max_len, lengths.max())
+    n = lengths.size(0)
+    seq_range = torch.arange(0, max_len, device=lengths.device)
+    expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
+
+    return expaned_lengths >= lengths.unsqueeze(-1)
+
+def generate_partial_autoregressive_mask(sz, start, end):
+    mask = torch.zeros(sz, sz).bool()
+    mask[start:end, start:end] = torch.triu(torch.ones(end-start, end-start,dtype=torch.bool), diagonal=1)
+    mask[:start, start:end] = True
+    mask[end:, start:end] = True
+    return mask

models/voicecraft.py

@@ -0,0 +1,1406 @@
+import random
+
+import numpy as np
+import logging
+import argparse, copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchmetrics.classification import MulticlassAccuracy
+
+from .modules.utils import make_pad_mask
+
+from .modules.embedding import SinePositionalEmbedding, TokenEmbedding
+from .modules.transformer import (
+    LayerNorm,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
+from .codebooks_patterns import DelayedPatternProvider
+
+def top_k_top_p_filtering(
+    logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(
+            max(top_k, min_tokens_to_keep), logits.size(-1)
+        )  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(
+            F.softmax(sorted_logits, dim=-1), dim=-1
+        )
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+            ..., :-1
+        ].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = filter_value
+    return logits
+
+
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+    # temperature: (`optional`) float
+    #     The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+    # top_k: (`optional`) int
+    #     The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+    # top_p: (`optional`) float
+    #     The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+    # Temperature (higher temperature => more likely to sample low probability tokens)
+    if temperature != 1.0:
+        logits = logits / temperature
+    # Top-p/top-k filtering
+    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+    # Sample
+    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+    return token
+
+
+
+class VoiceCraft(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = copy.copy(args)
+        self.pattern = DelayedPatternProvider(n_q=self.args.n_codebooks)
+        if not getattr(self.args, "special_first", False):
+            self.args.special_first = 0
+        if not getattr(self.args, "n_special", False):
+            self.args.n_special = 3
+        self.args.eos = getattr(self.args, "eos", -1)
+        self.eog = nn.Parameter(torch.full((self.args.n_codebooks, 1), self.args.eog, dtype=torch.long), requires_grad=False) # [K 1]
+        if self.args.eos > 0:
+            assert self.args.eos != self.args.audio_pad_token and self.args.eos != self.args.empty_token, self.args.eos
+            self.eos = nn.Parameter(torch.full((self.args.n_codebooks, 1), self.args.eos, dtype=torch.long), requires_grad=False) # [K 1]
+        if type(self.args.audio_vocab_size) == str:
+            self.args.audio_vocab_size = eval(self.args.audio_vocab_size)
+
+        self.n_text_tokens = self.args.text_vocab_size + 1
+        assert self.args.text_pad_token == self.args.text_vocab_size, f"self.args.text_vocab_size: {self.args.text_vocab_size}, self.args.text_pad_token: {self.args.text_pad_token}"
+
+        self.n_audio_tokens = [self.args.audio_vocab_size + self.args.n_special] * self.args.n_codebooks # special tokens: empty token, EOG token, audio pad token
+        assert self.args.audio_vocab_size == self.args.empty_token, self.args.empty_token
+        assert self.args.eog == self.args.audio_vocab_size + 1, self.args.eog
+        assert self.args.audio_pad_token == self.args.audio_vocab_size + 2, self.args.audio_pad_token
+
+        self.text_embedding = TokenEmbedding(
+            dim_model=self.args.d_model,
+            vocab_size=self.n_text_tokens, 
+            dropout=self.args.text_embedding_dropout
+        )
+
+        self.audio_embedding = nn.ModuleList(
+            [
+                TokenEmbedding(
+                dim_model=self.args.audio_embedding_dim, 
+                vocab_size=self.n_audio_tokens[k], 
+                dropout=self.args.audio_embedding_dropout
+            ) for k in range(self.args.n_codebooks)
+            ]
+        )
+        self.mask_embedding = nn.Parameter(torch.randn(self.args.max_n_spans, self.args.d_model), requires_grad=True)
+        self.text_positional_embedding = SinePositionalEmbedding(
+            self.args.d_model,
+            dropout=self.args.text_positional_embedding_dropout,
+            scale=False,
+            alpha=True, # learnable scaler, scale the volume of positional embedding
+        )
+        self.audio_positional_embedding = SinePositionalEmbedding(
+            self.args.d_model,
+            dropout=self.args.audio_positional_embedding_dropout,
+            scale=False,
+            alpha=True, # learnable scaler, scale the volume of positional embedding
+        )
+
+        dec_layer = TransformerEncoderLayer(
+            self.args.d_model,
+            self.args.nhead,
+            dim_feedforward=self.args.d_model * 4,
+            dropout=self.args.trm_dropout,
+            batch_first=True,
+            norm_first=True,
+            layer_norm_cls=LayerNorm
+        )
+        self.decoder = TransformerEncoder(
+            dec_layer,
+            num_layers=self.args.num_decoder_layers,
+            norm=LayerNorm(self.args.d_model),
+        )
+        
+        self.predict_layer = nn.ModuleList(
+            [
+                nn.Sequential(nn.Linear(self.args.d_model, self.args.audio_vocab_size//2), nn.GELU(), nn.Linear(self.args.audio_vocab_size//2, self.n_audio_tokens[k])) for k in range(self.args.n_codebooks)
+            ]
+        )
+        
+        self.accuracy_metrics = nn.ModuleList(
+            [MulticlassAccuracy(
+                self.n_audio_tokens[k],
+                top_k=10,
+                average="micro",
+                multidim_average="global",
+                ignore_index=None,
+            ) for k in range(self.args.n_codebooks)]
+        )
+
+    
+    def prepare_mask_intervals(self, y_lens):
+        mask_intervals = []
+        non_mask_intervals = []
+
+        for i, y_len in enumerate(y_lens):
+            if self.args.mask_sample_dist == "uniform":
+                n_spans = random.choice(range(1, self.args.max_n_spans+1))
+            elif "poisson" in self.args.mask_sample_dist.lower():
+                param = float(self.args.mask_sample_dist[len("poisson"):])
+                poisson_sample = torch.poisson(torch.tensor([param]))
+                n_spans = int(poisson_sample.clamp(1, self.args.max_n_spans).item())
+
+            starts = random.sample(range(1, y_len-1-self.args.mask_len_min), n_spans)
+            starts = sorted(starts)
+
+            for j in range(len(starts)-1, 0, -1):
+                if starts[j] - starts[j-1] < self.args.min_gap:
+                    del starts[j] # If elements are too close, delete the later one
+            assert len(starts) > 0, f"there is no masked span left, y_len: {y_len}, sampled n_spans: {n_spans}"
+
+            temp_starts =  starts + [y_len]
+            gaps = [temp_starts[j+1] - temp_starts[j] for j in range(len(temp_starts)-1)]
+
+            ends = []
+
+            for j, (start, gap) in enumerate(zip(starts, gaps)):
+                mask_len = random.randint(self.args.mask_len_min, self.args.mask_len_max)
+                # if mask_len > gap * self.args.max_mask_portion: # make sure the masks are not overlapping with each other
+                if mask_len > gap - 1: # make sure the masks are not overlapping with each other
+                    # temp_mask_start = int(0.6*gap*self.args.max_mask_portion)
+                    # temp_mask_end = int(gap*self.args.max_mask_portion)
+                    temp_mask_start = 1
+                    temp_mask_end = gap - 1
+                    mask_len = random.randint(temp_mask_start, temp_mask_end)
+                ends.append(start + mask_len)
+            
+            mask_intervals.append([(s,e) for s,e in zip(starts, ends)])
+            non_mask_intervals.append([(ns,ne) for ns, ne in zip([0]+ends, starts+[y_len])])
+
+        return mask_intervals, non_mask_intervals
+    
+    def rearrange(self, y, non_mask_intervals, mask_intervals):
+        reduced_eog = getattr(self.args, "reduced_eog", 0)
+        rearranged_y = []
+        for i in range(len(y)):
+            if self.args.eos > 0:
+                assert reduced_eog
+                cur_y = [y[i, :, item[0]: item[1]] for item in non_mask_intervals[i][:-1]] + [torch.cat([y[i, :, non_mask_intervals[i][-1][0]: non_mask_intervals[i][-1][1]], self.eos], dim=-1)] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # only insert eog to the last non-mask-interval, which is when the utterance actual ends
+            else:
+                if reduced_eog:
+                    cur_y = [y[i, :, item[0]: item[1]] for item in non_mask_intervals[i][:-1]] + [torch.cat([y[i, :, non_mask_intervals[i][-1][0]: non_mask_intervals[i][-1][1]], self.eog], dim=-1)] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # only insert eog to the last non-mask-interval, which is when the utterance actual ends
+                else:
+                    cur_y = [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in non_mask_intervals[i]] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # eog is added to each section TODO this is not correct, I should add eog to non_mask_intervals if that segment is not the ending segment (as there is no way for the model to predict eog for those segments, and this will do harm to tts experiment, where the model randomly output eog for the first segment)
+            rearranged_y.append(cur_y)
+        return rearranged_y
+        
+    def shift(self, rearranged_y):
+        shifted_y = []
+        patterns = []
+        for i in range(len(rearranged_y)):
+            cur_patterns = [self.pattern.get_pattern(cur_y.shape[1]) for cur_y in rearranged_y[i]]
+            out = [cur_pattern.build_pattern_sequence(z=cur_y.unsqueeze(0).contiguous(), special_token=self.args.empty_token, keep_only_valid_steps=False) for cur_pattern, cur_y in zip(cur_patterns, rearranged_y[i])]
+            shifted_y.append([item[0].squeeze(0) for item in out]) # the first item is values, later two are indexes and mask
+            patterns.append(cur_patterns)
+        return shifted_y, patterns
+    
+    def insert_mask(self, shifted_y):
+        inserted_y = []
+        mask_position = []
+        mask_value = []
+        for i in range(len(shifted_y)):
+            num_masks = (len(shifted_y[i]) - 1) // 2
+            assert num_masks == (len(shifted_y[i]) - 1) / 2, len(shifted_y[i])
+            emb_inds = list(range(self.args.max_n_spans))
+            if self.args.shuffle_mask_embedding:
+                random.shuffle(emb_inds)
+            emb_inds_use = emb_inds[:num_masks]
+            emb_inds_use = emb_inds_use + emb_inds_use
+            mask_value.append(emb_inds_use)
+            cur_inserted_y = []
+            cur_mask_position = []
+            for j in range(len(shifted_y[i])-1):
+                cur_inserted_y.append(shifted_y[i][j])
+                cur_mask_position.append(sum([item.shape[1] for item in cur_inserted_y])) # each item is of shape [K S], so take shape[1]
+                cur_inserted_y.append(self.eog) # insert mask token of shape [K, 1], BUT we are actually using the eog token as a place holder here, as the real mask will be inserted in embed_y function
+
+            cur_inserted_y.append(shifted_y[i][-1])
+
+            inserted_y.append(cur_inserted_y)
+            mask_position.append(cur_mask_position)
+        return inserted_y, mask_position, mask_value
+    
+    def cat_y(self, inserted_y, mask_position, y_lens):
+        reduced_eog = getattr(self.args, "reduced_eog", 0)
+        cated_y = []
+        new_y_lens = []
+        for i in range(len(inserted_y)):
+            cur_cated_y = torch.cat(inserted_y[i], dim=1) #[K S]
+            cur_cated_y = cur_cated_y.transpose(1,0) # [S K]
+            cur_cated_y_len = cur_cated_y.shape[0]
+            if reduced_eog:
+                assert cur_cated_y_len == y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i])/2 + 1), f"cur_cated_y_len == {cur_cated_y_len}, but it should be y_lens[i] ({y_lens[i]}) + len(mask_position[i]) ({len(mask_position[i])}) + (len(mask_position[i]) + 1) * self.args.n_codebooks ({(len(mask_position[i]) + 1) * self.args.n_codebooks}) + (len(mask_position[i])/2 + 1) ({len(mask_position[i])/2 + 1})={y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i])/2 + 1)}"
+            else:
+                assert cur_cated_y_len == y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i]) + 1), f"cur_cated_y_len == {cur_cated_y_len}, but it should be y_lens[i] ({y_lens[i]}) + len(mask_position[i]) ({len(mask_position[i])}) + (len(mask_position[i]) + 1) * self.args.n_codebooks ({(len(mask_position[i]) + 1) * self.args.n_codebooks}) + (len(mask_position[i]) + 1) ({len(mask_position[i]) + 1})" # the last term represent the inserted eog token, originally it's inserted at the end of every token, but this is wrong
+            new_y_lens.append(cur_cated_y_len)
+            cated_y.append(cur_cated_y)
+
+        cated_y = torch.nn.utils.rnn.pad_sequence(cated_y, batch_first=False, padding_value=self.args.audio_pad_token)
+        assert cated_y.shape == torch.Size([max(new_y_lens),len(inserted_y), self.args.n_codebooks]), f"cated_y.shape: {cated_y.shape}, but it should be {torch.Size([max(new_y_lens,len(inserted_y), self.args.n_codebooks)])}"
+        cated_y = cated_y.permute(2,0,1) # [T,B,K]->[K,T,B]
+        assert cated_y.shape[0] == self.args.n_codebooks, cated_y.shape
+        return cated_y, torch.LongTensor(new_y_lens).to(cated_y.device)
+
+    def embed_y(self, cated_y, mask_position, mask_value):
+        embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, T, B, D]
+        assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
+        assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
+        embedded_y = embedded_y.sum(dim=0) # [K,T,B,D]->[T,B,D]
+        embedded_y = embedded_y.transpose(1,0) # [T,B,D]->[B,T,D]
+        for i in range(len(embedded_y)):
+            if len(mask_position[i]) > 0:
+                embedded_y[i, mask_position[i]] = self.mask_embedding[mask_value[i]] 
+        return embedded_y
+    
+    def prepare_input_target(self, y, y_lens):
+        # rearrange y
+        # assume y shape: [B T K], K is n_codebooks
+        assert y.shape[1] == self.args.n_codebooks, y.shape
+        # sample mask_intervals
+        mask_intervals, non_mask_intervals = self.prepare_mask_intervals(y_lens)
+
+        # need to have EOG in each section (SOG will be generated by the pattern class)
+        # but mask can be inserted later after we have shifted the input
+        # y could be rearranged in this way:
+        # [
+        # [tensor[4, 12], tensor[4, 45], tensor[4, 102], tensor[4, 32]], tensor[4, 22]],
+        # [tensor[4, 44], tensor[4, 56], tensor[4, 19]],
+        # ...
+        # ]
+        # for the first list of tensors (4 tensors), first 3 tensors are non_masked part, last 2 are masked part.
+        # NOTE #non_masked_part = #masked_part + 1
+        # NOTE *these are also the targets*
+        # added eog at the end of each segment (masked segment and unmasked segment)
+        rearranged_y = self.rearrange(y, non_mask_intervals, mask_intervals)
+        targets = rearranged_y # each element in each sample is of shape [K T]
+        assert targets[0][0].shape[0] == self.args.n_codebooks, targets[0][0].shape
+
+        # next we need to apply pattern shifting to each tensor, after which, we'll replace the starting tokens of each section with a token that's different from the special padding token
+        #  [[5, 1, 2, 3, 4, 5, 5],
+        #  [5, 5, 1, 2, 3, 4, 5],
+        #  [5, 5, 5, 1, 2, 3, 4]]
+        shifted_y, patterns = self.shift(rearranged_y) # each element [K S]
+        assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape[0]
+
+
+        # then, insert mask token at the intersection of each tensor (we want to decide the arrangement of the mask (shuffle or not)), we better have a separate nn.embedding for it
+        # we also need to record the position of the inserted mask
+        inserted_y, mask_position, mask_value = self.insert_mask(shifted_y)
+        assert inserted_y[0][0].shape[0] == self.args.n_codebooks, inserted_y[0][0].shape[0]
+        assert inserted_y[0][1].shape == torch.Size((self.args.n_codebooks, 1)), f"this should be a mask, so should have shape {(self.args.n_codebooks, 1)}, but it's {inserted_y[0][1].shape}"
+        
+        # then concat tensors that belong to the same sample (in order) then get the length of each sample, and then stack them in batch dimension, pad them with pad_token
+        cated_y, new_y_lens = self.cat_y(inserted_y, mask_position, y_lens) # KTB
+        assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], len(inserted_y)))
+        
+
+        # embed remember to separately embed the mask tokens
+        embedded_y = self.embed_y(cated_y, mask_position, mask_value) #BTD
+        assert embedded_y.shape[1:] == torch.Size((max(new_y_lens), self.args.d_model)), embedded_y.shape
+        
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y)
+
+        # make attention mask and padding mask
+        y_padding_mask = make_pad_mask(new_y_lens).to(y.device)
+        y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y_padding_mask.device)
+        return y_input, new_y_lens, targets, y_padding_mask, y_attention_mask, mask_position, patterns
+
+    def remove_mask(self, logits, mask_position, new_y_lens):
+        # logits: [B K S card]
+        logits_use = []
+        for i in range(len(logits)):
+            non_mask_positions = [-1] + mask_position[i] + [new_y_lens[i]]
+            non_mask_intervals = [[non_mask_positions[i]+1, non_mask_positions[i+1]] for i in range(len(non_mask_positions)-1)]
+            cur_logits_use = [logits[i, :, l:r] for l,r in non_mask_intervals]
+            logits_use.append(cur_logits_use)
+        
+        return logits_use
+    
+    def revert_pattern(self, patterns, logits_use):
+        logits_final = []
+        logit_masks = []
+        for i in range(len(logits_use)):
+            cur_logits = [
+                item.unsqueeze(0).permute(0, 3, 1, 2).contiguous() for item in logits_use[i]
+            ] # each item is of shape [1 K S card] [1 card K S]
+            cur_logits_final = [
+                cur_pattern.revert_pattern_logits(
+                item, 0, keep_only_valid_steps=False
+                )
+                for cur_pattern, item in zip(patterns[i], cur_logits)
+            ] # if input output order doesn't match, this step will give an error
+            cur_logits_final_ret = [item[0].permute(0,2,3,1).squeeze(0) for item in cur_logits_final] # each element is of shape [K,T,card]
+            logits_final.append(cur_logits_final_ret)
+            logit_masks.append([item[2] for item in cur_logits_final])
+
+        return logits_final, logit_masks
+
+    def dec_forward(
+            self, 
+            x_input, 
+            x_lens,
+            x_attention_mask,
+            x_padding_mask,
+            y_input,
+            new_y_lens,
+            y_attention_mask,
+            y_padding_mask,
+            past=None,
+            last_3_tokens=False
+        ):
+            x_attn_mask = F.pad(
+                x_attention_mask,
+                (0, new_y_lens.max()),
+                value=True,
+            ) # x attn to all x, doesn't attn to any y, this follow figure 3 of the valle paper
+            y_attn_mask = F.pad(
+                y_attention_mask,
+                (x_lens.max(), 0), # y is padded at the front
+                value=False,
+            ) # y attn to all x, for y itself use lower triangle mask to ensure autoregressive
+            xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)
+
+            # merge key padding and attention masks
+            bsz, src_len = x_input.shape[0], x_lens.max() + new_y_lens.max()
+            xy_padding_mask = torch.concat([x_padding_mask, y_padding_mask], dim=1)
+            _xy_padding_mask = (
+                xy_padding_mask.view(bsz, 1, 1, src_len)
+                .expand(-1, self.args.nhead, -1, -1)
+                .reshape(bsz * self.args.nhead, 1, src_len)
+            )
+            xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)
+
+            new_attn_mask = torch.zeros_like(xy_attn_mask)
+            new_attn_mask.masked_fill_(xy_attn_mask, float("-inf"))
+            xy_attn_mask = new_attn_mask
+
+            xy_input = torch.cat([x_input, y_input], dim=1)
+
+            if past == None: # do not use kvcache
+                out, _ =  self.decoder((xy_input, None), mask=xy_attn_mask)
+                return out[:, x_lens.max():], None
+            else: # use kvcache
+                if past.ndim > 3: # uses kvcache, only need to pass the last tokens, this doesn't work with multi-span speech editing yet
+                    if last_3_tokens:
+                        xy_input = xy_input[:, -3:]
+                        xy_attn_mask = xy_attn_mask[:, -3:]
+                    else:
+                        xy_input = xy_input[:, -1:]
+                        xy_attn_mask = xy_attn_mask[:, -1:]
+
+                out, present =  self.decoder((xy_input, None), mask=xy_attn_mask, past=past)
+                if isinstance(out, tuple): # get rid of stage_embedding
+                    out = out[0]
+
+                if out.shape[1] > x_lens.max(): # the first pass, not kvcache yet
+                    return out[:, x_lens.max():], present
+                else: # used kvcache
+                    return out, present
+
+    def forward(self, batch):
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (N, S).
+          x_lens:
+            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (N, K, T).
+            where K is the number of codebooks
+          y_lens:
+            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
+            before padding.
+        """
+        x, x_lens, y, y_lens = batch["x"], batch["x_lens"], batch["y"], batch["y_lens"]
+        x = x[:, :x_lens.max()] # this deal with gradient accumulation, where x_lens.max() might not be longer than the length of the current slice of x
+        y = y[:, :y_lens.max()]
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3 and y.shape[1] == self.args.n_codebooks, y.shape
+        assert y_lens.ndim == 1, y_lens.shape
+        # makes attention mask and padding mask for x
+        x_padding_mask = make_pad_mask(x_lens).to(x.device)
+        x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x_padding_mask.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input)
+        y_input, new_y_lens, targets, y_padding_mask, y_attention_mask, mask_position, patterns = self.prepare_input_target(y, y_lens)
+        y_out = self.dec_forward(
+                    x_input, 
+                    x_lens,
+                    x_attention_mask,
+                    x_padding_mask,
+                    y_input,
+                    new_y_lens,
+                    y_attention_mask,
+                    y_padding_mask
+                )
+        y_out = y_out[0] # no kv-caching during training
+        assert y_out.shape == y_input.shape, f"y_out.shape: {y_out.shape}, y_input.shape: {y_input.shape}" # [B S D]
+        
+        logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card]
+        # take out the mask token (using mask_position and new_y_lens) and revert (using function provided by self.pattern)
+        assert logits.shape[1] == self.args.n_codebooks and logits.shape[3] == self.n_audio_tokens[0], logits.shape
+
+        logits_use = self.remove_mask(logits, mask_position, new_y_lens)
+        
+        # revert the pattern shift for each logits section in each sample
+        logits_final, logit_masks = self.revert_pattern(patterns, logits_use)
+        assert logits_final[0][0].shape[0] == self.args.n_codebooks and logits_final[0][0].shape[2] == self.n_audio_tokens[0], f"it is: {logits_final[0][0].shape}, but should be [K, T, card]"
+        # testing
+        sample_to_test = 0
+        assert len(logits_final[sample_to_test]) == len(targets[sample_to_test]), f"{len(logits_final[sample_to_test])}, {len(targets[sample_to_test])}"
+        temp = sum([logits_final[sample_to_test][i].shape[:-1] != targets[sample_to_test][i].shape for i in range(len(targets[sample_to_test]))])
+        assert temp == 0, f"none equal positions: {temp}, total number of elements: {len(targets[sample_to_test])}"
+
+        logit_masked = sum([(item==False).any() for cur_mask in logit_masks for item in cur_mask])
+        assert logit_masked == 0, logit_masks
+
+        logits = torch.cat([torch.cat(item, dim=1) for item in logits_final], dim=1) # [K, T1+T2+T3+..., card]
+        targets = torch.cat([torch.cat(item, dim=1) for item in targets], dim=1) # [K, T1+T2+T3+...]
+        assert targets.shape[0] == logits.shape[0], f"{targets.shape}, {logits.shape}"
+        loss = []
+        ntokens = []
+        top10acc = []
+        for k, (logit, target) in enumerate(zip(logits, targets)):
+            loss.append(F.cross_entropy(logit, target, reduction='mean'))
+            top10acc.append(self.accuracy_metrics[k](logit.detach(), target))
+            ntokens.append(len(logit))
+        
+        all_ntokens = sum(ntokens)
+        if self.args.codebook_weight != None:
+            codebook_weight = eval(self.args.codebook_weight)
+        else:
+            codebook_weight = [1.] * self.args.n_codebooks
+        loss = sum([l*nt*cw for l, nt, cw in zip(loss, ntokens, codebook_weight)])
+        top10acc_by_codebook = [t10a*nt for t10a, nt in zip(top10acc, ntokens)]
+        top10acc = sum(top10acc_by_codebook)
+        ntokens = torch.tensor(all_ntokens).to(logits.device)
+
+        return {
+            "loss": loss,
+            "top10acc": top10acc,
+            "top10acc_by_codebook": top10acc_by_codebook,
+            "effective_ntoken": ntokens,
+        }
+    
+    def inference(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        mask_interval: list[torch.Tensor],
+        top_k: int=-100,
+        top_p: float=1.0,
+        temperature: float=1.0,
+        stop_repetition: int=-1,
+        kvcache: int=1,
+        silence_tokens: list[int]=[1388,1898,131],
+    ) -> torch.Tensor:
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (1, L).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, K).
+          mask_interval:
+            a list of tensors of shape (M, 2). contains M mask_start and mask_end. list length is actually 1, because we only support single sample inference for now
+          top_k: (`optional`) int
+            The number of highest probability tokens to keep for top-k-filtering. Default to -100.
+          top_p: (`optional`) float
+            For Neucleus sampling
+          temperature: (`optional`) float
+            The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+          eog_coef: (`optional`) float
+            if 0, no change to eog token logits, otherwise, will adjust eog token logit based on the difference between acoustic token and phn token length
+          stop_repetition (`optional`) int
+            if not -1, will set the logits of a token that repeated this many times to be -100000, to avoid generating it again. This only apply to tokens from the first codebook
+          allowed_repeat_tokens (`optional`) list of ints
+            by inspecting the validation set, get a few tokens that indeed repeat a significant amount of time, and exclude those tokens from prevent repetition
+          ultimate_stop_repetition (`optional`) int
+            no matter that token it is, stop repetition once after this number
+        """
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        if self.args.special_first:
+            y = y + int(self.args.n_special)
+        y = y.transpose(2,1) # [1,T,K] -> [1,K,T]
+        assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding
+        assert mask_interval.shape == torch.Size((1, mask_interval.shape[1], 2)), mask_interval
+
+        # make x attention mask and x_input
+        x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device)
+        # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input)
+
+        # make initial y_input
+
+        # make mask_interval and non_mask_interval
+        y_len = y.shape[2]
+        y_lens = torch.LongTensor([y_len]).to(y.device)
+        mask_interval = mask_interval[0]
+        starts =  [item[0].item() for item in mask_interval] + [y_len]
+        ends = [0] + [item[1].item() for item in mask_interval]
+        mask_intervals = [[
+            (item[0].item(), item[1].item()) for item in mask_interval
+        ]] # a werid name change, mask_interval is input, now is mask_intervals, with one more dimension
+        non_mask_intervals = [[
+            (ns, ne) for ns, ne in zip(ends, starts)
+        ]]
+        
+        # rearrange y
+        # will add have EOG in each section (SOG will be generated by the pattern class)
+        # but mask can be inserted later after we have shifted the input
+        # y could be rearranged in this way:
+        # [
+        # [tensor[4, 12], tensor[4, 45], tensor[4, 102], tensor[4, 32]], tensor[4, 22]],
+        # [tensor[4, 44], tensor[4, 56], tensor[4, 19]],
+        # ...
+        # ]
+        # for the first list of tensors (4 tensors), first 3 tensors are non_masked part, last 2 are masked part.
+        # NOTE #non_masked_part = #masked_part + 1
+        rearranged_y = self.rearrange(y, non_mask_intervals, mask_intervals)
+        assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape
+
+        # shift each element of y
+        # next we need to apply pattern shifting to each tensor, after which, we'll replace the starting tokens of each section with a token that's different from the special padding token
+        #  [
+        #  [empty, 1, 2, 3, eog, empty, empty, empty],
+        #  [empty, empty, 1, 2, 3, eog, empty, empty],
+        #  [empty, empty, empty, 1, 2, 3, eog, empty],
+        #  [empty, empty, empty, empty, 1, 2, 3, eog]
+        # ]
+        shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y
+        assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
+
+        # insert mask token at the intersction of each tensor, but *actually inserted eog as place holder*
+        # the position of inserted mask is also recorded
+        # and the mask_value, the index of the mask emb is recorded
+        inserted_y, mask_position, mask_value = self.insert_mask(shifted_y)
+        assert inserted_y[0][0].shape[0] == self.args.n_codebooks, inserted_y[0][0].shape[0]
+        assert inserted_y[0][1].shape == torch.Size((self.args.n_codebooks, 1)), f"this should be a mask, so should have shape {(self.args.n_codebooks, 1)}, but it's {inserted_y[0][1].shape}"
+        
+        # then concat tensors that belong to the same sample (in order) then get the length of each sample, and then stack them in batch dimension, pad them with pad_token
+        cated_y, new_y_lens = self.cat_y(inserted_y, mask_position, y_lens) # KTB
+        assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], len(inserted_y)))
+        assert not (cated_y == self.args.audio_pad_token).any(), cated_y
+
+        ### NOTE this is different from forward, as we will remove the masked tokens
+        ### say there are two masked region
+        ### the cated_y should be like
+        ### [empty a a a a mask0 empty b b b mask1 empty c c mask0 empty]
+        ### which means we need to take the part after the last empty out
+        num_mask = len(mask_position[0])//2
+        assert num_mask == len(mask_position[0])/2, mask_position
+        cated_y = cated_y[:, :mask_position[0][num_mask]+2] # of shape [K,T,B]
+        # logging.info(f"mask_position[0][num_mask]+2: {mask_position[0][num_mask]+2}")
+        more_mask_value = mask_value[0][num_mask+1:] # NOTE this will be used in the generation loop for reference for inserting mask embedding
+        new_y_lens[0] = mask_position[0][num_mask]+2
+        mask_position[0] = mask_position[0][:num_mask+1]
+        assert mask_position[0][num_mask]+2 == cated_y.shape[1], f"num_mask: {num_mask}, mask_position: {mask_position}, cated_y.shape: {cated_y.shape}"
+
+        # embed: remember to separately embed the mask tokens
+        embedded_y = self.embed_y(cated_y, mask_position, [mask_value[0][:num_mask+1]]) #BTD
+        # assert embedded_y.shape == torch.Size((y.shape[0], max(new_y_lens), self.args.d_model)), embedded_y.shape
+
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y)
+
+        # make attention mask and padding mask
+        y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
+        # y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
+
+        x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device)
+        y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
+
+
+        codebook_eog = [False] * self.args.n_codebooks
+        generated = [] # doesn't contain any empty_token, contains eog
+        cur_generated = []
+        # say 0 is empty, 4 is eog
+        # tensor([[ 1,  2,  3,  4,  0,  0],
+        #         [ 0,  1,  2,  3,  4,  0],
+        #         [ 0,  0,  1,  2,  3,  4]])
+        num_gen = []
+        cur_num_gen = 0
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
+        consec_silence_count = 0
+        prev_token = None
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        # prepare the cache placeholder
+        # n_layers, 2, bsz, num_heads, src_len, head_dim
+        past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None
+        # handle multi-span kv-cache
+        new_masked_span = False
+        
+        def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen):
+            if n_eog == 0:
+                logits_adjust = logits
+                for jj in range(1,self.args.n_codebooks):
+                    logits_adjust[jj][self.args.eog] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                ##################### silence repetition handling #####################
+                if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition:
+                    if logits_adjust[0, prev_token] < 0:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (consec_silence_count - (stop_repetition-1))
+                    else:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (consec_silence_count - (stop_repetition-1))
+                ##################### silence repetition handling #####################
+                if type(logits_adjust) == list:
+                    samples_list= []
+                    for logit in logits_adjust:
+                        # print(logit)
+                        # print(logit.shape)
+                        cur_sample = topk_sampling(
+                            logit.unsqueeze(0), top_k=top_k, top_p=top_p, temperature=temperature
+                        ) # [1, 1]
+                        samples_list.append(cur_sample)
+                    samples = torch.cat(samples_list, dim=0) # [K, 1]
+                else:
+                    samples = topk_sampling(
+                            logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
+                        ) # [K, 1]
+                assert samples.shape == torch.Size((self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}"
+                if cur_num_gen < self.args.n_codebooks-1:
+                    for jj in range(1, self.args.n_codebooks - cur_num_gen):
+                        samples[-jj, 0] = self.args.empty_token
+
+                if (
+                    samples[0,0] == self.args.eog or torch.argmax(logits[0], dim=-1) == self.args.eog or y_input.shape[1] > x_lens[0] * 10
+                ): # last one means y is already too long, shouldn't happen, but put it here
+                    samples[0,0] = self.args.eog
+                    codebook_eog[0] = True
+                ##################### silence repetition handling #####################
+                ##################### silence repetition handling #####################
+                if samples[0,0] in silence_tokens and samples[0,0] == prev_token:
+                    consec_silence_count += 1
+                else:
+                    consec_silence_count = 0
+                prev_token = samples[0,0]
+                ##################### silence repetition handling #####################
+                ##################### silence repetition handling #####################
+                return samples, codebook_eog, prev_token, consec_silence_count
+            else:
+                assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
+                logits_adjust = logits
+                for jj in range(n_eog+1,self.args.n_codebooks):
+                    logits_adjust[jj][self.args.eog] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                if type(logits_adjust) == list:
+                    samples_list= []
+                    for logit in logits_adjust:
+                        cur_sample = topk_sampling(
+                            logit.unsqueeze(0), top_k=top_k, top_p=top_p, temperature=temperature
+                        ) # [1, 1]
+                        samples_list.append(cur_sample)
+                    samples = torch.cat(samples_list, dim=0) # [K, 1]
+                else:
+                    samples = topk_sampling(
+                            logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
+                        ) # [K, 1]
+                for jj in range(n_eog):
+                    samples[jj, 0] = self.args.empty_token
+                samples[n_eog, 0] = self.args.eog
+                codebook_eog[n_eog] = True
+                return samples, codebook_eog, prev_token, consec_silence_count
+
+        while True:
+            y_out, present = self.dec_forward(
+                                    x_input, 
+                                    x_lens,
+                                    x_attention_mask,
+                                    x_padding_mask,
+                                    y_input,
+                                    new_y_lens,
+                                    y_attention_mask,
+                                    y_padding_mask,
+                                    past=past,
+                                    last_3_tokens = new_masked_span
+                                    )
+            if new_masked_span:
+                new_masked_span = False
+
+            if past != None:
+                past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype)
+
+            y_out = y_out[:, -1:] # only take the last one
+
+            logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], B==S==1, so [1 K 1 card]
+            logits = logits.squeeze(0).squeeze(1) # [K card]
+            assert logits.shape == torch.Size((self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}"
+
+            n_eog = sum(codebook_eog)
+            assert n_eog < self.args.n_codebooks
+            if self.args.eos > 0: # eos stands for end-of-sentence, which shouldn't be used as we are doing speech editing
+                for jj in range(self.args.n_codebooks):
+                    logits[jj][self.args.eos] = -10000.
+            # need to use a helper function to hand different n_eog cases
+            samples, codebook_eog, prev_token, consec_silence_count = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen)
+            cur_num_gen += 1
+            cur_generated.append(samples.squeeze(-1)) # [K,1] -> [K]
+            # get samples_emb
+            samples_emb = torch.stack([self.audio_embedding[k](samples[k]) for k in range(self.args.n_codebooks)], dim=0) # [K,1,D]
+            samples_emb = samples_emb.sum(dim=0,keepdim=True) # [1,1,D]
+
+            if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done
+                # re-init
+                codebook_eog = [False] * self.args.n_codebooks
+                num_gen.append(cur_num_gen)
+                cur_num_gen = 0
+                generated.append(cur_generated)
+                cur_generated = []
+
+                # if the current mask span is the last span, then all done
+                # else
+                # append the next mask token and the four empty tokens to start the next generation 
+                if len(more_mask_value) > 0:
+                    next_mask_ind = more_mask_value.pop(0)
+                    mask_emb = self.mask_embedding[next_mask_ind].unsqueeze(0).unsqueeze(0) # [1,1,D]
+                    assert mask_emb.shape == torch.Size((1,1,self.args.d_model)), mask_emb.shape
+                    empty_token = torch.LongTensor([self.args.empty_token]).to(y.device)
+                    empty_emb = torch.stack([
+                        self.audio_embedding[k](empty_token) for k in range(self.args.n_codebooks)], dim=0
+                    ).sum(dim=0, keepdim=True) # [1,1,D]
+                    assert empty_emb.shape == torch.Size((1,1,self.args.d_model)), empty_emb.shape
+                    extra_emb = torch.cat([mask_emb, empty_emb], dim=1) # [1,2,D]
+                    samples_emb = torch.cat([samples_emb, extra_emb], dim=1) # [1,3,D] # prev_last_token, mask_token, empty token
+                    assert samples_emb.shape == torch.Size((1,3,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
+                    ##################### silence repetition handling #####################
+                    ##################### silence repetition handling #####################
+                    consec_silence_count = 0
+                    prev_token = None
+                    ##################### silence repetition handling #####################
+                    ##################### silence repetition handling #####################
+
+                    # handling kv-caching for multi-span editing
+                    new_masked_span = True
+                else:
+                    break
+            else:
+                assert samples_emb.shape == torch.Size((1,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
+
+            embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
+            # positional embedding
+            y_input = self.audio_positional_embedding(embedded_y) # [B T D]
+            # make attention mask and padding mask
+            y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
+            new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
+            y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
+        
+        assert len(generated) == num_mask, f"len(generated): {len(generated)}, num_mask: {num_mask}"
+
+        # # combine non_masked_span with generated spans
+        # first need to shift the generated part back
+        flatten_gen = []
+        for l, orig_span in enumerate(generated):
+            span = torch.stack(orig_span, dim=0) # [T K]
+            span = span.transpose(1,0) # [K, T]
+            assert span.shape[0] == self.args.n_codebooks, span.shape
+            unshifted_span = []
+            for j, s in enumerate(span):
+                start_from = j
+                end_at = - (self.args.n_codebooks - start_from)
+                unshifted_span.append(s[start_from:end_at])
+            unshifted_span = torch.stack(unshifted_span, dim=0)
+
+            assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
+            flatten_gen.append(unshifted_span)
+        # logging.info(f"unshfited_span: {unshifted_span.shape}")
+        # raise
+        assert len(non_mask_intervals[0]) - 1 == len(flatten_gen), f"len(non_mask_intervals[0]): {len(non_mask_intervals[0])}, len(flatten_gen): {len(flatten_gen)}"
+        res = []
+        for orig_interval, gen in zip(non_mask_intervals[0], flatten_gen):
+            res.append(y[0, :, orig_interval[0]:orig_interval[1]])
+            res.append(gen)
+        res.append(y[0, :, non_mask_intervals[0][-1][0]:non_mask_intervals[0][-1][1]])
+        res = torch.cat(res, dim=1).unsqueeze(0) # [K,new_T] -> [1, K, new_T]
+
+        expected_y_len = y_len - sum([item[1] - item[0] for item in mask_intervals[0]]) + sum([item - self.args.n_codebooks for item in num_gen])
+        assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len - sum([item[1] - item[0] for item in mask_interval]) + sum([item - self.args.n_codebooks for item in num_gen]): {y_len}-{sum([item[1] - item[0] for item in mask_interval])} + {sum([item - self.args.n_codebooks for item in num_gen])}"
+        
+        if self.args.special_first:
+            res = res - int(self.args.n_special)
+
+        return res
+
+    def inference_tts(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        top_k: int=-100,
+        top_p: float=1.0,
+        temperature: float=1.0,
+        stop_repetition: int=3,
+        kvcache: int=1,
+        silence_tokens: list[int]=[1388,1898,131],
+        *kargs
+    ) -> torch.Tensor:
+        """
+        different from inference_tts, this implementation uses kvcache, which should have significant speed up
+        Args:
+          x:
+            A 2-D tensor of shape (1, L).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, K).
+          top_k: (`optional`) int
+            The number of highest probability tokens to keep for top-k-filtering. Default to -100.
+          top_p: (`optional`) float
+            For Neucleus sampling
+          temperature: (`optional`) float
+            The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+        """
+        eog_inference = self.args.eos if self.args.eos>0 else self.args.eog
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        if self.args.special_first:
+            y = y + int(self.args.n_special)
+        y = y.transpose(2,1) # [1,T,K] -> [1,K,T]
+        assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding
+
+        # make x attention mask and x_input
+        x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device)
+        # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input)
+
+        y_len = y.shape[2]
+        y_lens = torch.LongTensor([y_len]).to(y.device)
+
+        # rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario
+        rearranged_y = [[y[0]]]
+        assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape
+
+        # shift y to create the delayed pattern
+        shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y
+        assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
+        assert len(shifted_y[0]) == 1, len(shifted_y[0])
+
+        # below is different from forward or inference
+        # where we cut this shifted part 
+        shifted_y[0][0] = shifted_y[0][0][:, :-(self.args.n_codebooks-1)]
+        assert not (shifted_y[0][0][self.args.n_codebooks:] == self.args.empty_token).any() and not (shifted_y[0][0][self.args.n_codebooks:] == self.args.eog).any(), shifted_y[0][0]
+
+        # next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that 
+        # next section is concate tensors of each sample to one tensor, which we also don't need
+        cated_y = shifted_y[0][0].unsqueeze(-1) #[K,S]->[K,S,B]
+        new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device)
+        assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1))
+        assert not (cated_y == self.args.audio_pad_token).any(), cated_y
+
+        # replace tokens in y with the embeddings, add sum codebooks up
+        embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, S, B, D]
+        assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
+        assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
+        embedded_y = embedded_y.sum(dim=0) # [K,S,B,D]->[S,B,D]
+        embedded_y = embedded_y.transpose(1,0) # [S,B,D]->[B,S,D]
+        
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y)
+
+        # make attention mask and padding mask
+        y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
+
+        x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device)
+        y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
+
+        # entering the generation stage
+        # starting from line 708
+        codebook_eog = [False] * self.args.n_codebooks
+        generated = [] # doesn't contain any empty token, contain eog
+        cur_generated = []
+        # say 0 is empty, 4 is eog
+        # tensor([[ 1,  2,  3,  4,  0,  0],
+        #         [ 0,  1,  2,  3,  4,  0],
+        #         [ 0,  0,  1,  2,  3,  4]])
+        num_gen = []
+        cur_num_gen = 0
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
+        consec_silence_count = 0
+        prev_token = None
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+
+        # prepare the cache placeholder
+        # n_layers, 2, bsz, num_heads, src_len, head_dim
+        past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen):
+            if n_eog == 0:
+                logits_adjust = logits
+                for jj in range(1,self.args.n_codebooks):
+                    logits_adjust[jj][eog_inference] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                if cur_num_gen <= self.args.encodec_sr // 5: # this shouldn't happen, but just in case the model stopped too early
+                    logits_adjust[0][eog_inference] = -10000
+                ##################### silence repetition handling #####################
+                if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition:
+                    if logits_adjust[0, prev_token] < 0:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (consec_silence_count - (stop_repetition-1))
+                    else:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (consec_silence_count - (stop_repetition-1))
+                ##################### silence repetition handling #####################
+                samples = topk_sampling(
+                        logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
+                    ) # [K, 1]
+                assert samples.shape == torch.Size((self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}"
+                if cur_num_gen < self.args.n_codebooks-1:
+                    for jj in range(1, self.args.n_codebooks - cur_num_gen):
+                        samples[-jj, 0] = self.args.empty_token
+
+                if (
+                    samples[0,0] == eog_inference or torch.argmax(logits[0], dim=-1) == eog_inference or y_input.shape[1] > x_lens[0] * (self.args.encodec_sr//5)
+                ): # last one means y is already too long, shouldn't happen, but put it here
+                    samples[0,0] = eog_inference
+                    codebook_eog[0] = True
+                ##################### silence repetition handling #####################
+                if samples[0,0] in silence_tokens and samples[0,0] == prev_token:
+                    consec_silence_count += 1
+                else:
+                    consec_silence_count = 0
+                prev_token = samples[0,0]
+                ##################### silence repetition handling #####################
+                return samples, codebook_eog, prev_token, consec_silence_count
+            else:
+                assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
+                logits_adjust = logits
+                for jj in range(n_eog+1,self.args.n_codebooks):
+                    logits_adjust[jj][eog_inference] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                samples = topk_sampling(
+                        logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature
+                    ) # [K, 1]
+                for jj in range(n_eog):
+                    samples[jj, 0] = self.args.empty_token
+                samples[n_eog, 0] = eog_inference
+                codebook_eog[n_eog] = True
+                return samples, codebook_eog, prev_token, consec_silence_count
+        while True:
+            y_out, present = self.dec_forward(
+                            x_input, 
+                            x_lens,
+                            x_attention_mask,
+                            x_padding_mask,
+                            y_input,
+                            new_y_lens,
+                            y_attention_mask,
+                            y_padding_mask,
+                            past=past
+                        )
+            if past != None:
+                past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype)
+
+
+            y_out = y_out[:, -1:] # only take the last token
+            logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], B==S==1, so [1 K 1 card]
+            logits = logits.squeeze(0).squeeze(1) # [K card]
+            assert logits.shape == torch.Size((self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}"
+
+            n_eog = sum(codebook_eog)
+            assert n_eog < self.args.n_codebooks
+            if self.args.eos > 0: # if we are using end-of-sentence token (which is used by default), eog shouldn't be used here, as there is no masked spans
+                for jj in range(self.args.n_codebooks):
+                    logits[jj][self.args.eog] = -10000.
+            
+            samples, codebook_eog, prev_token, consec_silence_count = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen)
+            
+            cur_num_gen += 1
+            cur_generated.append(samples.squeeze(-1)) # [K,1] -> [K]
+
+            # samples.shape is [K,1]
+            # ge samples_emb
+            samples_emb = torch.stack([self.audio_embedding[k](samples[k]) for k in range(self.args.n_codebooks)], dim=0) # [K,1,D]
+            samples_emb = samples_emb.sum(dim=0,keepdim=True) # [1,1,D]
+
+            if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done
+                codebook_eog = [False] * self.args.n_codebooks
+                num_gen.append(cur_num_gen)
+                cur_num_gen = 0
+                generated.append(cur_generated)
+                cur_generated = []
+                break
+            else:
+                assert samples_emb.shape == torch.Size((1,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
+            
+            embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
+            y_input = self.audio_positional_embedding(embedded_y) # [B T D]
+            # make attention mask and padding mask
+            y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
+            new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device)
+            y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
+        
+        assert len(generated) == 1, f"len(generated): {len(generated)}"
+
+        # revert the pattern
+        flatten_gen = []
+        for l, orig_span in enumerate(generated):
+            span = torch.stack(orig_span, dim=0) # [T, K]
+            span = span.transpose(1,0) # [K, T]
+            assert span.shape[0] == self.args.n_codebooks, span.shape
+            unshifted_span = []
+            for j, s in enumerate(span):
+                start_from = j
+                end_at = - (self.args.n_codebooks - start_from)
+                unshifted_span.append(s[start_from:end_at])
+            unshifted_span = torch.stack(unshifted_span, dim=0)
+
+            assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
+
+            flatten_gen.append(unshifted_span)
+        assert len(flatten_gen) == 1, len(flatten_gen)
+        
+        # combine 
+        res = [y[0], flatten_gen[0]]
+        res = torch.cat(res, dim=1).unsqueeze(0) # [K, new_t] -> [1, K, new_T]
+
+        expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen])
+        assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}"
+
+        if self.args.special_first:
+            res = res - int(self.args.n_special)
+            flatten_gen = flatten_gen - int(self.args.n_special)
+
+        return res, flatten_gen[0].unsqueeze(0)
+
+
+    def inference_tts_batch(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        top_k: int=-100,
+        top_p: float=1.0,
+        temperature: float=1.0,
+        stop_repetition: int=3,
+        kvcache: int=1,
+        batch_size: int=5,
+        silence_tokens: list[int]=[1388,1898,131],
+        *kargs
+    ) -> torch.Tensor:
+        """
+        have a batch size when forward passing, but they are equivalant to same example but different random seed, therefore as long as one example generated eog, we can drop all other samlpes
+        different from inference_tts, this implementation uses kvcache, which should have significant speed up
+        Args:
+          x:
+            A 2-D tensor of shape (1, L).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, K).
+          top_k: (`optional`) int
+            The number of highest probability tokens to keep for top-k-filtering. Default to -100.
+          top_p: (`optional`) float
+            For Neucleus sampling
+          temperature: (`optional`) float
+            The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+        """
+        eog_inference = self.args.eos if self.args.eos>0 else self.args.eog
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        if self.args.special_first:
+            y = y + int(self.args.n_special)
+        y = y.transpose(2,1) # [1,T,K] -> [1,K,T]
+        assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding
+
+        # make x attention mask and x_input
+        x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device)
+        # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input)
+
+        y_len = y.shape[2]
+        y_lens = torch.LongTensor([y_len]).to(y.device)
+
+        # rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario
+        rearranged_y = [[y[0]]]
+        assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape
+
+        # shift y to create the delayed pattern
+        shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y
+        assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape
+        assert len(shifted_y[0]) == 1, len(shifted_y[0])
+
+        # below is different from forward or inference
+        # where we cut this shifted part 
+        shifted_y[0][0] = shifted_y[0][0][:, :-(self.args.n_codebooks-1)]
+        assert not (shifted_y[0][0][self.args.n_codebooks:] == self.args.empty_token).any() and not (shifted_y[0][0][self.args.n_codebooks:] == self.args.eog).any(), shifted_y[0][0]
+
+        # next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that 
+        # next section is concate tensors of each sample to one tensor, which we also don't need
+        cated_y = shifted_y[0][0].unsqueeze(-1) #[K,S]->[K,S,B]
+        new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device)
+        assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1))
+        assert not (cated_y == self.args.audio_pad_token).any(), cated_y
+
+        # replace tokens in y with the embeddings, add sum codebooks up
+        embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, S, B, D]
+        assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
+        assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
+        embedded_y = embedded_y.sum(dim=0) # [K,S,B,D]->[S,B,D]
+        embedded_y = embedded_y.transpose(1,0) # [S,B,D]->[B,S,D]
+        
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y)
+
+        # make attention mask and padding mask
+        y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
+
+        x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device)
+        y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
+
+        # entering the generation stage
+        # starting from line 708
+        codebook_eog = [False] * self.args.n_codebooks
+        generated = [] # doesn't contain any empty token, contain eog
+        cur_generated = [[] for _ in range(batch_size)]
+        # say 0 is empty, 4 is eog
+        # tensor([[ 1,  2,  3,  4,  0,  0],
+        #         [ 0,  1,  2,  3,  4,  0],
+        #         [ 0,  0,  1,  2,  3,  4]])
+        num_gen = []
+        cur_num_gen = 0
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
+        consec_silence_counts = [0 for _ in range(batch_size)]
+        prev_tokens = [None for _ in range(batch_size)]
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+
+        # prepare the cache placeholder
+        # n_layers, 2, bsz, num_heads, src_len, head_dim
+        past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}")
+        keep = None # NOTE: this very important, tells which sample to keep
+        def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_tokens, consec_silence_counts, stop_repetition, silence_tokens, cur_num_gen, keep):
+            if n_eog == 0:
+                logits_adjust = logits
+                for jj in range(1,self.args.n_codebooks):
+                    logits_adjust[:,jj,eog_inference] = -10000
+                    logits_adjust[:,jj,self.args.empty_token] = -10000
+                if cur_num_gen <= self.args.encodec_sr // 5: # this shouldn't happen, but just in case the model stopped too early
+                    logits_adjust[:,:,eog_inference] = -10000
+                ##################### silence repetition handling #####################
+                for b in range(batch_size):
+                    prev_token = prev_tokens[b]
+                    consec_silence_count = consec_silence_counts[b]
+                    if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition:
+                        if logits_adjust[b, 0, prev_token] < 0:
+                            logits_adjust[b, 0, prev_token] = logits_adjust[b, 0, prev_token] * (consec_silence_count - (stop_repetition-1))
+                        else:
+                            logits_adjust[b, 0, prev_token] = logits_adjust[b, 0, prev_token] / (consec_silence_count - (stop_repetition-1))
+                ##################### silence repetition handling #####################
+                samples = topk_sampling(
+                        logits_adjust.reshape(batch_size * self.args.n_codebooks, logits_adjust.shape[-1]), top_k=top_k, top_p=top_p, temperature=temperature
+                    ) # [B*K, 1]
+                samples = samples.reshape(batch_size, self.args.n_codebooks, 1)
+                assert samples.shape == torch.Size((batch_size, self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}"
+                for b in range(batch_size):
+                    if cur_num_gen < self.args.n_codebooks-1:
+                        for jj in range(1, self.args.n_codebooks - cur_num_gen):
+                            samples[b, -jj, 0] = self.args.empty_token
+
+                    if (
+                        samples[b,0,0] == eog_inference or torch.argmax(logits[b,0], dim=-1) == eog_inference or y_input.shape[1] > x_lens[b] * (self.args.encodec_sr//5)
+                    ): # last one means y is already too long, shouldn't happen, but put it here
+                        samples[b,0,0] = eog_inference
+                        codebook_eog[0] = True
+                        keep = b # NOTE keep is a very important variable, we only return this one, note that if eog shows up in two samples, keep will be overwritten by the later one (or the last one)
+                    ##################### silence repetition handling #####################
+                    if samples[b,0,0] in silence_tokens and samples[b,0,0] == prev_tokens[b]:
+                        consec_silence_counts[b] += 1
+                    else:
+                        consec_silence_counts[b] = 0
+                    prev_tokens[b] = samples[b,0,0]
+                ##################### silence repetition handling #####################
+                return samples, codebook_eog, prev_tokens, consec_silence_counts, keep
+            else:
+                assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
+                logits_adjust = logits
+                for jj in range(n_eog+1,self.args.n_codebooks):
+                    logits_adjust[:,jj,eog_inference] = -10000
+                    logits_adjust[:,jj,self.args.empty_token] = -10000
+                samples = topk_sampling(
+                        logits_adjust.reshape(batch_size * self.args.n_codebooks, logits_adjust.shape[-1]), top_k=top_k, top_p=top_p, temperature=temperature
+                    ) # [B, K, 1]
+                samples = samples.reshape(batch_size, self.args.n_codebooks, 1)
+                for jj in range(n_eog):
+                    samples[keep, jj, 0] = self.args.empty_token
+                samples[keep, n_eog, 0] = eog_inference
+                codebook_eog[n_eog] = True
+                return samples, codebook_eog, prev_tokens, consec_silence_counts, keep
+        while True:
+            # if cur_num_gen > 0, should have everything in kvcache, so only pass in the last token
+            # in the first generation step, we repeat each tensor to make their first dimension of length the batch size 
+            if cur_num_gen == 0:
+                assert x_input.ndim == 3 and x_input.shape[0] == 1, x_input.shape
+                assert x_padding_mask.ndim == 2 and x_padding_mask.shape[0] == 1, x_padding_mask.shape
+                assert y_input.ndim == 3 and y_input.shape[0] == 1 and y_input.shape[1] == new_y_lens[0], y_input.shape
+                assert embedded_y.ndim == 3 and embedded_y.shape[0] == 1 and embedded_y.shape[1] == new_y_lens[0], embedded_y.shape
+                x_input = x_input.repeat(batch_size, 1, 1)
+                x_lens = x_lens.repeat(batch_size)
+                # x_attention_mask = x_attention_mask.repeat(batch_size, 1, 1) # no need to work with attention mask, it doesn't contain batch dimension
+                x_padding_mask = x_padding_mask.repeat(batch_size, 1)
+                y_input = y_input.repeat(batch_size, 1, 1)
+                new_y_lens = new_y_lens.repeat(batch_size)
+                # y_attention_mask = y_attention_mask.repeat(batch_size, 1, 1) # no need to work with attention mask, it doesn't contain batch dimension
+                y_padding_mask = y_padding_mask.repeat(batch_size, 1)
+                embedded_y = embedded_y.repeat(batch_size, 1, 1) # will be used to concat with newly generated token embedding
+                past = past.repeat(1, 1, batch_size) if past != None else None
+            else:
+                assert x_input.shape[0] == batch_size and x_padding_mask.shape[0] == batch_size and y_input.shape[0] == batch_size and new_y_lens.shape[0] == batch_size, f"x_input.shape: {x_input.shape}, x_padding_mask.shape: {x_padding_mask.shape}, y_input.shape: {y_input.shape}, new_y_lens.shape: {new_y_lens.shape}"
+            y_out, present = self.dec_forward(
+                            x_input, 
+                            x_lens,
+                            x_attention_mask,
+                            x_padding_mask,
+                            y_input,
+                            new_y_lens,
+                            y_attention_mask,
+                            y_padding_mask,
+                            past=past
+                        )
+            if past != None:
+                past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype)
+
+            # if no eog emerges, y_out should have batch size of batch_size
+            if sum(codebook_eog) == 0:
+                assert y_out.shape[0] == batch_size and y_out.ndim == 3, y_out.shape
+            y_out = y_out[:, -1:] # only take the last token
+            logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], S==1, so [B K 1 card]
+            logits = logits.squeeze(2) # [B K card]
+            assert logits.shape == torch.Size((batch_size, self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}"
+
+            n_eog = sum(codebook_eog)
+            if self.args.eos > 0:
+                for jj in range(self.args.n_codebooks):
+                    logits[:,jj,self.args.eog] = -10000.
+            samples, codebook_eog, prev_tokens, consec_silence_counts, keep = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_tokens, consec_silence_counts, stop_repetition, silence_tokens, cur_num_gen, keep)
+            
+            cur_num_gen += 1
+            if sum(codebook_eog) == 0: # no eog yet, keep batch_size of samples
+                assert keep == None
+                for b in range(batch_size):
+                    cur_generated[b].append(samples[b].squeeze(-1))
+            elif sum(codebook_eog) == 1: # the first eog just showed up in this step
+                assert keep != None
+                cur_generated = cur_generated[keep]
+                cur_generated.append(samples[keep].squeeze(-1))
+            else: # we are generating the rest eogs for the 'keep' sample 
+                cur_generated.append(samples[keep].squeeze(-1))
+
+            # samples.shape is [K,1]
+            # ge samples_emb
+            samples_emb = torch.stack([self.audio_embedding[k](samples[:, k]) for k in range(self.args.n_codebooks)], dim=1) # [B, K,1,D]
+            assert samples_emb.shape == torch.Size([batch_size, self.args.n_codebooks, 1, self.args.d_model])
+            samples_emb = samples_emb.sum(dim=1,keepdim=False) # [B,1,D]
+            if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done
+                codebook_eog = [False] * self.args.n_codebooks
+                num_gen.append(cur_num_gen)
+                cur_num_gen = 0
+                generated.append(cur_generated)
+                cur_generated = [[] for _ in range(batch_size)]
+                break
+            else:
+                assert samples_emb.shape == torch.Size((batch_size,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
+            
+            embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
+            y_input = self.audio_positional_embedding(embedded_y) # [B T D]
+            # make attention mask and padding mask
+            y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
+            new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device).repeat(batch_size)
+            y_padding_mask = torch.full((batch_size,new_y_lens[0]), False).to(y.device)
+        
+        assert len(generated) == 1, f"len(generated): {len(generated)}"
+
+        # revert the pattern
+        flatten_gen = []
+        for l, orig_span in enumerate(generated):
+            span = torch.stack(orig_span, dim=0) # [T, K]
+            span = span.transpose(1,0) # [K, T]
+            assert span.shape[0] == self.args.n_codebooks, span.shape
+            unshifted_span = []
+            for j, s in enumerate(span):
+                start_from = j
+                end_at = - (self.args.n_codebooks - start_from)
+                unshifted_span.append(s[start_from:end_at])
+            unshifted_span = torch.stack(unshifted_span, dim=0)
+
+            assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
+
+            flatten_gen.append(unshifted_span)
+        assert len(flatten_gen) == 1, len(flatten_gen)
+        
+        # combine 
+        res = [y[0], flatten_gen[0]]
+        res = torch.cat(res, dim=1).unsqueeze(0) # [K, new_t] -> [1, K, new_T]
+
+        expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen])
+        assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}"
+
+        if self.args.special_first:
+            res = res - int(self.args.n_special)
+            flatten_gen = flatten_gen - int(self.args.n_special)
+
+        return res, flatten_gen[0].unsqueeze(0)

pretrained_models/encodec_4cb2048_giga.th

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

pretrained_models/giga330M.pth

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

requirements.txt

@@ -0,0 +1,9 @@
+-e git+https://github.com/facebookresearch/audiocraft.git@c5157b5bf14bf83449c17ea1eeb66c19fb4bc7f0#egg=audiocraft
+xformers==0.0.22
+torchaudio==2.0.2 
+torch==2.0.1
+phonemizer==3.2.1
+gradio==3.50.2
+nltk>=3.8.1
+openai-whisper>=20231117
+spaces