add 1s

app.py

@@ -0,0 +1,284 @@
+# Environment setup
+from pathlib import Path
+import os
+import sys
+sys.path.append(str(Path(__file__).parent))
+# FIXME add weights_only=False in /usr/local/lib/python3.10/site-packages/fairseq/checkpoint_utils.py#315
+if os.path.exists('/usr/local/lib/python3.10/site-packages/fairseq/checkpoint_utils.py'):
+    file_lines = []
+    with open('/usr/local/lib/python3.10/site-packages/fairseq/checkpoint_utils.py', 'r') as f:
+        for line in f:
+            file_lines.append(line.strip('\n'))
+    file_lines[314] = file_lines[314].replace(
+        "state = torch.load(f, map_location=torch.device(\"cpu\"))",
+        "state = torch.load(f, map_location=torch.device(\"cpu\"), weights_only=False)"
+    )
+    with open('/usr/local/lib/python3.10/site-packages/fairseq/checkpoint_utils.py', 'w') as f:
+        for line in file_lines:
+            f.write(line+'\n')
+    print('[DEBUG] added weights_only=False')
+# Run
+import spaces
+import gradio as gr
+from zipfile import ZipFile
+from typing import Literal
+from huggingface_hub import snapshot_download
+from fireredtts.models.fireredtts import FireRedTTS
+# NOTE disable verbose INFO logs
+import logging
+httpx_logger = logging.getLogger("httpx")
+httpx_logger.setLevel(logging.WARNING)
+
+# NOTE Some launching setups
+# - install fairseq manually ("python -m pip install pip==24.0")
+# - manually add weights_only=False in /usr/local/lib/python3.10/site-packages/fairseq/checkpoint_utils.py#315
+
+
+# ================================================
+#                   FireRedTTS1s Model
+# ================================================
+# Global model instance
+tts_flow: FireRedTTS = None
+tts_acollm: FireRedTTS = None
+def initiate_model(pretrained_dir: str):
+    global tts_flow, tts_acollm
+    if tts_flow is None:
+        tts_flow = FireRedTTS(
+            config_path='configs/config_24k_flow.json',
+            pretrained_path=pretrained_dir,
+        )
+    if tts_acollm is None:
+        tts_acollm = FireRedTTS(
+            config_path='configs/config_24k.json',
+            pretrained_path=pretrained_dir,
+        )
+
+
+# ================================================
+#                   Gradio
+# ================================================
+
+# i18n
+_i18n_key2lang_dict = dict(
+    # Title markdown
+    title_md_desc=dict(
+        en="FireRedTTS-1s 🔥 Streamable TTS",
+        zh="FireRedTTS-1s 🔥 可流式TTS",
+    ),
+    # Decoder choice radio
+    decoder_choice_label=dict(
+        en="Decoder Choice",
+        zh="解码器选择",
+    ),
+    decoder_choice_1=dict(
+        en="Flow Matching",
+        zh="Flow Matching",
+    ),
+    decoder_choice_2=dict(
+        en="Acoustic LLM",
+        zh="Acoustic LLM",
+    ),
+    # Speaker Prompt
+    spk_prompt_audio_label=dict(
+        en="Speaker Prompt Audio",
+        zh="参考语音",
+    ),
+    spk_prompt_text_label=dict(
+        en="Speaker Prompt Text",
+        zh="参考语音的文本",
+    ),
+    spk_prompt_text_placeholder=dict(
+        en="Speaker Prompt Text",
+        zh="参考语音的文本",
+    ),
+    # Input textbox
+    target_text_input_label=dict(
+        en="Text To Synthesis",
+        zh="待合成文本",
+    ),
+    target_text_input_placeholder=dict(
+        en="Text To Synthesis",
+        zh="待合成文本",
+    ),
+    # Generate button
+    generate_btn_label=dict(
+        en="Generate Audio",
+        zh="合成",
+    ),
+    # Generated audio
+    generated_audio_label=dict(
+        en="Generated Audio",
+        zh="合成的音频",
+    ),
+    # Warining1: incomplete prompt info
+    warn_incomplete_prompt=dict(
+        en="Please provide prompt audio and text",
+        zh="请提供说话人参考语音与参考文本",
+    ),
+    # Warining2: invalid text for target text input
+    warn_invalid_target_text=dict(
+        en="Empty input text",
+        zh="待合成文本为空",
+    ),
+)
+
+global_lang: Literal['zh', 'en'] = 'zh'
+def i18n(key):
+    global global_lang
+    return _i18n_key2lang_dict[key][global_lang]
+
+
+def check_monologue_text(text:str, prefix:str=None)->bool:
+    text = text.strip()
+    # Check speaker tags
+    if prefix is not None and (not text.startswith(prefix)):
+        return False
+    # Remove prefix
+    if prefix is not None:
+        text = text.removeprefix(prefix)
+    text = text.strip()
+    # If empty?
+    if len(text) == 0:
+        return False
+    return True
+
+
+@spaces.GPU(duration=60)
+def synthesis_function(
+    spk_prompt_audio: str,
+    spk_prompt_text: str,
+    target_text: str,
+    decoder_choice: Literal[0, 1] = 0,  # 0 means flow matching decoder
+):
+    global tts_flow, tts_acollm
+
+    # Check prompt info
+    spk_prompt_text = spk_prompt_text.strip()
+    if spk_prompt_audio is None or spk_prompt_text == "":
+        gr.Warning(message=i18n('warn_incomplete_prompt'))
+        return None
+    # Check target text
+    target_text = target_text.strip()
+    if target_text == "":
+        gr.Warning(message=i18n('warn_invalid_target_text'))
+        return None
+    
+    # Go synthesis
+    if decoder_choice == 0:
+        audio = tts_flow.synthesize(
+            prompt_wav=spk_prompt_audio,
+            prompt_text=spk_prompt_text,
+            text=target_text,
+            lang="zh",
+            use_tn=True
+        )   
+    else:
+        audio = tts_acollm.synthesize(
+            prompt_wav=spk_prompt_audio,
+            prompt_text=spk_prompt_text,
+            text=target_text,
+            lang="zh",
+            use_tn=True
+        )
+    return (24000, audio.detach().cpu().squeeze(0).numpy())
+
+
+# UI rendering
+def render_interface()->gr.Blocks:
+    with gr.Blocks(title="FireRedTTS-2", theme=gr.themes.Default()) as page:
+        # ======================== UI ========================
+        # A large title
+        title_desc = gr.Markdown(value="# {}".format(i18n('title_md_desc')))
+        with gr.Row():
+            lang_choice = gr.Radio(
+                choices=['中文', 'English'],
+                value='中文',
+                label='Display Language/显示语言',
+                type="index",
+                interactive=True,
+            )
+            decoder_choice = gr.Radio(
+                choices=[i18n('decoder_choice_1'), i18n('decoder_choice_2')],
+                value=i18n('decoder_choice_1'),
+                label=i18n('decoder_choice_label'),
+                type="index",
+                interactive=True,
+            )
+        with gr.Row():
+            # ==== Speaker Prompt ====
+            spk_prompt_text = gr.Textbox(
+                label=i18n('spk_prompt_text_label'),
+                placeholder=i18n('spk_prompt_text_placeholder'),
+                lines=5,
+            )
+            spk_prompt_audio = gr.Audio(
+                label=i18n('spk_prompt_audio_label'),
+                type="filepath",
+                editable=False,
+                interactive=True,
+            )   # Audio component returns tmp audio path
+            # ==== Target Text ====
+            target_text_input = gr.Textbox(
+                label=i18n('target_text_input_label'),
+                placeholder=i18n('target_text_input_placeholder'),
+                lines=5,
+            )
+        # Generate button
+        generate_btn = gr.Button(value=i18n('generate_btn_label'), variant="primary", size="lg")
+        # Long output audio
+        generate_audio = gr.Audio(
+            label=i18n('generated_audio_label'), 
+            interactive=False,
+        )
+
+        # ======================== Action ========================
+        # Language action
+        def _change_component_language(lang):
+            global global_lang 
+            global_lang = ['zh', 'en'][lang]
+            return [
+                # title_desc
+                gr.update(value="# {}".format(i18n('title_md_desc'))),
+                # decoder_choice
+                gr.update(label=i18n('decoder_choice_label')),
+                # spk_prompt_{audio,text}
+                gr.update(label=i18n('spk_prompt_text_label'), placeholder=i18n('spk_prompt_text_placeholder')), 
+                gr.update(label=i18n('spk_prompt_audio_label')), 
+                # target_text_input
+                gr.update(label=i18n('target_text_input_label'), placeholder=i18n('target_text_input_placeholder')),
+                # generate_btn
+                gr.update(value=i18n('generate_btn_label')), 
+                # generate_audio
+                gr.update(label=i18n('generated_audio_label')), 
+            ]
+        lang_choice.change(
+            fn=_change_component_language,
+            inputs=[lang_choice],
+            outputs=[
+                title_desc, decoder_choice,
+                spk_prompt_text, spk_prompt_audio, 
+                target_text_input, 
+                generate_btn, generate_audio,
+            ]
+        )
+        generate_btn.click(
+            fn=synthesis_function,
+            inputs=[spk_prompt_audio, spk_prompt_text, target_text_input, decoder_choice],
+            outputs=[generate_audio]
+        )
+    return page
+
+
+if __name__ == '__main__':
+    # Download model
+    snapshot_download(repo_id='FireRedTeam/FireRedTTS-1S', local_dir='pretrained_models/FireRedTTS-1S')
+    # Unzip model, weights under "pretrained_models/FireRedTTS-1S/pretrained_models"
+    with ZipFile('pretrained_models/FireRedTTS-1S/pretrained_models.zip', 'r')  as zipf:
+        zipf.extractall('pretrained_models/FireRedTTS-1S')
+    # Init model
+    initiate_model('pretrained_models/FireRedTTS-1S/pretrained_models')
+    print('[INFO] model loaded')
+    # UI
+    page = render_interface()
+    page.launch()
+

configs/config_24k.json

@@ -0,0 +1,171 @@
+{
+    "semantic_llm": {
+        "start_text_token": 32000,
+        "stop_text_token": 32001,
+        "num_text_tokens": 32002,
+        "start_audio_token": 16384,
+        "stop_audio_token": 16385,
+        "num_audio_tokens": 16386,
+        "llm_hidden_size": 1024,
+        "llm_intermediate_size": 4096,
+        "llm_num_layers": 30,
+        "llm_num_heads": 16,
+        "llm_max_audio_seq_len": 630,
+        "llm_max_text_seq_len": 402,
+        "llm_max_prompt_len": 250,
+        "code_stride_len": 640,
+        "EOS_TOKEN": 16385
+    },
+    "acoustic_llm": {
+        "n_stacks": 1,
+        "layers": 24,
+        "model_dim": 1536,
+        "heads": 16,
+        "max_text_tokens": 2048,
+        "max_speech_tokens": 2048,
+        "max_conditioning_inputs": 1,
+        "number_text_tokens": 16386,
+        "start_text_token": 16384,
+        "stop_text_token": 16385,
+        "n_frames_per_step": 1,
+        "n_heads_per_frame": 8,
+        "delay_prediction": 1,
+        "upsample_factors": 1,
+        "streaming_delayed_frames": 8,
+        "number_speech_tokens": 16386,
+        "start_speech_token": 16384,
+        "stop_speech_token": 16385,
+        "speaker_embedding_pretrained": true,
+        "speaker_embedding_ckpt": null,
+        "speaker_embedding_dim": 512,
+        "temperature": 0.5,
+        "repetition_penalty": 2.0,
+        "top_p": 0.5,
+        "top_k": 25
+    },
+    "acoustic_codec": {
+        "n_model_size": 1024,
+        "encoder_config": {
+            "ngf": 48,
+            "up_ratios": [
+                2,
+                4,
+                4,
+                4,
+                5
+            ],
+            "causal": true
+        },
+        "decoder_config": {
+            "upsample_initial_channel": 1536,
+            "ngf": 48,
+            "up_ratios": [
+                6,
+                5,
+                4,
+                4,
+                2
+            ],
+            "causal": true
+        },
+        "vq_config": {
+            "n_groups": 8,
+            "ordered": true,
+            "codebook_size": [
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128,
+                128
+            ],
+            "codebook_dim": [
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8,
+                8
+            ],
+            "requires_projection": true,
+            "decay": 0.99,
+            "threshold_ema_dead_code": 0,
+            "commitment_weight": 0.01
+        },
+        "resampler_config": {
+            "source_sr": 16000,
+            "target_sr": 16000
+        }
+    },
+    "semantic_tokenizer": {
+        "in_dim": 1024,
+        "out_dim": 80,
+        "n_model_size": 512,
+        "downsample_scales": [
+            1,
+            1,
+            1,
+            2
+        ],
+        "upsample_scales": [
+            [
+                2,
+                1
+            ],
+            [
+                2,
+                1,
+                1,
+                1
+            ]
+        ],
+        "mel_config": {
+            "style": "BigVGAN",
+            "filter_length": 1024,
+            "hop_length": 160,
+            "win_length": 640,
+            "n_mel_channels": 80,
+            "sampling_rate": 16000
+        },
+        "vq_config": {
+            "codebook_size": [
+                128,
+                128
+            ],
+            "codebook_dim": [
+                128,
+                128
+            ],
+            "requires_projection": true
+        },
+        "tree_config": [
+            {
+                "downsample_rate": 1,
+                "n_groups": 1,
+                "dropout": 0
+            }
+        ],
+        "n_samples_per_token": 640,
+        "checkpointing": true
+    }
+}

configs/config_24k_flow.json

@@ -0,0 +1,125 @@
+{
+    "semantic_llm": {
+        "start_text_token": 32000,
+        "stop_text_token": 32001,
+        "num_text_tokens": 32002,
+        "start_audio_token": 16384,
+        "stop_audio_token": 16385,
+        "num_audio_tokens": 16386,
+        "llm_hidden_size": 1024,
+        "llm_intermediate_size": 4096,
+        "llm_num_layers": 30,
+        "llm_num_heads": 16,
+        "llm_max_audio_seq_len": 630,
+        "llm_max_text_seq_len": 402,
+        "llm_max_prompt_len": 250,
+        "code_stride_len": 640,
+        "EOS_TOKEN": 16385
+    },
+    "flow": {
+        "spk_channels": 512,
+        "spk_enc_channels": 80,
+        "infer_cfg_rate": 0.7,
+        "token_emb": {
+            "channels": 512
+        },
+        "encoder": {
+            "input_size": 512,
+            "output_size": 512,
+            "num_blocks": 6,
+            "num_up_blocks": 4,
+            "normalize_before": true,
+            "up_stride": 2,
+            "pre_lookahead_len": 3,
+            "attention_heads": 4,
+            "key_bias": true,
+            "linear_units": 2048,
+            "dropout_rate": 0.0,
+            "positional_dropout_rate": 0.0,
+            "attention_dropout_rate": 0.0
+        },
+        "estimator": {
+            "in_channels": 320,
+            "out_channels": 80,
+            "mlp_ratio": 4,
+            "depth": 16,
+            "num_heads": 8,
+            "head_dim": 64,
+            "hidden_size": 512
+        }
+    },
+    "mel": {
+        "num_mels": 80,
+        "n_fft": 1920,
+        "hop_size": 480,
+        "win_size": 1920,
+        "sampling_rate": 24000,
+        "fmin": 0,
+        "fmax": 8000,
+        "center": false
+    },
+    "bigvgan": {
+        "num_mels": 80,
+        "upsample_initial_channel": 1536,
+        "resblock_kernel_sizes": [3, 7, 11],
+        "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        "upsample_rates": [5, 4, 3, 2, 2, 2],
+        "upsample_kernel_sizes": [11, 8, 7, 4, 4, 4],
+        "resblock_type": "1",
+        "snake_logscale": true,
+        "activation": "snakebeta",
+        "use_tanh_at_final": false,
+        "use_bias_at_final": false
+    },
+    "semantic_tokenizer": {
+        "in_dim": 1024,
+        "out_dim": 80,
+        "n_model_size": 512,
+        "downsample_scales": [
+            1,
+            1,
+            1,
+            2
+        ],
+        "upsample_scales": [
+            [
+                2,
+                1
+            ],
+            [
+                2,
+                1,
+                1,
+                1
+            ]
+        ],
+        "mel_config": {
+            "style": "BigVGAN",
+            "filter_length": 1024,
+            "hop_length": 160,
+            "win_length": 640,
+            "n_mel_channels": 80,
+            "sampling_rate": 16000
+        },
+        "vq_config": {
+            "codebook_size": [
+                128,
+                128
+            ],
+            "codebook_dim": [
+                128,
+                128
+            ],
+            "requires_projection": true
+        },
+        "tree_config": [
+            {
+                "downsample_rate": 1,
+                "n_groups": 1,
+                "dropout": 0
+            }
+        ],
+        "n_samples_per_token": 640,
+        "checkpointing": true
+    }
+}

fireredtts/models/fireredtts.py

@@ -0,0 +1,266 @@
+import os
+import json
+import time
+from traceback import format_exc
+import torch
+
+from transformers import AutoTokenizer
+
+from fireredtts.utils.utils import load_audio
+from fireredtts.modules.text_normalizer.utils import text_split
+from fireredtts.utils.spliter import clean_text
+from fireredtts.modules.text_normalizer.normalize import TextNormalizer
+from fireredtts.modules.semantic_tokenizer import SemanticTokenizer
+from fireredtts.modules.semantic_llm.llm_gpt2 import Speech_LLM_GPT2
+from fireredtts.models.token2audio import TwoStageCodec, FlowToken2Audio
+
+
+class FireRedTTS:
+    def __init__(self, config_path, pretrained_path, device="cuda"):
+        self.device = device
+        self.config = json.load(open(config_path))
+        self.EOS_TOKEN = self.config["semantic_llm"]["EOS_TOKEN"]
+
+        # pretrained models
+        self.tokenizer_path = os.path.join(pretrained_path, "tokenizer")
+        self.speech_tokenizer_path = os.path.join(pretrained_path, "speech_tokenizer")
+        self.semantic_llm_path = os.path.join(pretrained_path, "semantic_llm.pt")
+        assert os.path.exists(self.tokenizer_path)
+        assert os.path.exists(self.speech_tokenizer_path)
+        assert os.path.exists(self.semantic_llm_path)
+        if 'acoustic_llm' in self.config:
+            self.acoustic_llm_path = os.path.join(pretrained_path, "acoustic_llm.bin")
+            self.acoustic_codec_path = os.path.join(pretrained_path, "acoustic_codec.bin")
+            assert os.path.exists(self.acoustic_llm_path)
+            assert os.path.exists(self.acoustic_codec_path)
+        else:
+            self.flow_path = os.path.join(pretrained_path, "flow.pt")
+            self.bigvgan_path = os.path.join(pretrained_path, "bigvgan.pt")
+            assert os.path.exists(self.flow_path)
+            assert os.path.exists(self.bigvgan_path)
+
+        # text normalizer
+        self.text_normalizer = TextNormalizer()
+        # text tokenizer
+        self.text_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path)
+
+        # semantic llm
+        self.semantic_llm = Speech_LLM_GPT2(
+            start_text_token=self.config["semantic_llm"]["start_text_token"],
+            stop_text_token=self.config["semantic_llm"]["stop_text_token"],
+            num_text_tokens=self.config["semantic_llm"]["num_text_tokens"],
+            start_audio_token=self.config["semantic_llm"]["start_audio_token"],
+            stop_audio_token=self.config["semantic_llm"]["stop_audio_token"],
+            num_audio_tokens=self.config["semantic_llm"]["num_audio_tokens"],
+            llm_hidden_size=self.config["semantic_llm"]["llm_hidden_size"],
+            llm_intermediate_size=self.config["semantic_llm"]["llm_intermediate_size"],
+            llm_num_layers=self.config["semantic_llm"]["llm_num_layers"],
+            llm_num_heads=self.config["semantic_llm"]["llm_num_heads"],
+            llm_max_audio_seq_len=self.config["semantic_llm"]["llm_max_audio_seq_len"],
+            llm_max_text_seq_len=self.config["semantic_llm"]["llm_max_text_seq_len"],
+            llm_max_prompt_len=self.config["semantic_llm"]["llm_max_prompt_len"],
+            code_stride_len=self.config["semantic_llm"]["code_stride_len"],
+        )
+
+        sd = torch.load(self.semantic_llm_path, map_location=device)["model"]
+        self.semantic_llm.load_state_dict(sd, strict=True)
+        self.semantic_llm = self.semantic_llm.to(device=device)
+        self.semantic_llm.eval()
+        self.semantic_llm.init_gpt_for_inference(kv_cache=True)
+
+        # Speech tokenizer
+        self.speech_tokenizer = SemanticTokenizer(
+            config=self.config["semantic_tokenizer"], path=self.speech_tokenizer_path
+        )
+
+        # Acoustic decoder
+        if 'acoustic_llm' in self.config:
+            self.acoustic_decoder = TwoStageCodec(self.config)
+            self.acoustic_decoder.load_model(self.acoustic_llm_path, self.acoustic_codec_path)
+        else:
+            self.acoustic_decoder = FlowToken2Audio(self.config)
+            self.acoustic_decoder.load_model(self.flow_path, self.bigvgan_path)
+        self.acoustic_decoder.eval()
+        self.acoustic_decoder = self.acoustic_decoder.to(device)
+
+    def extract_spk_embeddings(self, prompt_wav):
+        audio, lsr, audio_resampled = load_audio(
+            audiopath=prompt_wav,
+            sampling_rate=16000,
+        )
+        _, _, audio_resampled24k = load_audio(
+            audiopath=prompt_wav,
+            sampling_rate=24000,
+        )
+
+        audio_resampled = audio_resampled.to(self.device)
+        audio_len = torch.tensor(
+            data=[audio_resampled.shape[1]], dtype=torch.long, requires_grad=False
+        )
+
+        # spk_embeddings：[1, 512]
+        prompt_tokens, token_lengths, spk_embeddings = self.speech_tokenizer(
+            audio_resampled, audio_len
+        )
+
+        prompt_acoustic_tokens, acoustic_llm_spk = self.acoustic_decoder.extract(
+            audio_resampled if isinstance(self.acoustic_decoder, TwoStageCodec) else audio_resampled24k, 
+            audio_len, spk_embeddings.unsqueeze(0)
+        )
+
+        return prompt_tokens, spk_embeddings, prompt_acoustic_tokens, acoustic_llm_spk
+
+    def synthesize_base(
+        self,
+        prompt_semantic_tokens,
+        prompt_acoustic_tokens,
+        spk_semantic_llm,
+        spk_acoustic_llm,
+        prompt_text,
+        text,
+        lang="auto",
+    ):
+        """_summary_
+
+        Args:
+            prompt_wav (_type_): _description_
+            prompt_text (_type_): _description_
+            text (_type_): _description_
+            lang (str, optional): _description_. Defaults to "auto".
+
+        Returns:
+            _type_: _description_
+        """
+        if lang == "en":
+            text = prompt_text + " " + text
+        else:
+            text = prompt_text + text
+
+        print("---text:\n", text)
+
+        # Pre-process prompt tokens
+        # text to tokens
+        text_tokens = self.text_tokenizer.encode(
+            text=text,
+            add_special_tokens=False,
+            max_length=10**6,
+            truncation=False,
+        )
+        # print("---decode", [self.text_tokenizer.decode([c]) for c in text_tokens])
+
+        text_tokens = torch.IntTensor(text_tokens).unsqueeze(0).to(self.device)
+
+        assert text_tokens.shape[-1] < 200
+        with torch.no_grad():
+            gpt_codes = self.semantic_llm.generate_ic(
+                cond_latents=spk_semantic_llm,
+                text_inputs=text_tokens,
+                prompt_tokens=prompt_semantic_tokens[:, :-3],
+                do_sample=True,
+                top_p=0.85,
+                top_k=30,
+                temperature=0.75,
+                num_return_sequences=7,
+                num_beams=1,
+                length_penalty=2.0,
+                repetition_penalty=5.0,
+                output_attentions=False,
+            )
+
+        seqs = []
+        for seq in gpt_codes:
+            index = (seq == self.EOS_TOKEN).nonzero(as_tuple=True)[0][0]
+            seq = seq[:index]
+            seqs.append(seq)
+
+        sorted_seqs = sorted(seqs, key=lambda i: len(i), reverse=False)
+        sorted_len = [len(l) for l in sorted_seqs]
+
+        gpt_codes = sorted_seqs[2].unsqueeze(0)
+
+        # Acoustic decoder
+        rec_wavs = self.acoustic_decoder(
+            gpt_codes, prompt_semantic_tokens, prompt_acoustic_tokens, spk_acoustic_llm
+        )
+
+        rec_wavs = rec_wavs.detach().cpu()
+        return rec_wavs
+
+    @torch.no_grad()
+    def synthesize(self, prompt_wav, prompt_text, text, lang="auto", use_tn=False):
+        """audio synthesize
+
+        Args:
+            prompt_wav (_type_): _description_
+            prompt_text (_type_): _description_
+            text (_type_): _description_
+            lang (str, optional): _description_. Defaults to "auto".
+
+        Returns:
+            _type_: _description_
+        """
+        assert lang in ["zh", "en", "auto"]
+        assert os.path.exists(prompt_wav)
+
+        (
+            prompt_semantic_tokens,
+            spk_embeddings,
+            prompt_acoustic_tokens,
+            spk_acoustic_llm,
+        ) = self.extract_spk_embeddings(prompt_wav=prompt_wav)
+
+        spk_embeddings = spk_embeddings.unsqueeze(0)
+        spk_semantic_llm = self.semantic_llm.reference_embedding(spk_embeddings)
+
+        # print("---prompt_semantic_tokens:\n", prompt_semantic_tokens)
+        # print("---spk_embeddings:\n", spk_embeddings)
+
+        # clean text
+        prompt_text = clean_text(prompt_text)
+        text = clean_text(text=text)
+
+        if use_tn:
+            substrings = text_split(text=text)
+
+            out_wavs = []
+            try:
+                for sub in substrings:
+
+                    res_lang = self.text_normalizer.tn(text=sub)[1]
+
+                    chunk = self.synthesize_base(
+                        prompt_semantic_tokens=prompt_semantic_tokens,
+                        prompt_acoustic_tokens=prompt_acoustic_tokens,
+                        spk_semantic_llm=spk_semantic_llm,
+                        spk_acoustic_llm=spk_acoustic_llm,
+                        prompt_text=prompt_text,
+                        text=sub,
+                        lang=res_lang,
+                    )
+
+                    out_wavs.append(chunk)
+                out_wav = torch.concat(out_wavs, axis=-1)
+                return out_wav
+            except:
+                print('[ERROR] ', format_exc())
+                return None
+        else:
+            out_wavs = []
+            try:
+                res_lang = self.text_normalizer.tn(text=text)[1]
+
+                chunk = self.synthesize_base(
+                    prompt_semantic_tokens=prompt_semantic_tokens,
+                    prompt_acoustic_tokens=prompt_acoustic_tokens,
+                    spk_semantic_llm=spk_semantic_llm,
+                    spk_acoustic_llm=spk_acoustic_llm,
+                    prompt_text=prompt_text,
+                    text=text,
+                    lang=res_lang,
+                )
+
+                out_wavs.append(chunk)
+                out_wav = torch.concat(out_wavs, axis=-1)
+                return out_wav
+            except:
+                return None

fireredtts/models/token2audio.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn.functional as F
+from fireredtts.modules.acoustic_llm import AcousticLLM
+from fireredtts.modules.acoustic_codec import AcousticCodec
+from fireredtts.modules.flowmatching import FlowToken2Mel
+from fireredtts.modules.bigvgan import BigVGAN, MelExtractor
+
+
+class TwoStageCodec(torch.nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.acoustic_llm = AcousticLLM(**config["acoustic_llm"])
+        self.acoustic_codec = AcousticCodec(**config["acoustic_codec"])
+    
+    def load_model(self, acoustic_llm_path, acoustic_codec_path):
+        self.acoustic_llm.load_state_dict(
+            torch.load(acoustic_llm_path, map_location="cpu"), strict=True
+        )
+        self.acoustic_codec.load_state_dict(
+            torch.load(acoustic_codec_path, map_location="cpu"), strict=True
+        )
+
+    @torch.inference_mode()
+    def forward(
+        self, semantic_token, prompt_semantic_token, prompt_acoustic_token, spk_gpt
+    ):
+        # print('Before: ', semantic_token.shape)
+        token_pred = torch.cat((prompt_semantic_token, semantic_token), dim=1)
+
+        # Fine LLM inference
+        token_pred = self.acoustic_llm.inference_speech(
+            speech_conditioning_latent=spk_gpt,
+            text_inputs=token_pred,
+            num_return_sequences=1,
+            input_tokens=prompt_acoustic_token,
+        )[0]
+
+        if isinstance(token_pred, (tuple, list)):
+            token_pred = [x.unsqueeze(0) for x in token_pred]
+        else:
+            token_pred = token_pred.unsqueeze(0)
+
+        acoustic_outputs = self.acoustic_codec.reconstruct_wav(token=token_pred)
+        wav = acoustic_outputs["wav_pred"].squeeze(1)
+
+        return wav
+
+    def extract(self, wavs, wav_lengths, spk):
+        if torch.cuda.is_available():
+            wavs = wavs.cuda()
+        cond_tok = self.acoustic_codec.extract_speech_tokens(wavs, wav_lengths)[
+            "token"
+        ][0]
+        spk_gpt = self.acoustic_llm.get_conditioning(spk)
+        return cond_tok, spk_gpt
+
+
+"""For FlowToken2Audio, keep interface consistant with TwoStageCodec to minimize code changes.
+prompt_acoustic_token alias to prompt_mel
+spk_gpt alias to spk_embeddings
+"""
+class FlowToken2Audio(torch.nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.flow = FlowToken2Mel(config['flow'])
+        self.bigvgan = BigVGAN(**config['bigvgan'])
+        self.mel_extractor = MelExtractor(**config['mel'])
+
+    def load_model(self, flow_path, bigvgan_path):
+        self.flow.load_state_dict(
+            torch.load(flow_path, map_location="cpu"), strict=True
+        )
+        self.bigvgan.load_state_dict(
+            torch.load(bigvgan_path, map_location="cpu")['generator'], strict=True
+        )
+        self.bigvgan.remove_weight_norm()
+
+    @torch.inference_mode()
+    def forward(
+        self, semantic_token, prompt_semantic_token, prompt_acoustic_token, spk_gpt
+    ):
+        # Align prompt token & prompt_mel
+        target_mel_length = prompt_semantic_token.shape[1] * 2
+        if target_mel_length > prompt_acoustic_token.shape[1]:
+            prompt_acoustic_token = F.pad(
+                prompt_acoustic_token, (0, 0, 0, target_mel_length-prompt_acoustic_token.shape[1]),
+                mode='constant', value=-11.5
+            )
+        elif target_mel_length < prompt_acoustic_token.shape[1]:
+            prompt_acoustic_token = prompt_acoustic_token[:, :target_mel_length]
+        # prompt_acoustic_token = F.interpolate(
+        #     prompt_acoustic_token.transpose(1, 2),
+        #     size=prompt_semantic_token.shape[1] * 2, mode='nearest'
+        # ).transpose(1, 2)
+        mel_pred = self.flow.inference(
+            prompt_token=prompt_semantic_token,
+            prompt_xvec=spk_gpt,
+            prompt_feat=prompt_acoustic_token,
+            token=semantic_token
+        )
+        wav = self.bigvgan(mel_pred.transpose(1, 2)).squeeze(1)
+        return wav
+
+    def extract(self, wavs, wav_lengths, spk):
+        mel = self.mel_extractor(wavs, 24000).transpose(1, 2)
+        if torch.cuda.is_available():
+            mel = mel.cuda()
+        return mel, spk.squeeze(0)

fireredtts/modules/acoustic_codec/__init__.py

@@ -0,0 +1 @@
+from .bigcodec import BigCodec as AcousticCodec

fireredtts/modules/acoustic_codec/alias_free_torch/__init__.py

@@ -0,0 +1,6 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+from .filter import *
+from .resample import *
+from .act import *

fireredtts/modules/acoustic_codec/alias_free_torch/act.py

@@ -0,0 +1,35 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from .resample import UpSample1d, DownSample1d
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+        causal: bool = False,
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+        self.causal = causal
+
+    # x: [B,C,T]
+    def forward(self, x):
+        if self.causal:
+            x = self.act(x)
+        else:
+            x = self.upsample(x)
+            x = self.act(x)
+            x = self.downsample(x)
+
+        return x

fireredtts/modules/acoustic_codec/alias_free_torch/filter.py

@@ -0,0 +1,99 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+if "sinc" in dir(torch):
+    sinc = torch.sinc
+else:
+    # This code is adopted from adefossez's julius.core.sinc under the MIT License
+    # https://adefossez.github.io/julius/julius/core.html
+    #   LICENSE is in incl_licenses directory.
+    def sinc(x: torch.Tensor):
+        """
+        Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+        __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+        """
+        return torch.where(
+            x == 0,
+            torch.tensor(1.0, device=x.device, dtype=x.dtype),
+            torch.sin(math.pi * x) / math.pi / x,
+        )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+#   LICENSE is in incl_licenses directory.
+def kaiser_sinc_filter1d(
+    cutoff, half_width, kernel_size
+):  # return filter [1,1,kernel_size]
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    # For kaiser window
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+        # Normalize filter to have sum = 1, otherwise we will have a small leakage
+        # of the constant component in the input signal.
+        filter_ /= filter_.sum()
+        filter = filter_.view(1, 1, kernel_size)
+
+    return filter
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        # kernel_size should be even number for stylegan3 setup,
+        # in this implementation, odd number is also possible.
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filter)
+
+    # input [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+        return out

fireredtts/modules/acoustic_codec/alias_free_torch/resample.py

@@ -0,0 +1,58 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from torch.nn import functional as F
+from .filter import LowPassFilter1d
+from .filter import kaiser_sinc_filter1d
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
+        self.register_buffer("filter", filter)
+
+    # x: [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(
+            x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+        )
+        x = x[..., self.pad_left : -self.pad_right]
+
+        return x
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
+
+    def forward(self, x):
+        xx = self.lowpass(x)
+
+        return xx

fireredtts/modules/acoustic_codec/bigcodec.py

@@ -0,0 +1,698 @@
+from einops import rearrange
+from torch import sin, pow
+from torch.nn import Parameter
+from torch.nn.utils import spectral_norm, weight_norm
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+import typing as tp
+import warnings
+
+from .alias_free_torch import *
+from .vector_quantization import VectorQuantization
+
+
+CONV_NORMALIZATIONS = frozenset(
+    [
+        "none",
+        "weight_norm",
+        "spectral_norm",
+        "time_layer_norm",
+        "layer_norm",
+        "time_group_norm",
+    ]
+)
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == "weight_norm":
+        return weight_norm(module)
+    elif norm == "spectral_norm":
+        return spectral_norm(module)
+    else:
+        return module
+
+
+def get_norm_module(
+    module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs
+) -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == "time_group_norm":
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(
+    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
+) -> int:
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(
+    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
+):
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+def pad1d(
+    x: torch.Tensor,
+    paddings: tp.Tuple[int, int],
+    mode: str = "zero",
+    value: float = 0.0,
+):
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == "reflect":
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left:end]
+
+
+class NormConv1d(nn.Module):
+
+    def __init__(
+        self,
+        *args,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+
+    def __init__(
+        self,
+        *args,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(
+            nn.ConvTranspose1d(*args, **kwargs), norm
+        )
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class SConv1d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        pad_mode: str = "reflect",
+        **kwargs,
+    ):
+        super().__init__()
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn(
+                "SConv1d has been initialized with stride > 1 and dilation > 1"
+                f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
+            )
+        self.conv = NormConv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            causal=causal,
+            norm=norm,
+            norm_kwargs=norm_kwargs,
+        )
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        B, C, T = x.shape
+        kernel_size = self.conv.conv.kernel_size[0]
+        stride = self.conv.conv.stride[0]
+        dilation = self.conv.conv.dilation[0]
+        kernel_size = (
+            kernel_size - 1
+        ) * dilation + 1  # effective kernel size with dilations
+        padding_total = kernel_size - stride
+        extra_padding = get_extra_padding_for_conv1d(
+            x, kernel_size, stride, padding_total
+        )
+        if self.causal:
+            # Left padding for causal
+            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(
+                x, (padding_left, padding_right + extra_padding), mode=self.pad_mode
+            )
+        return self.conv(x)
+
+
+class SConvTranspose1d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        causal: bool = False,
+        norm: str = "none",
+        trim_right_ratio: float = 1.0,
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            causal=causal,
+            norm=norm,
+            norm_kwargs=norm_kwargs,
+        )
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert (
+            self.causal or self.trim_right_ratio == 1.0
+        ), "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0.0 and self.trim_right_ratio <= 1.0
+
+    def forward(self, x):
+        kernel_size = self.convtr.convtr.kernel_size[0]
+        stride = self.convtr.convtr.stride[0]
+        padding_total = kernel_size - stride
+
+        y = self.convtr(x)
+
+        # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+        # removed at the very end, when keeping only the right length for the output,
+        # as removing it here would require also passing the length at the matching layer
+        # in the encoder.
+        if self.causal:
+            # Trim the padding on the right according to the specified ratio
+            # if trim_right_ratio = 1.0, trim everything from right
+            padding_right = math.ceil(padding_total * self.trim_right_ratio)
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        return y
+
+
+def WNConv1d(*args, **kwargs):
+    if kwargs.get("causal", False):
+        kwargs["norm"] = "weight_norm"
+        conv1d = SConv1d(*args, **kwargs)
+    else:
+        kwargs.pop("causal")
+        conv1d = weight_norm(nn.Conv1d(*args, **kwargs))
+    return conv1d
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    if kwargs.get("causal", False):
+        kwargs["norm"] = "weight_norm"
+        transposed_conv1d = SConvTranspose1d(*args, **kwargs)
+    else:
+        kwargs.pop("causal")
+        transposed_conv1d = weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+    return transposed_conv1d
+
+
+class SnakeBeta(nn.Module):
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        super(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+
+class ResidualUnit(nn.Module):
+
+    def __init__(self, dim: int = 16, dilation: int = 1, causal: bool = False):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True), causal=causal),
+            WNConv1d(
+                dim, dim, kernel_size=7, dilation=dilation, padding=pad, causal=causal
+            ),
+            Activation1d(activation=SnakeBeta(dim, alpha_logscale=True), causal=causal),
+            WNConv1d(dim, dim, kernel_size=1, causal=causal),
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+
+class EncoderBlock(nn.Module):
+
+    def __init__(
+        self, dim: int = 16, stride: int = 1, dilations=(1, 3, 9), causal: bool = False
+    ):
+        super().__init__()
+        runits = [ResidualUnit(dim // 2, dilation=d, causal=causal) for d in dilations]
+        self.block = nn.Sequential(
+            *runits,
+            Activation1d(
+                activation=SnakeBeta(dim // 2, alpha_logscale=True), causal=causal
+            ),
+            WNConv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=stride // 2 + stride % 2,
+                causal=causal,
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class DecoderBlock(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int = 16,
+        output_dim: int = 8,
+        stride: int = 1,
+        dilations=(1, 3, 9),
+        causal: bool = False,
+    ):
+        super().__init__()
+        self.block = nn.Sequential(
+            Activation1d(
+                activation=SnakeBeta(input_dim, alpha_logscale=True), causal=causal
+            ),
+            WNConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=stride // 2 + stride % 2,
+                output_padding=stride % 2,
+                causal=causal,
+            ),
+        )
+        self.block.extend(
+            [ResidualUnit(output_dim, dilation=d, causal=causal) for d in dilations]
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class ResLSTM(nn.Module):
+
+    def __init__(
+        self,
+        dimension: int,
+        num_layers: int = 2,
+        bidirectional: bool = False,
+        skip: bool = True,
+    ):
+        super().__init__()
+        self.skip = skip
+        self.lstm = nn.LSTM(
+            dimension,
+            dimension if not bidirectional else dimension // 2,
+            num_layers,
+            batch_first=True,
+            bidirectional=bidirectional,
+        )
+
+    def forward(self, x):
+        x = rearrange(x, "b f t -> b t f")
+        y, _ = self.lstm(x)
+        if self.skip:
+            y = y + x
+        y = rearrange(y, "b t f -> b f t")
+        return y
+
+
+class Resampler(nn.Module):
+
+    def __init__(self, source_sr=24000, target_sr=24000):
+        super().__init__()
+        self.source_sr = source_sr
+        self.target_sr = target_sr
+
+    def forward(self, wav, wav_length):
+        if self.source_sr != self.target_sr:
+            wav = torchaudio.functional.resample(wav, self.source_sr, self.target_sr)
+            wav_length = (wav_length * (self.source_sr / self.target_sr)).int()
+        return wav, wav_length
+
+
+class CodecEncoder(nn.Module):
+
+    def __init__(
+        self,
+        ngf=48,
+        use_rnn=True,
+        rnn_bidirectional=False,
+        rnn_num_layers=2,
+        up_ratios=(2, 2, 2, 5, 5),
+        dilations=(1, 3, 9),
+        out_channels=1024,
+        causal=False,
+    ):
+        super().__init__()
+        self.hop_length = np.prod(up_ratios)
+        self.ngf = ngf
+        self.up_ratios = up_ratios
+
+        # Create first convolution
+        d_model = ngf
+        self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3, causal=causal)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for i, stride in enumerate(up_ratios):
+            d_model *= 2
+            self.block += [
+                EncoderBlock(d_model, stride=stride, dilations=dilations, causal=causal)
+            ]
+        # RNN
+        if use_rnn:
+            self.block += [
+                ResLSTM(
+                    d_model, num_layers=rnn_num_layers, bidirectional=rnn_bidirectional
+                )
+            ]
+        # Create last convolution
+        self.block += [
+            Activation1d(
+                activation=SnakeBeta(d_model, alpha_logscale=True), causal=causal
+            ),
+            WNConv1d(d_model, out_channels, kernel_size=3, padding=1, causal=causal),
+        ]
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        out = self.block(x)
+        return out
+
+    def remove_weight_norm(self):
+        def _remove_weight_norm(m):
+            try:
+                torch.nn.utils.remove_weight_norm(m)
+            except ValueError:  # this module didn't have weight norm
+                return
+
+        self.apply(_remove_weight_norm)
+
+    def apply_weight_norm(self):
+        def _apply_weight_norm(m):
+            if isinstance(m, nn.Conv1d):
+                torch.nn.utils.weight_norm(m)
+
+        self.apply(_apply_weight_norm)
+
+    def reset_parameters(self):
+        self.apply(init_weights)
+
+
+class CodecDecoder(nn.Module):
+
+    def __init__(
+        self,
+        in_channels=1024,
+        upsample_initial_channel=1536,
+        ngf=48,
+        use_rnn=True,
+        rnn_bidirectional=False,
+        rnn_num_layers=2,
+        up_ratios=(5, 5, 2, 2, 2),
+        dilations=(1, 3, 9),
+        causal=False,
+        delay=0,
+    ):
+        super().__init__()
+        self.hop_length = np.prod(up_ratios)
+        self.ngf = ngf
+        self.up_ratios = up_ratios
+        self.delay = delay
+
+        channels = upsample_initial_channel
+        layers = [
+            WNConv1d(in_channels, channels, kernel_size=7, padding=3, causal=causal)
+        ]
+
+        if use_rnn:
+            layers += [
+                ResLSTM(
+                    channels, num_layers=rnn_num_layers, bidirectional=rnn_bidirectional
+                )
+            ]
+
+        for i, stride in enumerate(up_ratios):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            layers += [
+                DecoderBlock(input_dim, output_dim, stride, dilations, causal=causal)
+            ]
+
+        layers += [
+            Activation1d(
+                activation=SnakeBeta(output_dim, alpha_logscale=True), causal=causal
+            ),
+            WNConv1d(output_dim, 1, kernel_size=7, padding=3, causal=causal),
+            nn.Tanh(),
+        ]
+
+        self.model = nn.Sequential(*layers)
+        self.reset_parameters()
+
+    def forward(self, x):
+        # Time delay
+        if self.delay > 0:
+            x = F.pad(x, (0, self.delay), mode="constant", value=0)
+
+        x = self.model(x)
+
+        # De-delay
+        if self.delay > 0:
+            x = x[..., self.delay :]
+
+        return x
+
+    def remove_weight_norm(self):
+        def _remove_weight_norm(m):
+            try:
+                torch.nn.utils.remove_weight_norm(m)
+            except ValueError:  # this module didn't have weight norm
+                return
+
+        self.apply(_remove_weight_norm)
+
+    def apply_weight_norm(self):
+        def _apply_weight_norm(m):
+            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
+                torch.nn.utils.weight_norm(m)
+
+        self.apply(_apply_weight_norm)
+
+    def reset_parameters(self):
+        self.apply(init_weights)
+
+
+class BigCodec(nn.Module):
+
+    def __init__(
+        self,
+        n_model_size: int,
+        encoder_config: dict,
+        decoder_config: dict,
+        vq_config: dict,
+        resampler_config: dict = None,
+    ):
+        super(BigCodec, self).__init__()
+        self.n_model_size = n_model_size
+
+        self.encoder = CodecEncoder(out_channels=n_model_size, **encoder_config)
+        self.decoder = CodecDecoder(in_channels=n_model_size, **decoder_config)
+        self.quantizer = VectorQuantization(n_model_size, **vq_config)
+
+        # Optional modules
+        if resampler_config:
+            self.resampler = Resampler(**resampler_config)
+
+    def forward(
+        self, wav, wav_length=None, enable_vq=True, decode=True, update_codebook=True
+    ):
+        # Preprocess wav
+        if len(wav.shape) == 2:
+            wav = wav.unsqueeze(1)
+        if wav_length is None:
+            wav_length = torch.full([wav.shape[0]], max(wav.shape)).to(wav.device)
+
+        # (Optional) Resample
+        processed_wav, processed_wav_length = wav, wav_length
+        if hasattr(self, "resampler"):
+            processed_wav, processed_wav_length = self.resampler(
+                processed_wav, processed_wav_length
+            )
+
+        # Update VQ parameters
+        quant_length = torch.ceil(processed_wav_length / self.encoder.hop_length).int()
+        update_codebook = update_codebook and self.training
+
+        # Encode
+        encoder_outputs = self.encoder(processed_wav)
+
+        # Quantize
+        quant, diff, embed_ind = self.quantizer(
+            encoder_outputs.transpose(1, 2),
+            quant_length.clamp(max=encoder_outputs.shape[2]),
+            enable_vq=enable_vq,
+            update_codebook=update_codebook,
+        )
+
+        if decode:
+            # Decode
+            decoder_outputs = self.decoder(quant.transpose(1, 2))
+        else:
+            decoder_outputs = None
+
+        output_dict = {
+            "quant": quant,
+            "token": embed_ind,
+            "token_length": quant_length,
+            "encoder_diffs": diff,
+            "wav_pred": decoder_outputs,
+        }
+        return output_dict
+
+    @torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
+    def extract_speech_tokens(
+        self, wav, wav_length, serialize=True, extract_spk=True, shuffle=False
+    ):
+        output_dict = self.forward(wav, wav_length, enable_vq=True, decode=False)
+        token_seqs, token_length = [output_dict["token"]], [output_dict["token_length"]]
+        output_dict.update(
+            {
+                "token": token_seqs,
+                "token_length": token_length,
+            }
+        )
+
+        return output_dict
+
+    @torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
+    def reconstruct_wav(self, token=None, quant=None, spk=None):
+        if token is not None:
+            # De-tokenization
+            quant = self.quantizer.decode(token)
+
+            # Speaker embedding
+            if hasattr(self, "global_encoder"):
+                quant = quant + spk.unsqueeze(2)
+        else:
+            assert quant is not None
+
+        # Decode
+        wav_pred = self.decoder(quant)
+
+        return {
+            "wav_pred": wav_pred,
+        }

fireredtts/modules/acoustic_codec/vector_quantization.py

@@ -0,0 +1,580 @@
+from einops import rearrange, repeat
+from torch import nn
+from typing import Union
+
+import torch
+import torch.nn.functional as F
+import typing as tp
+import numpy as np
+import warnings
+
+
+def default(val: tp.Any, d: tp.Any) -> tp.Any:
+    return val if val is not None else d
+
+
+def flatten(x, x_len):
+    x_f = x.view(-1, *x.shape[2:])
+    return x_f
+
+
+def ema_inplace(moving_avg, new, decay):
+    if isinstance(decay, torch.Tensor):
+        moving_avg.data.mul_(decay).add_(new * (1 - decay))
+    else:
+        moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
+    return (x + epsilon) / (x.sum() + n_categories * epsilon)
+
+
+def uniform_init(*shape: int):
+    t = torch.empty(shape)
+    nn.init.kaiming_uniform_(t)
+    return t
+
+
+def sample_vectors(samples, num: int):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+class EuclideanCodebook(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        decay: float = 0.99,
+        epsilon: float = 1e-5,
+        threshold_ema_dead_code: float = 1.0,
+        n_cache_iters: int = 1,
+    ):
+        super().__init__()
+        self.decay = decay
+        init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = uniform_init
+        embed = init_fn(codebook_size, dim)
+
+        self.codebook_size = codebook_size
+
+        self.epsilon = epsilon
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.update_iter = 0
+
+        self.n_cache_iters = n_cache_iters
+        self.cache_vectors = []
+        self.cache_indices = []
+
+        if isinstance(self.decay, (tuple, list)):
+            self.embed_avg_cache = []
+            self.register_buffer("diff_avg_long", torch.zeros(codebook_size) + 1e-5)
+            self.register_buffer("diff_avg_short", torch.zeros(codebook_size) + 1e-5)
+        self.register_buffer("inited", torch.Tensor([True]))
+        self.register_buffer("cluster_size", torch.zeros(codebook_size))
+        self.register_buffer("embed", embed)
+        self.register_buffer("embed_avg", embed.clone())
+
+    @torch.jit.ignore
+    def init_embed_(self, data):
+        if self.inited:
+            return
+
+    def replace_(self, samples, mask, dists=None):
+        reset_cluster_size = min(
+            self.threshold_ema_dead_code + 1, self.threshold_ema_dead_code * 1.1
+        )
+
+        modified_codebook = torch.where(
+            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
+        )
+        modified_codebook_avg = torch.where(
+            mask[..., None], modified_codebook * reset_cluster_size, self.embed_avg
+        )
+        modified_cluster_size = torch.where(
+            mask,
+            torch.full_like(self.cluster_size, reset_cluster_size),
+            self.cluster_size,
+        )
+
+        self.embed.data.copy_(modified_codebook)
+        self.embed_avg.data.copy_(modified_codebook_avg)
+        self.cluster_size.data.copy_(modified_cluster_size)
+
+    def expire_codes_(self, batch_samples, dists=None):
+        self.update_iter += 1
+        if self.threshold_ema_dead_code == 0:
+            return
+        elif self.threshold_ema_dead_code < 1:
+            threshold_ema_dead_code = (
+                sum(self.cluster_size) * self.threshold_ema_dead_code
+            )
+        else:
+            threshold_ema_dead_code = self.threshold_ema_dead_code
+
+        expired_codes = self.cluster_size < threshold_ema_dead_code
+        if not torch.any(expired_codes):
+            return
+
+        batch_samples = rearrange(batch_samples, "... d -> (...) d")
+        self.replace_(batch_samples, mask=expired_codes, dists=dists)
+
+    def preprocess(self, x):
+        x = rearrange(x, "... d -> (...) d")
+        return x
+
+    def quantize(self, x):
+        embed = self.embed.t()
+        dist = -(
+            x.pow(2).sum(1, keepdim=True)
+            - 2 * x @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+        embed_ind = dist.max(dim=-1).indices
+
+        return embed_ind, dist
+
+    def postprocess_emb(self, embed_ind, shape):
+        return embed_ind.view(*shape[:-1])
+
+    def dequantize(self, embed_ind):
+        quantize = F.embedding(embed_ind, self.embed)
+        return quantize
+
+    def encode(self, x):
+        shape = x.shape
+        # pre-process
+        x = self.preprocess(x)
+        # quantize
+        embed_ind, dist = self.quantize(x)
+        # post-process
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+
+        return embed_ind, dist
+
+    def decode(self, embed_ind):
+        quantize = self.dequantize(embed_ind)
+        return quantize
+
+    def forward(self, x, x_len, enable_vq=True, update_codebook=True, masking=False):
+        x_org, shape, dtype = x, x.shape, x.dtype
+
+        x = self.preprocess(x)
+
+        embed_ind, dist = self.quantize(x)
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        dist = dist.view(shape[0], shape[1], dist.shape[-1])
+
+        quantize = self.dequantize(embed_ind)
+
+        if self.training and update_codebook:
+            if enable_vq:
+                quantize = x_org + (quantize - x_org).detach()
+            else:
+                quantize = x_org
+
+            # Get flatten embedding indices and distances
+            if masking:
+                x_f = torch.cat(
+                    [e[: int(e_len)] for e, e_len in zip(x_org, x_len)], dim=0
+                )
+                embed_ind_f = torch.cat(
+                    [e[: int(e_len)] for e, e_len in zip(embed_ind, x_len)], dim=0
+                )
+                dist_f = torch.cat(
+                    [e[: int(e_len)] for e, e_len in zip(dist, x_len)], dim=0
+                )
+                q_f = torch.cat(
+                    [e[: int(e_len)] for e, e_len in zip(quantize.detach(), x_len)],
+                    dim=0,
+                )
+                commit_loss = F.mse_loss(q_f, x_f)
+            else:
+                x_f = x_org.view(-1, x_org.shape[-1]).contiguous()
+                embed_ind_f = embed_ind.view(-1).contiguous()
+                dist_f = dist.view(-1).contiguous()
+                commit_loss = F.mse_loss(quantize.detach(), x_org)
+            self.init_embed_(x_f)
+
+            # We do the expiry of code at that point as buffers are in sync
+            # and all the workers will take the same decision.
+            self.expire_codes_(x_f, dist_f)
+
+            # Calculate codebook statistics
+            embed_onehot = F.one_hot(embed_ind_f, self.codebook_size).type(dtype)
+            embed_onehot_sum = embed_onehot.sum(0)
+            embed_sum = x_f.t() @ embed_onehot
+
+            # EMA updating
+            ema_inplace(self.cluster_size, embed_onehot_sum, self.decay)
+            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
+
+            cluster_size = (
+                laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)
+                * self.cluster_size.sum()
+            )
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embed.data.copy_(embed_normalized)
+        else:
+            commit_loss = torch.tensor(
+                0.0, device=quantize.device, requires_grad=self.training
+            )
+
+        return quantize, commit_loss, embed_ind
+
+
+class MultiHeadEuclideanCodebook(nn.Module):
+
+    def __init__(
+        self,
+        dim: Union[int, list],
+        codebook_size: list,
+        n_groups: int = 1,
+        dropout_rate_per_group: float = 0,
+        ordered: bool = False,
+        ordered_axis: str = "sequence",
+        method: str = "product",
+        **kwargs,
+    ):
+        super().__init__()
+        self.codebook_sizes = codebook_size
+        self.codebook_dims = dim
+        self.n_groups = n_groups
+        self.n_heads_per_group = len(codebook_size) // n_groups
+        self.dropout_rate_per_group = dropout_rate_per_group
+        self.ordered = ordered
+        self.ordered_axis = ordered_axis
+        self.method = method
+        assert len(codebook_size) % n_groups == 0
+
+        self.codebooks = nn.ModuleList()
+        dim = self.codebook_dims
+        for i, size in enumerate(self.codebook_sizes):
+            if isinstance(self.codebook_dims, list):
+                dim = (
+                    self.codebook_dims[i]
+                    if method == "product"
+                    else sum(self.codebook_dims)
+                )
+            self.codebooks.append(EuclideanCodebook(dim, size, **kwargs))
+
+    def decode(self, embed_ind):
+        if self.n_groups == 1 or len(embed_ind.shape) == 2:
+            embed_ind = embed_ind.unsqueeze(-1)
+
+        actual_n_groups = embed_ind.shape[-1]
+        if actual_n_groups < self.n_groups:
+            print(
+                f"The actual number of heads ({actual_n_groups}) is smaller than the pre-designed ({self.n_groups})!"
+            )
+            embed_ind = F.pad(
+                embed_ind, (0, self.n_groups - actual_n_groups), "replicate"
+            )
+        # assert embed_ind.shape[-1] == self.n_groups
+
+        index_heads, codebook_heads, scale_heads = zip(
+            *[
+                (
+                    embed_ind[..., i // self.n_heads_per_group],
+                    self.codebooks[i : i + self.n_heads_per_group],
+                    self.codebook_sizes[i : i + self.n_heads_per_group],
+                )
+                for i in range(0, len(self.codebook_sizes), self.n_heads_per_group)
+            ]
+        )
+
+        quantize_heads, quantize_groups = [], []
+        for i in range(self.n_groups):
+            embed_ind, codebooks, scales = (
+                index_heads[i],
+                codebook_heads[i],
+                scale_heads[i],
+            )
+
+            inv_scales = list(torch.tensor([1] + scales[:-1]).cumprod(dim=0))[::-1]
+            inv_quantizes = []
+            for codebook, scale in zip(codebooks[::-1], inv_scales):
+                index, embed_ind = embed_ind // scale, embed_ind % scale
+                quantize = codebook.dequantize(index)
+                inv_quantizes.append(quantize)
+            quantizes = inv_quantizes[::-1]
+            group_embeddings = torch.cat(quantizes, dim=-1)
+            quantize_groups.append(group_embeddings)
+            quantize_heads += quantizes
+
+        if self.method == "product":
+            if actual_n_groups < self.n_groups:
+                for i in range(actual_n_groups, self.n_groups):
+                    quantize_groups[i].zero_()
+            quantize = torch.cat(quantize_groups, dim=-1)
+        elif self.method == "residual":
+            quantize = sum(quantize_heads)
+
+        return quantize
+
+    def forward(self, x, x_len, enable_vq=True, update_codebook=True):
+        # Pre-process
+        x = self._preprocess(x)
+
+        # Quantize
+        quants, losses, indices = self._quantize(
+            x, x_len, enable_vq=enable_vq, update_codebook=update_codebook
+        )
+
+        # Integrate
+        quant, loss, index = self._integrate(
+            quants, losses, indices, update_codebook=update_codebook
+        )
+
+        return quant, loss, index
+
+    def _preprocess(self, x):
+        if self.method == "product" and isinstance(self.codebook_dims, (list, tuple)):
+            x = torch.split(x, self.codebook_dims, dim=-1)
+        return x
+
+    def _quantize(self, x, x_len, enable_vq, update_codebook):
+        if self.method == "product":
+            quants, losses, indices = zip(
+                *[
+                    codebook(
+                        chunk,
+                        x_len,
+                        enable_vq=enable_vq,
+                        update_codebook=update_codebook,
+                    )
+                    for chunk, codebook in zip(x, self.codebooks)
+                ]
+            )
+        elif self.method == "residual":
+            quants, losses, indices = [], [], []
+            residual = x
+            for codebook in self.codebooks:
+                quant, loss, index = codebook(
+                    residual,
+                    x_len,
+                    enable_vq=enable_vq,
+                    update_codebook=update_codebook,
+                )
+                residual = residual - quant
+                quants.append(quant)
+                losses.append(loss)
+                indices.append(index)
+
+        return quants, losses, indices
+
+    def _integrate(self, quants, losses, indices, update_codebook=True):
+        (B, T, D), M = quants[0].shape, len(quants)
+        device = quants[0].device
+
+        # Average loss
+        loss = sum(losses) / len(losses)
+
+        # Get indices
+        if self.n_groups == 1:
+            scale = (
+                torch.tensor([1] + self.codebook_sizes[:-1]).cumprod(dim=0).to(device)
+            )
+            index = (torch.stack(indices, dim=-1) * scale).sum(dim=-1)
+        else:
+            index_heads, scale_heads = zip(
+                *[
+                    (
+                        indices[i : i + self.n_heads_per_group],
+                        torch.tensor(
+                            [1]
+                            + self.codebook_sizes[i : i + self.n_heads_per_group - 1]
+                        )
+                        .cumprod(dim=0)
+                        .to(device),
+                    )
+                    for i in range(0, len(quants), self.n_heads_per_group)
+                ]
+            )
+            index = torch.stack(
+                [
+                    (torch.stack(x, dim=-1) * s).sum(dim=-1)
+                    for x, s in zip(index_heads, scale_heads)
+                ],
+                dim=-1,
+            )
+
+        # Add dropout
+        quant_groups = self._dropout(quants, enabled=update_codebook)
+
+        # Combine quantized features
+        if self.method == "product":
+            quant = torch.cat(quant_groups, dim=-1)
+        elif self.method == "residual":
+            quant = torch.cat(quant_groups, dim=-1).view(B, T, M, D).sum(dim=2)
+
+        return quant, loss, index
+
+    def _dropout(self, quants, enabled=True):
+        if enabled and self.training and self.ordered:
+            if self.dropout_rate_per_group == 0:
+                threshold = [
+                    (i // self.n_heads_per_group * 1.0 / self.n_groups)
+                    for i in range(0, len(quants), self.n_heads_per_group)
+                ]
+            elif self.dropout_rate_per_group == "exp":
+                x = [np.exp(4 * i / self.n_groups) for i in range(self.n_groups)]
+                x = np.asarray(x) / sum(x)
+                threshold = np.cumsum(np.asarray([0] + x))
+            else:
+                x = np.asarray(self.dropout_rate_per_group) / sum(
+                    self.dropout_rate_per_group
+                )
+                threshold = np.cumsum(np.asarray([0] + x))
+
+            if self.ordered_axis == "sequence":
+                rate = torch.rand((quants[0].shape[0], 1, 1), device=quants[0].device)
+            elif self.ordered_axis == "frame":
+                rate = torch.rand(
+                    (quants[0].shape[0], quants[0].shape[1], 1), device=quants[0].device
+                )
+
+            quant_groups = []
+            for i in range(0, len(quants), self.n_heads_per_group):
+                quant_group = torch.cat(quants[i : i + self.n_heads_per_group], dim=-1)
+                is_kept = threshold[i // self.n_heads_per_group] <= rate
+                quant_group = torch.where(
+                    is_kept, quant_group, torch.zeros_like(quant_group)
+                )
+                quant_groups.append(quant_group)
+        elif self.ordered:
+            quant_groups = []
+            for i in range(0, len(quants), self.n_heads_per_group):
+                quant_group = torch.cat(quants[i : i + self.n_heads_per_group], dim=-1)
+                quant_groups.append(quant_group)
+        else:
+            quant_groups = quants
+
+        return quant_groups
+
+
+class VectorQuantization(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: Union[int, list],
+        codebook_dim: Union[int, list] = None,
+        decay: float = 0.99,
+        epsilon: float = 1e-5,
+        threshold_ema_dead_code: float = 1.0,
+        commitment_weight: float = 1.0,
+        requires_projection: bool = False,
+        norm: str = "none",
+        **kwargs,
+    ):
+        super().__init__()
+        _codebook_dim: Union[int, list] = default(codebook_dim, dim)
+
+        requires_projection = _codebook_dim != dim or requires_projection
+        proj_dim = (
+            sum(_codebook_dim) if isinstance(_codebook_dim, list) else _codebook_dim
+        )
+        if requires_projection:
+            self.project_in = nn.Linear(dim, proj_dim)
+            self.project_out = nn.Linear(proj_dim, dim)
+            if norm == "weight_norm":
+                self.project_in = torch.nn.utils.weight_norm(self.project_in)
+                self.project_out = torch.nn.utils.weight_norm(self.project_out)
+        else:
+            self.norm = None
+            self.project_in = nn.Identity()
+            self.project_out = nn.Identity()
+
+        self.epsilon = epsilon
+        self.commitment_weight = commitment_weight
+        self.codebook_size = codebook_size
+
+        codebook_class = (
+            EuclideanCodebook
+            if isinstance(codebook_size, int)
+            else MultiHeadEuclideanCodebook
+        )
+        self._codebook = codebook_class(
+            dim=_codebook_dim,
+            codebook_size=codebook_size,
+            decay=decay,
+            epsilon=epsilon,
+            threshold_ema_dead_code=threshold_ema_dead_code,
+            **kwargs,
+        )
+        self.codebook_size = codebook_size
+
+    @property
+    def codebook(self):
+        return self._codebook.embed
+
+    def encode(self, x, x_len=None):
+        x = rearrange(x, "b d n -> b n d")
+        x = self.project_in(x)
+        embed_in = self._codebook.encode(x)
+        return embed_in
+
+    def decode(self, embed_ind, embed_len=None):
+        quantize = self._codebook.decode(embed_ind)
+        quantize = self.project_out(quantize)
+        quantize = rearrange(quantize, "b n d -> b d n")
+        return quantize
+
+    def decode_latent(self, latent, latent_len=None):
+        if latent_len is None:
+            latent_len = (
+                torch.Tensor([latent.shape[1]] * latent.shape[0])
+                .to(latent.device)
+                .int()
+            )
+
+        quantize, _, _ = self._codebook(latent, latent_len)
+        quantize = self.project_out(quantize)
+        return quantize
+
+    @torch.cuda.amp.autocast(dtype=torch.float32)
+    def forward(
+        self,
+        x,
+        x_len,
+        enable_vq=True,
+        update_codebook=True,
+        return_pre_quant=False,
+        return_dict=False,
+    ):
+        device = x.device
+
+        x = self.project_in(x)
+
+        quantize, commit_loss, embed_ind = self._codebook(
+            x, x_len, enable_vq=enable_vq, update_codebook=update_codebook
+        )
+        if self.training and update_codebook:
+            loss = torch.tensor(0.0, device=device, requires_grad=True)
+            if self.commitment_weight > 0:
+                loss = loss + commit_loss * self.commitment_weight
+        else:
+            loss = torch.tensor(0.0, device=device, requires_grad=False)
+
+        embed = quantize
+        quantize = self.project_out(quantize)
+
+        if return_dict:
+            return {
+                "quantize": quantize,
+                "loss": loss,
+                "embed": embed,
+                "embed_ind": embed_ind,
+            }
+        elif return_pre_quant:
+            pre_quantize = self.project_out(x)
+            return (pre_quantize, quantize), loss, embed_ind
+        else:
+            return quantize, loss, embed_ind

fireredtts/modules/acoustic_llm/__init__.py

@@ -0,0 +1 @@
+from .acoustic_llm import AcousticLLM

fireredtts/modules/acoustic_llm/acoustic_llm.py

@@ -0,0 +1,876 @@
+from einops import rearrange
+from time import time
+from torch.utils.checkpoint import checkpoint
+from transformers import (
+    GPT2Config,
+    GPT2Model,
+    GPT2PreTrainedModel,
+    LogitsProcessorList,
+    LogitsWarper,
+    StoppingCriteria,
+    StoppingCriteriaList,
+)
+from transformers.generation.streamers import BaseStreamer
+from transformers.generation.utils import (
+    GenerationConfig,
+    GenerateDecoderOnlyOutput,
+    GenerateEncoderDecoderOutput,
+    GenerateNonBeamOutput,
+)
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from transformers.utils.model_parallel_utils import get_device_map, assert_device_map
+from typing import Any, Dict, Optional, Tuple, Union
+
+import functools
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MultiHeadRepetitionPenaltyLogitsProcessor(LogitsWarper):
+
+    def __init__(
+        self, penalty: float = 2.0, n_heads: int = 4, n_frames: int = -1, start_index=0
+    ):
+        if not isinstance(penalty, float) or not (penalty > 0):
+            raise ValueError(
+                f"`penalty` has to be a strictly positive float, but is {penalty}"
+            )
+
+        self.penalty = penalty
+        self.n_heads = n_heads
+        self.n_frames = n_frames
+        self.start_index = start_index
+
+    def __call__(
+        self, input_ids: torch.LongTensor, scores: torch.FloatTensor
+    ) -> torch.FloatTensor:
+        input_ids = input_ids[:, self.start_index :]
+        if input_ids.size(1) == 0:
+            return scores
+
+        if self.n_frames <= 0:
+            input_ids = torch.flip(input_ids, [1])[:, self.n_heads - 1 :: self.n_heads]
+        else:
+            input_ids = torch.flip(input_ids, [1])[
+                :, self.n_heads - 1 : self.n_heads * self.n_frames : self.n_heads
+            ]
+        score = torch.gather(scores, 1, input_ids)
+
+        # if score < 0 then repetition penalty has to be multiplied to reduce the token probabilities
+        if self.penalty > 100:
+            score = torch.full_like(score, -1e3)
+        else:
+            score = torch.where(score < 0, score * self.penalty, score / self.penalty)
+
+        scores.scatter_(1, input_ids, score)
+        return scores
+
+
+def null_position_embeddings(range, dim):
+    return torch.zeros((range.shape[0], range.shape[1], dim), device=range.device)
+
+
+class FixedStoppingCriteria(StoppingCriteria):
+
+    def __init__(self, running_steps, start_index=0):
+        self.running_steps = running_steps
+        self.start_index = start_index
+
+    def __call__(
+        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
+    ) -> torch.BoolTensor:
+        assert input_ids.shape[0] == 1, input_ids.shape
+        if input_ids.shape[1] - self.start_index >= self.running_steps:
+            return torch.tensor([True]).to(input_ids.device)
+        return torch.tensor([False]).to(input_ids.device)
+
+
+class DelayStoppingCriteria(StoppingCriteria):
+
+    def __init__(self, eos_token_id, delay_steps):
+        self.delay_steps = delay_steps
+        self.eos_token_id = torch.tensor(eos_token_id)
+
+    def __call__(
+        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
+    ) -> torch.BoolTensor:
+        assert input_ids.shape[0] == 1, input_ids.shape
+
+        if (input_ids == self.eos_token_id).any():
+            index = (input_ids[0] == self.eos_token_id).nonzero(as_tuple=True)[0][0]
+            if index + self.delay_steps < input_ids.shape[1]:
+                return torch.tensor([True]).to(input_ids.device)
+        return torch.tensor([False]).to(input_ids.device)
+
+
+class SuppressionLogitsProcessor(LogitsWarper):
+
+    def __init__(self, suppressed_ids=[]):
+        self.suppressed_ids = suppressed_ids
+
+    def __call__(
+        self, input_ids: torch.LongTensor, scores: torch.FloatTensor
+    ) -> torch.FloatTensor:
+        for sid in self.suppressed_ids:
+            scores[..., sid] = scores.min()
+        return scores
+
+
+class MHGPT2InferenceModel(GPT2PreTrainedModel):
+
+    def __init__(
+        self, config, gpt, text_pos_emb, embeddings, norm, linear, kv_cache=True
+    ):
+        super().__init__(config)
+        self.transformer = gpt
+        self.text_pos_embedding = text_pos_emb
+        self.embeddings = embeddings
+        self.lm_head = nn.ModuleList([norm, linear])  # nn.Sequential(norm, linear)
+        self.kv_cache = kv_cache
+
+        # Multi-head configuration
+        self.n_heads = len(linear)
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+        self.cached_mel_emb = None
+        self.cached_mel_parallel_emb = None
+
+    def store_mel_emb(self, mel_emb):
+        self.cached_mel_emb = mel_emb
+
+    def store_mel_parallel_emb(self, mel_emb):
+        self.cached_mel_parallel_emb = mel_emb
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)  # usually None
+        if not self.kv_cache:
+            past_key_values = None
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        assert self.cached_mel_emb is not None
+        assert inputs_embeds is None  # Not supported by this inference model.
+        assert labels is None  # Training not supported by this inference model.
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        # Create embedding
+        mel_len = self.cached_mel_emb.shape[1]
+        attention_mask = None
+        position_ids = None
+
+        if input_ids.shape[1] != 1 and past_key_values is None:
+            text_inputs = input_ids[:, mel_len:]
+            text_emb = sum(
+                [self.embeddings[i](text_inputs[:, :, i]) for i in range(self.n_heads)]
+            )
+            text_emb = text_emb + self.text_pos_embedding(text_emb)
+
+            if self.cached_mel_emb.shape[0] != text_emb.shape[0]:
+                mel_emb = self.cached_mel_emb.repeat_interleave(
+                    text_emb.shape[0] // self.cached_mel_emb.shape[0], 0
+                )
+            else:  # this outcome only occurs once per loop in most cases
+                mel_emb = self.cached_mel_emb
+
+            if self.cached_mel_parallel_emb is not None:
+                text_emb = (
+                    text_emb + self.cached_mel_parallel_emb[:, : text_emb.shape[1]]
+                )
+
+            emb = torch.cat([mel_emb, text_emb], dim=1)
+        else:  # KV-cache mode
+            text_inputs = input_ids[:, mel_len:]
+            emb = sum(
+                [self.embeddings[i](text_inputs[:, -1, i]) for i in range(self.n_heads)]
+            )
+            emb = emb + self.text_pos_embedding.get_fixed_embedding(
+                text_inputs.shape[1] - 1, emb.device
+            )
+
+            if self.cached_mel_parallel_emb is not None:
+                emb = emb + self.cached_mel_parallel_emb[:, text_inputs.shape[1] - 1]
+
+        transformer_outputs = self.transformer(
+            inputs_embeds=emb,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=True,  # output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        past_key_values = transformer_outputs.past_key_values
+        output_hidden_states = transformer_outputs.hidden_states
+        output_attentions = transformer_outputs.attentions
+
+        # Set device for model parallelism
+        if self.model_parallel:
+            if torch.backends.mps.is_available():
+                self.to(self.transformer.first_device)
+            else:
+                torch.cuda.set_device(self.transformer.first_device)
+            hidden_states = hidden_states.to(self.lm_head.weight.device)
+
+        lm_logits = self.lm_head[0](hidden_states)
+        lm_logits = [head(lm_logits) for head in self.lm_head[1]]
+        lm_logits = torch.stack(lm_logits, dim=2)
+
+        if not return_dict:
+            return (lm_logits,) + transformer_outputs[1:]
+
+        output = CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=lm_logits,
+            past_key_values=past_key_values,
+            hidden_states=output_hidden_states,
+            attentions=output_attentions,
+        )
+        return output
+
+    def _sample(
+        self,
+        input_ids: torch.LongTensor,
+        logits_processor: LogitsProcessorList,
+        stopping_criteria: StoppingCriteriaList,
+        generation_config: GenerationConfig,
+        synced_gpus: bool,
+        streamer: Optional["BaseStreamer"],
+        **model_kwargs,
+    ) -> Union[GenerateNonBeamOutput, torch.LongTensor]:
+        r"""
+        Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and
+        can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.
+
+        Parameters:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                The sequence used as a prompt for the generation.
+            logits_processor (`LogitsProcessorList`):
+                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                used to modify the prediction scores of the language modeling head applied at each generation step.
+            stopping_criteria (`StoppingCriteriaList`):
+                An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                used to tell if the generation loop should stop.
+            generation_config ([`~generation.GenerationConfig`]):
+                The generation configuration to be used as parametrization of the decoding method.
+            synced_gpus (`bool`):
+                Whether to continue running the while loop until max_length (needed to avoid deadlocking with
+                `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3).
+            streamer (`BaseStreamer`, *optional*):
+                Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+                through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+            model_kwargs:
+                Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is
+                an encoder-decoder model the kwargs should include `encoder_outputs`.
+
+        Return:
+            [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or `torch.LongTensor`:
+            A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+            [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+            `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if
+            `model.config.is_encoder_decoder=True`.
+        """
+        # init values
+        pad_token_id = generation_config._pad_token_tensor
+        output_attentions = generation_config.output_attentions
+        output_hidden_states = generation_config.output_hidden_states
+        output_scores = generation_config.output_scores
+        output_logits = generation_config.output_logits
+        return_dict_in_generate = generation_config.return_dict_in_generate
+        max_length = generation_config.max_length
+        has_eos_stopping_criteria = any(
+            hasattr(criteria, "eos_token_id") for criteria in stopping_criteria
+        )
+        do_sample = generation_config.do_sample
+
+        # init attention / hidden states / scores tuples
+        scores = () if (return_dict_in_generate and output_scores) else None
+        raw_logits = () if (return_dict_in_generate and output_logits) else None
+        decoder_attentions = (
+            () if (return_dict_in_generate and output_attentions) else None
+        )
+        cross_attentions = (
+            () if (return_dict_in_generate and output_attentions) else None
+        )
+        decoder_hidden_states = (
+            () if (return_dict_in_generate and output_hidden_states) else None
+        )
+
+        # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+        if return_dict_in_generate and self.config.is_encoder_decoder:
+            encoder_attentions = (
+                model_kwargs["encoder_outputs"].get("attentions")
+                if output_attentions
+                else None
+            )
+            encoder_hidden_states = (
+                model_kwargs["encoder_outputs"].get("hidden_states")
+                if output_hidden_states
+                else None
+            )
+
+        # keep track of which sequences are already finished
+        batch_size, cur_len, num_streams = input_ids.shape
+        this_peer_finished = False
+        unfinished_sequences = torch.ones(
+            batch_size, dtype=torch.long, device=input_ids.device
+        )
+        model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+
+        while self._has_unfinished_sequences(
+            this_peer_finished,
+            synced_gpus,
+            device=input_ids.device,
+            cur_len=cur_len,
+            max_length=max_length,
+        ):
+            # prepare model inputs
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+
+            # prepare variable output controls (note: some models won't accept all output controls)
+            model_inputs.update(
+                {"output_attentions": output_attentions} if output_attentions else {}
+            )
+            model_inputs.update(
+                {"output_hidden_states": output_hidden_states}
+                if output_hidden_states
+                else {}
+            )
+
+            # forward pass to get next token
+            outputs = self(**model_inputs, return_dict=True)
+
+            # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs,
+                model_kwargs,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+            )
+            if synced_gpus and this_peer_finished:
+                continue
+
+            # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
+            # (the clone itself is always small)
+            next_token_logits = outputs.logits.clone()[:, -1].float()
+            next_token_logits = next_token_logits.to(input_ids.device)
+
+            # pre-process distribution
+            batch_size, seq_len, num_streams = input_ids.shape
+            rearrange_input_ids = rearrange(input_ids, "b l n -> (b n) l")
+            next_token_logits = rearrange(next_token_logits, "b n d -> (b n) d")
+            next_token_scores = logits_processor(rearrange_input_ids, next_token_logits)
+            next_token_scores = rearrange(
+                next_token_scores, "(b n) d -> b n d", b=batch_size
+            )
+
+            # Store scores, attentions and hidden_states when required
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_token_scores,)
+                if output_logits:
+                    raw_logits += (next_token_logits,)
+                if output_attentions:
+                    decoder_attentions += (
+                        (outputs.decoder_attentions,)
+                        if self.config.is_encoder_decoder
+                        else (outputs.attentions,)
+                    )
+                    if self.config.is_encoder_decoder:
+                        cross_attentions += (outputs.cross_attentions,)
+
+                if output_hidden_states:
+                    decoder_hidden_states += (
+                        (outputs.decoder_hidden_states,)
+                        if self.config.is_encoder_decoder
+                        else (outputs.hidden_states,)
+                    )
+
+            # token selection
+            if do_sample:
+                probs = nn.functional.softmax(next_token_scores, dim=-1)
+                # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution
+                probs = probs.view(-1, probs.shape[-1])
+                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+                next_tokens = next_tokens.view(*next_token_scores.shape[:-1])
+            else:
+                next_tokens = torch.argmax(next_token_scores, dim=-1)
+
+            # finished sentences should have their next token be a padding token
+            if has_eos_stopping_criteria:
+                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (
+                    1 - unfinished_sequences
+                )
+
+            # update generated ids, model inputs, and length for next step
+            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=1)
+            if streamer is not None:
+                streamer.put(next_tokens.cpu())
+
+            unfinished_sequences = unfinished_sequences & ~stopping_criteria(
+                input_ids, scores
+            )
+            this_peer_finished = unfinished_sequences.max() == 0
+            cur_len += 1
+
+            # This is needed to properly delete outputs.logits which may be very large for first iteration
+            # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration
+            del outputs
+
+        if streamer is not None:
+            streamer.end()
+
+        if return_dict_in_generate:
+            if self.config.is_encoder_decoder:
+                return GenerateEncoderDecoderOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    logits=raw_logits,
+                    encoder_attentions=encoder_attentions,
+                    encoder_hidden_states=encoder_hidden_states,
+                    decoder_attentions=decoder_attentions,
+                    cross_attentions=cross_attentions,
+                    decoder_hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+            else:
+                return GenerateDecoderOnlyOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    logits=raw_logits,
+                    attentions=decoder_attentions,
+                    hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+        else:
+            return input_ids
+
+
+class LearnedPositionEmbeddings(nn.Module):
+
+    def __init__(self, seq_len, model_dim, init=0.02):
+        super().__init__()
+        self.emb = nn.Embedding(seq_len, model_dim)
+        self.emb.weight.data.normal_(mean=0.0, std=init)
+
+    def forward(self, x):
+        sl = x.shape[1]
+        return self.emb(torch.arange(0, sl, device=x.device))
+
+    def get_fixed_embedding(self, ind, dev):
+        return self.emb(torch.tensor([ind], device=dev)).unsqueeze(0)
+
+
+def build_hf_gpt_transformer(
+    layers, model_dim, heads, max_mel_seq_len, max_text_seq_len, checkpointing
+):
+    gpt_config = GPT2Config(
+        vocab_size=256,
+        n_positions=max_mel_seq_len + max_text_seq_len,
+        n_ctx=max_mel_seq_len + max_text_seq_len,
+        n_embd=model_dim,
+        n_layer=layers,
+        n_head=heads,
+        use_cache=not checkpointing,
+        scale_attn_by_inverse_layer_idx=True,
+        reorder_and_upcast_attn=True,
+        attn_implementation="sdpa",
+    )
+    gpt = GPT2Model(gpt_config)
+
+    if checkpointing:
+        gpt.gradient_checkpointing_enable()
+
+    del gpt.wpe, gpt.wte
+    gpt.wpe = functools.partial(null_position_embeddings, dim=model_dim)
+    mel_pos_embs = LearnedPositionEmbeddings(max_mel_seq_len, model_dim)
+
+    return gpt, mel_pos_embs
+
+
+class AcousticLLM(nn.Module):
+
+    def __init__(
+        self,
+        # Model
+        n_stacks=2,
+        layers=12,
+        model_dim=1024,
+        heads=16,
+        # Text
+        max_text_tokens=120,
+        number_text_tokens=8194,
+        start_text_token=8192,
+        stop_text_token=8193,
+        # Speech
+        n_frames_per_step=4,
+        n_heads_per_frame=1,
+        max_speech_tokens=250,
+        number_speech_tokens=8194,
+        start_speech_token=8192,
+        stop_speech_token=8193,
+        # CoS Prediction
+        streaming=False,
+        streaming_delayed_frames=4,
+        accumulative_speech_embedding=False,
+        upsample_factors=2,
+        # Reference embedding
+        max_conditioning_inputs=1,
+        speaker_embedding_pretrained=True,
+        speaker_embedding_ckpt=None,
+        speaker_embedding_dim=256,
+        # For training
+        checkpointing=True,
+        loss_weights=1.0,
+        # For inference
+        delay_prediction=1,
+        temperature=0.3,
+        length_penalty=1.0,
+        repetition_penalty=2.0,
+        top_p=0.2,
+        top_k=50,
+    ):
+        super().__init__()
+        self.n_stacks = n_stacks
+        self.number_text_tokens = number_text_tokens
+        self.start_text_token = start_text_token
+        self.stop_text_token = stop_text_token
+        self.number_speech_tokens = number_speech_tokens
+        self.start_speech_token = start_speech_token
+        self.stop_speech_token = stop_speech_token
+        self.layers = layers
+        self.heads = heads
+
+        self.streaming = streaming
+        self.streaming_delayed_frames = streaming_delayed_frames
+        self.accumulative_speech_embedding = accumulative_speech_embedding
+        self.upsample_factors = upsample_factors
+
+        self.n_frames_per_step = n_frames_per_step
+        self.n_heads_per_frame = n_heads_per_frame
+        self.number_speech_heads = n_heads_per_frame * n_frames_per_step
+
+        self.max_speech_tokens = max_speech_tokens
+        self.max_text_tokens = max_text_tokens
+        self.model_dim = model_dim
+        self.max_conditioning_inputs = max_conditioning_inputs
+
+        self.speaker_embedding_pretrained = speaker_embedding_pretrained
+        self.speaker_embedding_ckpt = speaker_embedding_ckpt
+        self.speaker_embedding_dim = speaker_embedding_dim
+
+        # For training
+        self.loss_weights = loss_weights
+
+        # For inference
+        self.delay_prediction = delay_prediction
+        self.temperature = temperature
+        self.length_penalty = length_penalty
+        self.repetition_penalty = repetition_penalty
+        self.top_p = top_p
+        self.top_k = top_k
+
+        # Conditional embedding
+        self.reference_embedding = nn.Sequential(
+            nn.Linear(speaker_embedding_dim, 256),
+            nn.Tanh(),
+            nn.Linear(256, model_dim),
+        )
+
+        self.text_embedding = nn.Embedding(self.number_text_tokens + 1, model_dim)
+        self.text_embedding.weight.data.normal_(mean=0.0, std=0.02)
+
+        self.mel_embedding = nn.ModuleList(
+            [
+                nn.Embedding(self.number_speech_tokens, model_dim)
+                for _ in range(self.number_speech_heads)
+            ]
+        )
+        for module in self.mel_embedding:
+            module.weight.data.normal_(mean=0.0, std=0.02)
+
+        # Build GPTs
+        self.gpt, self.mel_pos_embedding = build_hf_gpt_transformer(
+            layers,
+            model_dim,
+            heads,
+            self.max_speech_tokens + 2 + self.max_conditioning_inputs,
+            self.max_text_tokens + 2,
+            checkpointing,
+        )
+        self.final_norm = nn.LayerNorm(model_dim)
+        self.mel_head = nn.ModuleList(
+            [
+                nn.Linear(model_dim, self.number_speech_tokens)
+                for _ in range(self.number_speech_heads)
+            ]
+        )
+
+    def post_init_gpt2_config(self, use_deepspeed=False, kv_cache=True, half=False):
+        seq_length = self.max_speech_tokens + self.max_text_tokens + 2
+        gpt_config = GPT2Config(
+            vocab_size=self.max_speech_tokens,
+            n_positions=seq_length,
+            n_ctx=seq_length,
+            n_embd=self.model_dim,
+            n_layer=self.layers,
+            n_head=self.heads,
+            gradient_checkpointing=False,
+            use_cache=True,
+        )
+
+        self.inference_model = MHGPT2InferenceModel(
+            gpt_config,
+            self.gpt,
+            self.mel_pos_embedding,
+            self.mel_embedding,
+            self.final_norm,
+            self.mel_head,
+            kv_cache=kv_cache,
+        )
+        self.inference_model.eval()
+
+    def build_aligned_inputs_and_targets(
+        self, seqs, lens, start_token, stop_token, delay=0
+    ):
+        for i in range(seqs.shape[0]):
+            seqs[i, lens[i] :] = stop_token
+
+        if len(seqs.shape) == 2:
+            inp = F.pad(
+                seqs, (self.streaming_delayed_frames, 0), value=start_token
+            ).type_as(seqs)
+            inp = F.pad(inp, (0, 1), value=stop_token).type_as(inp)
+            tar = F.pad(inp[:, 1:], (0, 1), value=stop_token).type_as(seqs)
+        else:
+            inp = F.pad(
+                seqs, (0, 0, self.streaming_delayed_frames, 0), value=start_token
+            ).type_as(seqs)
+            inp = F.pad(inp, (0, 0, 0, 1), value=stop_token).type_as(inp)
+            tar = F.pad(inp[:, 1:], (0, 0, 0, 1), value=stop_token).type_as(seqs)
+
+        if delay > 0:
+            pad_size = delay * (inp.shape[2] - 1)
+            L = inp.shape[1] + pad_size
+            inp = F.pad(inp, (0, 0, pad_size, 0), value=start_token).type_as(inp)
+            inp = F.pad(inp, (0, 0, 0, pad_size), value=stop_token).type_as(inp)
+            inp = torch.stack(
+                [
+                    inp[:, pad_size - i * delay : pad_size - i * delay + L, i]
+                    for i in range(inp.shape[-1])
+                ],
+                dim=-1,
+            )
+
+            tar = F.pad(tar, (0, 0, pad_size, 0), value=start_token).type_as(tar)
+            tar = F.pad(tar, (0, 0, 0, pad_size), value=stop_token).type_as(tar)
+            tar = torch.stack(
+                [
+                    tar[:, pad_size - i * delay : pad_size - i * delay + L, i]
+                    for i in range(tar.shape[-1])
+                ],
+                dim=-1,
+            )
+
+            lens += pad_size
+
+        return inp, tar, lens + self.streaming_delayed_frames + 1
+
+    def get_logits(
+        self,
+        final_norm,
+        first_inputs,
+        first_head,
+        speech_conditioning_inputs=None,
+        attention_mask=None,
+        get_attns=False,
+        return_latent=False,
+    ):
+        emb = first_inputs
+        if speech_conditioning_inputs is not None:
+            emb = torch.cat([speech_conditioning_inputs, emb], dim=1)
+
+        gpt_out = self.gpt(
+            inputs_embeds=emb,
+            return_dict=True,
+            attention_mask=attention_mask,
+            output_attentions=get_attns,
+        )
+
+        enc = gpt_out.last_hidden_state
+        if speech_conditioning_inputs is not None:
+            enc = enc[:, 1:]
+        enc = final_norm(enc)
+
+        first_logits = [head(enc).permute(0, 2, 1) for head in first_head]
+
+        return first_logits
+
+    @torch.cuda.amp.autocast()
+    def get_conditioning(self, speech_conditioning_input):
+        if hasattr(self, "reference_encoder"):
+            if len(speech_conditioning_input.shape) == 2:
+                speech_conditioning_input = speech_conditioning_input.unsqueeze(1)
+            speech_conditioning_input = self.reference_encoder(
+                speech_conditioning_input
+            )
+        conds = self.reference_embedding(speech_conditioning_input)
+        return conds
+
+    def inference_speech(
+        self,
+        speech_conditioning_latent,
+        text_inputs,
+        input_tokens=None,
+        num_return_sequences=1,
+        max_generate_length=None,
+        **hf_generate_kwargs,
+    ):
+        if not hasattr(self, "inference_model"):
+            self.post_init_gpt2_config()
+
+        # Cond
+        emb = speech_conditioning_latent
+        self.inference_model.store_mel_emb(emb)
+
+        # Text
+        text = torch.repeat_interleave(text_inputs, self.upsample_factors, dim=1)
+        text = F.pad(
+            text,
+            (0, self.streaming_delayed_frames + self.number_speech_heads - 1),
+            value=self.stop_speech_token,
+        )
+        text_embedding = self.text_embedding(text)
+        self.inference_model.store_mel_parallel_emb(text_embedding)
+
+        fake_inputs = torch.full(
+            (
+                emb.shape[0],  # should be 1 for stable inference
+                emb.shape[1] + 1,  # + 1 for the start_speech_token
+                self.number_speech_heads,
+            ),
+            fill_value=self.start_speech_token,
+            dtype=torch.long,
+            device=text_inputs.device,
+        )
+        if input_tokens is None:
+            inputs = fake_inputs
+            prompt_index = 0
+        else:
+            prompt, _, _ = self.build_aligned_inputs_and_targets(
+                input_tokens,
+                torch.Tensor([len(input_tokens[0])]).int(),
+                self.start_speech_token,
+                self.stop_speech_token,
+                self.delay_prediction,
+            )
+            prompt = prompt[:, 1 : 1 + input_tokens.shape[1]]
+            inputs = torch.cat([fake_inputs, prompt], dim=1)
+            prompt_index = input_tokens.shape[1]
+        trunc_index = fake_inputs.shape[1]
+
+        stop_criteria = StoppingCriteriaList(
+            [FixedStoppingCriteria(text_embedding.shape[1], start_index=emb.shape[1])]
+        )
+
+        logits_processor = (
+            LogitsProcessorList(
+                [
+                    MultiHeadRepetitionPenaltyLogitsProcessor(
+                        penalty=self.repetition_penalty,
+                        n_heads=self.n_heads_per_frame,
+                        n_frames=-1,
+                        start_index=trunc_index + prompt_index,
+                    )
+                ]
+            )
+            if self.repetition_penalty > 1.0
+            else LogitsProcessorList()
+        )
+        logits_processor.append(
+            SuppressionLogitsProcessor(suppressed_ids=[self.stop_speech_token])
+        )
+
+        max_length = (
+            trunc_index + self.max_speech_tokens - 1
+            if max_generate_length is None
+            else trunc_index + max_generate_length
+        )
+
+        # Recommandation of temp & top_p: (0.8, 0.8), (0.5, 0.5), (0.3, 0.2), (0.2, 0.1)
+        gen = self.inference_model.generate(
+            inputs,
+            bos_token_id=self.start_speech_token,
+            pad_token_id=self.stop_speech_token,
+            eos_token_id=self.stop_speech_token + 2,
+            max_length=max_length,
+            stopping_criteria=stop_criteria,
+            logits_processor=logits_processor,
+            num_return_sequences=num_return_sequences,
+            do_sample=True,
+            temperature=self.temperature,
+            length_penalty=self.length_penalty,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            **hf_generate_kwargs,
+        )
+
+        seq = gen[0][trunc_index:]
+
+        start, heads = 0, []
+        for j in range(self.number_speech_heads):
+            head = seq[j * self.delay_prediction :, j]
+            start_indices = (head == self.start_speech_token).nonzero(as_tuple=True)[0]
+            start = max(start, start_indices[-1] + 1 if len(start_indices) > 0 else 0)
+            stop = (head == self.stop_speech_token).nonzero(as_tuple=True)[0]
+            stop = stop[0] if len(stop) > 0 else len(head)
+            heads.append(head[:stop])
+
+        min_length = min([len(x) for x in heads])
+        seq = torch.stack(
+            [head[start + prompt_index : min_length] for head in heads], dim=-1
+        )
+        return [seq]

fireredtts/modules/bigvgan/__init__.py

@@ -0,0 +1,2 @@
+from .bigvgan import BigVGAN
+from .mel_spectrogram import MelExtractor

fireredtts/modules/bigvgan/activations.py

@@ -0,0 +1,126 @@
+# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
+
+
+import torch
+from torch import nn, sin, pow
+from torch.nn import Parameter
+
+
+class Snake(nn.Module):
+    """
+    Implementation of a sine-based periodic activation function
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter
+    References:
+        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snake(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha: trainable parameter
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            alpha will be trained along with the rest of your model.
+        """
+        super(Snake, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        Snake ∶= x + 1/a * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+
+class SnakeBeta(nn.Module):
+    """
+    A modified Snake function which uses separate parameters for the magnitude of the periodic components
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter that controls frequency
+        - beta - trainable parameter that controls magnitude
+    References:
+        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snakebeta(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha - trainable parameter that controls frequency
+            - beta - trainable parameter that controls magnitude
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            beta is initialized to 1 by default, higher values = higher-magnitude.
+            alpha will be trained along with the rest of your model.
+        """
+        super(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        SnakeBeta := x + 1/b * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x

fireredtts/modules/bigvgan/alias_free_torch/__init__.py

@@ -0,0 +1,5 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+from .filter import *
+from .resample import *
+from .act import *

fireredtts/modules/bigvgan/alias_free_torch/act.py

@@ -0,0 +1,29 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch.nn as nn
+from .resample import UpSample1d, DownSample1d
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    # x: [B,C,T]
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+
+        return x

fireredtts/modules/bigvgan/alias_free_torch/filter.py

@@ -0,0 +1,98 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+if "sinc" in dir(torch):
+    sinc = torch.sinc
+else:
+    # This code is adopted from adefossez's julius.core.sinc under the MIT License
+    # https://adefossez.github.io/julius/julius/core.html
+    #   LICENSE is in incl_licenses directory.
+    def sinc(x: torch.Tensor):
+        """
+        Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+        __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+        """
+        return torch.where(
+            x == 0,
+            torch.tensor(1.0, device=x.device, dtype=x.dtype),
+            torch.sin(math.pi * x) / math.pi / x,
+        )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+#   LICENSE is in incl_licenses directory.
+def kaiser_sinc_filter1d(
+    cutoff, half_width, kernel_size
+):  # return filter [1,1,kernel_size]
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    # For kaiser window
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+        # Normalize filter to have sum = 1, otherwise we will have a small leakage
+        # of the constant component in the input signal.
+        filter_ /= filter_.sum()
+        filter = filter_.view(1, 1, kernel_size)
+
+    return filter
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        # kernel_size should be even number for stylegan3 setup,
+        # in this implementation, odd number is also possible.
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filter)
+
+    # input [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+        return out

fireredtts/modules/bigvgan/alias_free_torch/resample.py

@@ -0,0 +1,57 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch.nn as nn
+from torch.nn import functional as F
+from .filter import LowPassFilter1d
+from .filter import kaiser_sinc_filter1d
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
+        self.register_buffer("filter", filter)
+
+    # x: [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(
+            x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+        )
+        x = x[..., self.pad_left : -self.pad_right]
+
+        return x
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
+
+    def forward(self, x):
+        xx = self.lowpass(x)
+
+        return xx

fireredtts/modules/bigvgan/bigvgan.py

@@ -0,0 +1,369 @@
+import typing as tp
+import torch
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from .alias_free_torch import Activation1d as TorchActivation1d
+from .activations import Snake, SnakeBeta
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class AMPBlock1(torch.nn.Module):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3, 5),
+        activation=None,
+        snake_logscale=True,
+    ):
+        super(AMPBlock1, self).__init__()
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+        self.num_layers = len(self.convs1) + len(
+            self.convs2
+        )  # total number of conv layers
+
+        Activation1d = TorchActivation1d
+        if (
+            activation == "snake"
+        ):  # periodic nonlinearity with snake function and anti-aliasing
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=Snake(channels, alpha_logscale=snake_logscale)
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif (
+            activation == "snakebeta"
+        ):  # periodic nonlinearity with snakebeta function and anti-aliasing
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=SnakeBeta(channels, alpha_logscale=snake_logscale)
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class AMPBlock2(torch.nn.Module):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3),
+        activation=None,
+        snake_logscale=True,
+    ):
+        super(AMPBlock2, self).__init__()
+
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+        self.convs.apply(init_weights)
+
+        self.num_layers = len(self.convs)  # total number of conv layers
+
+        Activation1d = TorchActivation1d
+
+        if (
+            activation == "snake"
+        ):  # periodic nonlinearity with snake function and anti-aliasing
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=Snake(channels, alpha_logscale=snake_logscale)
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif (
+            activation == "snakebeta"
+        ):  # periodic nonlinearity with snakebeta function and anti-aliasing
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=SnakeBeta(channels, alpha_logscale=snake_logscale)
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        for c, a in zip(self.convs, self.activations):
+            xt = a(x)
+            xt = c(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class BigVGAN(torch.nn.Module):
+    # this is our main BigVGAN model. Applies anti-aliased periodic activation for resblocks.
+    def __init__(
+        self,
+        num_mels: int,
+        upsample_initial_channel: int,
+        resblock_kernel_sizes: tp.List[int],
+        resblock_dilation_sizes: tp.List[tp.List[int]],
+        upsample_rates: tp.List[int],
+        upsample_kernel_sizes: tp.List[int],
+        resblock_type: str = "1",
+        snake_logscale: bool = True,
+        activation: str = "snakebeta",
+        use_tanh_at_final: bool = False,
+        use_bias_at_final: bool = False,
+        **kwargs,
+    ):
+        super(BigVGAN, self).__init__()
+
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+
+        # pre conv
+        self.conv_pre = weight_norm(
+            Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3)
+        )
+
+        # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+        resblock = AMPBlock1 if resblock_type == "1" else AMPBlock2
+
+        # transposed conv-based upsamplers. does not apply anti-aliasing
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                upsample_initial_channel // (2**i),
+                                upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                padding=(k - u) // 2,
+                            )
+                        )
+                    ]
+                )
+            )
+
+        # residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(
+                    resblock(
+                        ch,
+                        k,
+                        d,
+                        activation=activation,
+                        snake_logscale=snake_logscale,
+                    )
+                )
+
+        Activation1d = TorchActivation1d
+
+        # post conv
+        if (
+            activation == "snake"
+        ):  # periodic nonlinearity with snake function and anti-aliasing
+            activation_post = Snake(ch, alpha_logscale=snake_logscale)
+            self.activation_post = Activation1d(activation=activation_post)
+        elif (
+            activation == "snakebeta"
+        ):  # periodic nonlinearity with snakebeta function and anti-aliasing
+            activation_post = SnakeBeta(ch, alpha_logscale=snake_logscale)
+            self.activation_post = Activation1d(activation=activation_post)
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+        # whether to use bias for the final conv_post. Defaults to True for backward compatibility
+        self.use_bias_at_final = use_bias_at_final
+        self.conv_post = weight_norm(
+            Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)
+        )
+
+        # weight initialization
+        for i in range(len(self.ups)):
+            self.ups[i].apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+        # final tanh activation. Defaults to True for backward compatibility
+        self.use_tanh_at_final = use_tanh_at_final
+
+    def forward(self, x):
+        # pre conv
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            # upsampling
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            # AMP blocks
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+        # post conv
+        x = self.activation_post(x)
+        x = self.conv_post(x)
+        # final tanh activation
+        if self.use_tanh_at_final:
+            x = torch.tanh(x)
+        else:
+            x = torch.clamp(x, min=-1.0, max=1.0)  # bound the output to [-1, 1]
+
+        return x
+
+    def remove_weight_norm(self):
+        print("Removing weight norm...")
+        for l in self.ups:
+            for l_i in l:
+                remove_weight_norm(l_i)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)

fireredtts/modules/bigvgan/mel_spectrogram.py

@@ -0,0 +1,111 @@
+import numpy as np
+import torch
+import torchaudio
+from librosa.filters import mel as librosa_mel_fn
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def mel_spectrogram(
+    y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
+):
+    global mel_basis, hann_window  # pylint: disable=global-statement
+    if f"{str(fmax)}_{str(y.device)}" not in mel_basis:
+        mel = librosa_mel_fn(
+            sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
+        )
+        mel_basis[str(fmax) + "_" + str(y.device)] = (
+            torch.from_numpy(mel).float().to(y.device)
+        )
+        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
+
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
+    y = y.squeeze(1)
+
+    spec = torch.view_as_real(
+        torch.stft(
+            y,
+            n_fft,
+            hop_length=hop_size,
+            win_length=win_size,
+            window=hann_window[str(y.device)],
+            center=center,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=True,
+        )
+    )
+
+    spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+    spec = torch.matmul(mel_basis[str(fmax) + "_" + str(y.device)], spec)
+    spec = spectral_normalize_torch(spec)
+
+    return spec
+
+
+class MelExtractor(object):
+    def __init__(
+        self,
+        num_mels: int = 80,
+        n_fft: int = 1920,
+        hop_size: int = 480,
+        win_size: int = 1920,
+        sampling_rate: int = 24000,
+        fmin: int = 0,
+        fmax: int = 8000,
+        center: bool = False,
+    ):
+        super().__init__()
+        self.num_mels = num_mels
+        self.n_fft = n_fft
+        self.hop_size = hop_size
+        self.win_size = win_size
+        self.sampling_rate = sampling_rate
+        self.fmin = fmin
+        self.fmax = fmax
+        self.center = center
+
+    def __call__(self, audio: torch.Tensor, audio_sr: int):
+        """Args:
+            audio(torch.Tensor): shape (1, t)
+        Returns:
+            mel(torch.Tensor): shape (1, num_mels, t')
+        """
+        if audio_sr != self.sampling_rate:
+            audio = torchaudio.functional.resample(
+                audio, orig_freq=audio_sr, new_freq=self.sampling_rate
+            )
+            audio_sr = self.sampling_rate
+        mel = mel_spectrogram(
+            audio,
+            self.n_fft,
+            self.num_mels,
+            self.sampling_rate,
+            self.hop_size,
+            self.win_size,
+            self.fmin,
+            self.fmax,
+            self.center,
+        )  # (1, num_mels, t)
+        return mel

fireredtts/modules/flowmatching/__init__.py

@@ -0,0 +1,18 @@
+from .estimator_dit import DiT
+from .upsample_encoder import UpsampleConformerEncoder
+from .flow import CausalFmWithSpkCtx, DualEmbedding
+
+
+class FlowToken2Mel(CausalFmWithSpkCtx):
+    def __init__(self, config):
+        token_emb = DualEmbedding(**config['token_emb'])
+        encoder = UpsampleConformerEncoder(**config['encoder'])
+        estimator = DiT(**config['estimator'])
+        super().__init__(
+            spk_channels=config['spk_channels'],
+            spk_enc_channels=config['spk_enc_channels'],
+            infer_cfg_rate=config['infer_cfg_rate'],
+            token_emb=token_emb,
+            encoder=encoder,
+            estimator=estimator,
+        )

fireredtts/modules/flowmatching/estimator_dit.py

@@ -0,0 +1,356 @@
+import math
+import torch
+from typing import Optional
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MLP(torch.nn.Module):
+    def __init__(
+            self,
+            in_features:int,
+            hidden_features:Optional[int]=None,
+            out_features:Optional[int]=None,
+            act_layer=nn.GELU,
+            norm_layer=None,
+            bias=True,
+            drop=0.,
+    ):
+        super().__init__()
+        hidden_features = hidden_features or in_features
+        out_features = out_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop)
+        self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop2 = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.norm(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+
+
+class Attention(torch.nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            head_dim: int = 64,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.inner_dim = num_heads * head_dim
+        self.scale = head_dim ** -0.5
+
+        self.to_q = nn.Linear(dim, self.inner_dim, bias=qkv_bias)
+        self.to_k = nn.Linear(dim, self.inner_dim, bias=qkv_bias)
+        self.to_v = nn.Linear(dim, self.inner_dim, bias=qkv_bias)
+
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.proj = nn.Linear(self.inner_dim, dim)
+
+    def to_heads(self, ts:torch.Tensor):
+        b, t, c = ts.shape
+        # (b, t, nh, c)
+        ts = ts.reshape(b, t, self.num_heads, c // self.num_heads)
+        ts = ts.transpose(1, 2)
+        return ts
+    
+    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
+        """Args:
+            x(torch.Tensor): shape (b, t, c)
+            attn_mask(torch.Tensor): shape (b, t, t)
+        """
+        b, t, c = x.shape
+
+        q = self.to_q(x)
+        k = self.to_k(x)
+        v = self.to_v(x)
+
+        q = self.to_heads(q)    # (b, nh, t, c)
+        k = self.to_heads(k)
+        v = self.to_heads(v)
+    
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(1)
+        
+        x = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_drop.p if self.training else 0.,
+        )   # (b, nh, t, c)
+        x = x.transpose(1, 2).reshape(b, t, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+   
+
+def modulate(x, shift, scale):
+    return x * (1 + scale) + shift
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+        # from SinusoidalPosEmb
+        self.scale = 1000
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t * self.scale, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+# Convolution related
+class Transpose(torch.nn.Module):
+    def __init__(self, dim0: int, dim1: int):
+        super().__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, x: torch.Tensor):
+        x = torch.transpose(x, self.dim0, self.dim1)
+        return x
+
+
+class CausalConv1d(torch.nn.Conv1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+    ) -> None:
+        super(CausalConv1d, self).__init__(in_channels, out_channels, kernel_size)
+        self.causal_padding = (kernel_size - 1, 0)
+
+    def forward(self, x: torch.Tensor):
+        x = F.pad(x, self.causal_padding)
+        x = super(CausalConv1d, self).forward(x)
+        return x
+
+
+class CausalConvBlock(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: int = 3,
+                 ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+
+        self.block = torch.nn.Sequential(
+            # norm
+            # conv1
+            Transpose(1, 2),
+            CausalConv1d(in_channels, out_channels, kernel_size),
+            Transpose(1, 2),
+            # norm & act
+            nn.LayerNorm(out_channels),
+            nn.Mish(),
+            # conv2
+            Transpose(1, 2),
+            CausalConv1d(out_channels, out_channels, kernel_size),
+            Transpose(1, 2),
+        )
+    
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
+        """
+        Args:
+            x: shape (b, t, c)
+            mask: shape (b, t, 1)
+        """
+        if mask is not None: x = x * mask
+        x = self.block(x)
+        if mask is not None: x = x * mask
+        return x
+    
+
+class DiTBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_size, num_heads, head_dim, mlp_ratio=4.0, **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, head_dim=head_dim, qkv_bias=True, qk_norm=True, **block_kwargs)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = MLP(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+        self.norm3 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.conv = CausalConvBlock(in_channels=hidden_size, out_channels=hidden_size, kernel_size=3)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 9 * hidden_size, bias=True)
+        )
+
+    def forward(self, x:torch.Tensor, c:torch.Tensor, attn_mask:torch.Tensor=None, conv_mask:torch.Tensor=None):
+        """Args
+            x: shape (b, t, c)
+            c: shape (b, 1, c)
+            attn_mask: shape (b, t, t), bool type attention mask
+            conv_mask: shape (b, 1, t), bool type non-pad mask
+        """
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_conv, scale_conv, gate_conv \
+              = self.adaLN_modulation(c).chunk(9, dim=-1)
+        # attention
+        x = x + gate_msa * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), attn_mask=attn_mask)
+        # conv
+        x = x + gate_conv * self.conv(modulate(self.norm3(x), shift_conv, scale_conv), mask=conv_mask)
+        # mlp
+        x = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class DiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        mlp_ratio: float = 4.0,
+        depth: int = 28,
+        num_heads: int = 8,
+        head_dim: int = 64,
+        hidden_size: int = 256,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.t_embedder = TimestepEmbedder(hidden_size)
+
+        self.in_proj = nn.Linear(in_channels, hidden_size)
+
+        self.blocks = nn.ModuleList([
+            DiTBlock(hidden_size, num_heads, head_dim, mlp_ratio=mlp_ratio) for _ in range(depth)
+        ])
+        self.final_layer = FinalLayer(hidden_size, self.out_channels)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    """For non-streaming inference.
+    """
+    def forward(self, x:torch.Tensor, c:torch.Tensor, t:torch.Tensor, attn_mask:torch.Tensor=None, conv_mask:torch.Tensor=None):
+        """
+        Args:
+            x: shape (b, c, t)
+            c: aux condition, shape (b, c, t)
+            t: shape (b,)
+            attn_mask: (b, t, t)
+            conv_mask: (b, 1, t)
+        Returns:
+            pred: shape (b, c, t)
+        """
+        # time
+        t = self.t_embedder(t.view(-1)).unsqueeze(1)  # (b, 1, c)
+
+        # CausalConvBlock mask is (b, t, 1)
+        conv_mask = conv_mask if conv_mask is None else conv_mask.transpose(1, 2)
+        
+        x = torch.cat([x, c], dim=1)
+        # forward blocks
+        x = x.transpose(1, 2)
+        x = self.in_proj(x)
+        for block in self.blocks:
+            x = block(x, t, attn_mask=attn_mask, conv_mask=conv_mask)
+        x = self.final_layer(x, t)
+        x = x.transpose(1, 2)
+        return x
+

fireredtts/modules/flowmatching/flow.py

@@ -0,0 +1,138 @@
+import torch
+import torch.nn.functional as F
+from typing import Dict, List
+from einops import pack, repeat
+from .estimator_dit import DiT
+from .upsample_encoder import UpsampleConformerEncoder
+
+
+class DualEmbedding(torch.nn.Module):
+    def __init__(
+        self, 
+        channels:int=512,
+    ):
+        super().__init__()
+        self.codebook_size = 128
+        self.codebook_dim = 128
+        self.codebook = torch.nn.ModuleList([
+            torch.nn.Embedding(self.codebook_size, self.codebook_dim),
+            torch.nn.Embedding(self.codebook_size, self.codebook_dim),
+        ])
+        self.out_proj = torch.nn.Linear(self.codebook_dim * 2, channels)
+
+    def forward(self, tokens):
+        """
+        Args:
+            tokens: shape (b, t)
+        Returns:
+            token_embs: shape (b, t, c)
+        """
+        token_embs = torch.cat([
+            self.codebook[0](tokens % self.codebook_size), 
+            self.codebook[1](tokens // self.codebook_size)
+        ], dim=-1)
+        token_embs = self.out_proj(token_embs)
+        return token_embs
+
+
+class CausalFmWithSpkCtx(torch.nn.Module):
+    def __init__(
+        self,
+        # Basic in-out
+        spk_channels: int,
+        spk_enc_channels: int,  # out channels of spk & encoder projection
+        # Module
+        token_emb: DualEmbedding,
+        encoder: UpsampleConformerEncoder,
+        estimator: DiT,
+        # Flow cfg
+        infer_cfg_rate: float = 0.7,
+    ):
+        super().__init__()
+        # Variants
+        self.up_stride = encoder.up_stride
+        self.infer_cfg_rate = infer_cfg_rate
+        # Module
+        self.spk_proj = torch.nn.Linear(spk_channels, spk_enc_channels)
+        self.token_emb = token_emb
+        self.encoder = encoder
+        self.encoder_proj = torch.nn.Linear(encoder.output_size, spk_enc_channels)
+        self.estimator = estimator
+        # Initial noise, maximum of 600s
+        self.register_buffer(
+            "x0",
+            torch.randn([1, self.estimator.out_channels, 50 * 600]),
+            persistent=False,
+        )
+
+    def _euler(
+        self,
+        x0: torch.Tensor,
+        c: torch.Tensor,
+        n_timesteps: int = 10,
+    ):
+        # time steps
+        t_span = torch.linspace(0, 1, n_timesteps + 1).to(x0)
+        # cosine time schduling
+        t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+        # euler solver
+        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+        t = t.unsqueeze(dim=0)
+
+        xt = x0
+        for step in range(1, len(t_span)):
+            # pack input
+            x_in = torch.cat([xt, xt], dim=0)
+            c_in = torch.cat([c, torch.zeros_like(c)], dim=0)
+            t_in = torch.cat([t, t], dim=0)
+
+            # model call
+            with torch.no_grad():
+                vt = self.estimator.forward(x_in, c_in, t_in)
+            # cfg
+            vt_cond, vt_cfg = vt.chunk(2, dim=0)
+            vt = (1.0 + self.infer_cfg_rate) * vt_cond - self.infer_cfg_rate * vt_cfg
+
+            xt = xt + dt * vt
+            t = t + dt
+            if step < len(t_span) - 1:
+                dt = t_span[step + 1] - t
+        return xt
+
+    def inference(
+        self,
+        prompt_token: torch.Tensor,
+        prompt_xvec: torch.Tensor,
+        prompt_feat: torch.Tensor,
+        token: torch.Tensor,
+    ):
+        # NOTE align prompt_token, prompt_feat in advance
+
+        # Spk condition
+        embedding = F.normalize(prompt_xvec, dim=1)
+        spks = self.spk_proj(embedding)
+
+        # Token condition
+        token = torch.concat([prompt_token, token], dim=1)
+        xs = self.token_emb(token)
+
+        xs_lens = torch.tensor([xs.shape[1]]).to(token)
+        xs = self.encoder(xs, xs_lens)
+        mu = self.encoder_proj(xs)
+
+        # Mel context
+        ctx = torch.zeros_like(mu)
+        ctx[:, : prompt_feat.shape[1]] = prompt_feat
+
+        # Compose condition
+        cond = mu.transpose(1, 2)
+        ctx = ctx.transpose(1, 2)
+        spks = repeat(spks, "b c -> b c t", t=cond.shape[-1])
+        cond = pack([cond, spks, ctx], "b * t")[0]
+
+        # FM inference
+        x0 = self.x0[..., : mu.shape[1]]
+        x1 = self._euler(x0, cond, n_timesteps=10)
+
+        feat = x1.transpose(1, 2)[:, prompt_feat.shape[1] :]
+        return feat

fireredtts/modules/flowmatching/upsample_encoder.py

@@ -0,0 +1,617 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+from typing import Tuple, List, Union
+
+
+"""Attention modules.
+"""
+class MultiHeadedAttention(nn.Module):
+    def __init__(self,
+                 n_head: int,
+                 n_feat: int,
+                 dropout_rate: float,
+                 key_bias: bool = True):
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat, bias=key_bias)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self, 
+                    query: torch.Tensor, 
+                    key: torch.Tensor, 
+                    value: torch.Tensor):
+        """
+        Args:
+            query,key,value: shape (b, t, c)
+        Returns:
+            query,key,value: shape (b, nh, t, c//nh)
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+        return q, k, v
+
+    def forward_attention(self,
+                          value: torch.Tensor,
+                          scores: torch.Tensor,
+                          mask: torch.Tensor = None):
+        """Compute attention context vector.
+        Args:
+            value (torch.Tensor): shape: (b, nh, t2, c//nh).
+            scores (torch.Tensor): shape: (b, nh, t1, t2).
+            mask (torch.Tensor): attention padded mask, size (b, 1, t2) or (b, t1, t2)
+        Returns:
+            shape: (b, t1, c)
+        """
+        b = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1).eq(0)
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = scores.softmax(dim=-1).masked_fill(mask, 0.0)
+        else:
+            attn = scores.softmax(dim=-1)
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = x.transpose(1, 2).contiguous().view(b, -1, self.h * self.d_k)
+        return self.linear_out(x)
+
+
+class RelPositionMultiHeadedAttention(MultiHeadedAttention):
+    def __init__(self,
+                 n_head: int,
+                 n_feat: int,
+                 dropout_rate: float,
+                 key_bias: bool = True):
+        """Multi-Head Attention layer with relative position encoding.
+        Paper: https://arxiv.org/abs/1901.02860
+        Args:
+            n_head (int): The number of heads.
+            n_feat (int): The number of features.
+            dropout_rate (float): Dropout rate.
+        """
+        super().__init__(n_head, n_feat, dropout_rate, key_bias)
+        # linear transformation for positional encoding
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        # these two learnable bias are used in matrix c and matrix d
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        torch.nn.init.xavier_uniform_(self.pos_bias_u)
+        torch.nn.init.xavier_uniform_(self.pos_bias_v)
+
+    def rel_shift(self, x: torch.Tensor) -> torch.Tensor:
+        """Compute relative positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1).
+            time1 means the length of query vector.
+
+        Returns:
+            torch.Tensor: Output tensor.
+
+        """
+        zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1),
+                               device=x.device,
+                               dtype=x.dtype)
+        x_padded = torch.cat([zero_pad, x], dim=-1)
+
+        x_padded = x_padded.view(x.size()[0],
+                                 x.size()[1],
+                                 x.size(3) + 1, x.size(2))
+        x = x_padded[:, :, 1:].view_as(x)[
+            :, :, :, : x.size(-1) // 2 + 1
+        ]  # only keep the positions from 0 to time2
+        return x
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask: torch.Tensor = None,
+        cache: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            query (torch.Tensor): shape (b, t1, c).
+            key (torch.Tensor): shape (b, t2, c).
+            value (torch.Tensor): shape (b, t2, c).
+            mask (torch.Tensor): attention padded mask, shape (b, 1, t2) or (b, t1, t2).
+            pos_emb (torch.Tensor): Positional embedding tensor (b, 2*t1-1, c).
+            cache (torch.Tensor): Cache tensor (1, nh, cache_t, d_k * 2).
+        Returns:
+            torch.Tensor: Output tensor (b, t1, d_model).
+            torch.Tensor: Cache tensor (1, nh, cache_t + t1, d_k * 2)
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
+
+        if cache is not None and cache.size(0) > 0:
+            key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        new_cache = torch.cat((k, v), dim=-1)
+
+        n_batch_pos = pos_emb.size(0)
+        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)    # (batch, 2*time1-1, head, d_k)
+        p = p.transpose(1, 2)  # (batch, head, 2*time1-1, d_k)
+
+        # (batch, head, time1, d_k)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        # (batch, head, time1, d_k)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+        # compute attention score
+        # first compute matrix a and matrix c
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        # (batch, head, time1, time2)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+
+        # compute matrix b and matrix d
+        # matrix_bd: (batch, head, time1, 2*time1-1)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        # NOTE(Xiang Lyu): Keep rel_shift since espnet rel_pos_emb is used
+        if matrix_ac.shape != matrix_bd.shape:
+            matrix_bd = self.rel_shift(matrix_bd)
+
+        scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k)  # (batch, head, time1, time2)
+
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+class EspnetRelPositionalEncoding(torch.nn.Module):
+    """Relative positional encoding module (new implementation).
+
+    Details can be found in https://github.com/espnet/espnet/pull/2816.
+
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+
+    """
+
+    def __init__(self, d_model: int, dropout_rate: float=0.0, max_len: int = 5000):
+        """Construct an PositionalEncoding object."""
+        super(EspnetRelPositionalEncoding, self).__init__()
+        self.d_model = d_model
+        self.xscale = math.sqrt(self.d_model)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+
+    def extend_pe(self, x: torch.Tensor):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            # self.pe contains both positive and negative parts
+            # the length of self.pe is 2 * input_len - 1
+            if self.pe.size(1) >= x.size(1) * 2 - 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
+        # Suppose `i` means to the position of query vecotr and `j` means the
+        # position of key vector. We use position relative positions when keys
+        # are to the left (i>j) and negative relative positions otherwise (i<j).
+        pe_positive = torch.zeros(x.size(1), self.d_model)
+        pe_negative = torch.zeros(x.size(1), self.d_model)
+        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe_positive[:, 0::2] = torch.sin(position * div_term)
+        pe_positive[:, 1::2] = torch.cos(position * div_term)
+        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
+        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
+
+        # Reserve the order of positive indices and concat both positive and
+        # negative indices. This is used to support the shifting trick
+        # as in https://arxiv.org/abs/1901.02860
+        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
+        pe_negative = pe_negative[1:].unsqueeze(0)
+        pe = torch.cat([pe_positive, pe_negative], dim=1)
+        self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+    def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0) \
+            -> Tuple[torch.Tensor, torch.Tensor]:
+        """Add positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+
+        """
+        self.extend_pe(x)
+        x = x * self.xscale
+        pos_emb = self.position_encoding(size=x.size(1), offset=offset)
+        return self.dropout(x), self.dropout(pos_emb)
+
+    def position_encoding(self,
+                          offset: Union[int, torch.Tensor],
+                          size: int) -> torch.Tensor:
+        """ For getting encoding in a streaming fashion
+
+        Attention!!!!!
+        we apply dropout only once at the whole utterance level in a none
+        streaming way, but will call this function several times with
+        increasing input size in a streaming scenario, so the dropout will
+        be applied several times.
+
+        Args:
+            offset (int or torch.tensor): start offset
+            size (int): required size of position encoding
+
+        Returns:
+            torch.Tensor: Corresponding encoding
+        """
+        pos_emb = self.pe[
+            :,
+            self.pe.size(1) // 2 - size + 1: self.pe.size(1) // 2 + size,
+        ]
+        return pos_emb
+
+
+"""Other modules.
+"""
+class Upsample1D(nn.Module):
+    """A 1D upsampling layer with an optional convolution.
+
+    Parameters:
+        channels (`int`):
+            number of channels in the inputs and outputs.
+        use_conv (`bool`, default `False`):
+            option to use a convolution.
+        use_conv_transpose (`bool`, default `False`):
+            option to use a convolution transpose.
+        out_channels (`int`, optional):
+            number of output channels. Defaults to `channels`.
+    """
+
+    def __init__(self, channels: int, out_channels: int, stride: int = 2):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels
+        self.stride = stride
+        self.conv = nn.Conv1d(self.channels, self.out_channels, stride * 2 + 1, stride=1, padding=0)
+
+    def forward(self, inputs: torch.Tensor, input_lengths: torch.Tensor):
+        outputs = F.interpolate(inputs, scale_factor=self.stride, mode="nearest")
+        outputs = F.pad(outputs, (self.stride * 2, 0), value=0.0)
+        outputs = self.conv(outputs)
+        return outputs, input_lengths * self.stride
+
+
+class PreLookaheadLayer(nn.Module):
+    def __init__(self, channels: int, pre_lookahead_len: int = 1):
+        super().__init__()
+        self.channels = channels
+        self.pre_lookahead_len = pre_lookahead_len
+        self.conv1 = nn.Conv1d(
+            channels, channels,
+            kernel_size=pre_lookahead_len + 1,
+            stride=1, padding=0,
+        )
+        self.conv2 = nn.Conv1d(
+            channels, channels,
+            kernel_size=3, stride=1, padding=0,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        inputs: (batch_size, seq_len, channels)
+        """
+        outputs = inputs.transpose(1, 2).contiguous()
+        # look ahead
+        outputs = F.pad(outputs, (0, self.pre_lookahead_len), mode='constant', value=0.0)
+        outputs = F.leaky_relu(self.conv1(outputs))
+        # outputs
+        outputs = F.pad(outputs, (2, 0), mode='constant', value=0.0)
+        outputs = self.conv2(outputs)
+        outputs = outputs.transpose(1, 2).contiguous()
+
+        # residual connection
+        outputs = outputs + inputs
+        return outputs
+    
+
+class PositionwiseFeedForward(torch.nn.Module):
+    """Positionwise feed forward layer.
+
+    FeedForward are appied on each position of the sequence.
+    The output dim is same with the input dim.
+
+    Args:
+        idim (int): Input dimenstion.
+        hidden_units (int): The number of hidden units.
+        dropout_rate (float): Dropout rate.
+        activation (torch.nn.Module): Activation function
+    """
+
+    def __init__(
+            self,
+            idim: int,
+            hidden_units: int,
+            dropout_rate: float,
+            activation: torch.nn.Module = torch.nn.ReLU(),
+    ):
+        """Construct a PositionwiseFeedForward object."""
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = torch.nn.Linear(idim, hidden_units)
+        self.activation = activation
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.w_2 = torch.nn.Linear(hidden_units, idim)
+
+    def forward(self, xs: torch.Tensor) -> torch.Tensor:
+        """Forward function.
+
+        Args:
+            xs: input tensor (B, L, D)
+        Returns:
+            output tensor, (B, L, D)
+        """
+        return self.w_2(self.dropout(self.activation(self.w_1(xs))))
+
+
+class LinearNoSubsampling(torch.nn.Module):
+    """Linear transform the input without subsampling
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+    """
+    def __init__(self, 
+                 idim: int, 
+                 odim: int, 
+                 dropout_rate: float,
+                 pos_enc_class: torch.nn.Module
+        ):
+        """Construct an linear object."""
+        super().__init__()
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(idim, odim),
+            torch.nn.LayerNorm(odim, eps=1e-5),
+            torch.nn.Dropout(dropout_rate),
+        )
+        self.pos_enc = pos_enc_class
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        offset: int = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Input x.
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: linear input tensor (#batch, time', odim),
+                where time' = time .
+            torch.Tensor: linear input mask (#batch, 1, time'),
+                where time' = time .
+        """
+        x = self.out(x)
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb
+
+
+"""Encoder layer & encoder
+"""
+class ConformerEncoderLayer(nn.Module):
+    """Encoder layer module.
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
+            instance can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward` instance can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool):
+            True: use layer_norm before each sub-block.
+            False: use layer_norm after each sub-block.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        self_attn: torch.nn.Module,
+        feed_forward: torch.nn.Module,
+        dropout_rate: float = 0.1,
+        normalize_before: bool = True,
+    ):
+        """Construct an EncoderLayer object."""
+        super().__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.norm_ff = nn.LayerNorm(size, eps=1e-12)  # for the FNN module
+        self.norm_mha = nn.LayerNorm(size, eps=1e-12)  # for the MHA module
+        self.ff_scale = 1.0
+        self.dropout = nn.Dropout(dropout_rate)
+        self.size = size
+        self.normalize_before = normalize_before
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        pos_emb: torch.Tensor,
+        att_cache: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x: shape (b, t, c)
+            mask: self-attention padded mask, shape (b, 1, t) or (b, t, t)
+            pos_emb: relative positional embedding, shape (b, t, 2t-1)
+            att_cache: shape (1, nh, cache_t, d_k * 2)
+        """
+        # multi-headed self-attention module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_mha(x)
+        # att_cache: (b, head, cache_t, d_k*2)
+        x_att, new_att_cache = self.self_attn(x, x, x, pos_emb, mask, att_cache)
+        x = residual + self.dropout(x_att)
+        if not self.normalize_before:
+            x = self.norm_mha(x)
+
+        # feed forward module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_ff(x)
+        x_ffn = self.feed_forward(x)
+        x = residual + self.ff_scale * self.dropout(x_ffn)
+        if not self.normalize_before:
+            x = self.norm_ff(x)
+
+        return x, new_att_cache
+
+
+class UpsampleConformerEncoder(torch.nn.Module):
+
+    def __init__(
+        self,
+        # Common
+        input_size: int = 512,
+        output_size: int = 512,
+        num_blocks: int = 6,
+        num_up_blocks: int = 4,
+        normalize_before: bool = True,
+        # Input & upsampling
+        up_stride: int = 2,
+        pre_lookahead_len: int = 3,
+        # Attention
+        attention_heads: int = 4,
+        key_bias: bool = True,
+        # MLP
+        linear_units: int = 2048,
+        # Dropouts
+        dropout_rate: float = 0.0,
+        positional_dropout_rate: float = 0.0,
+        attention_dropout_rate: float = 0.0,
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.up_stride = up_stride
+        # Input embedding
+        self.embed = LinearNoSubsampling(
+            input_size,
+            output_size,
+            dropout_rate,
+            # Positional encoding
+            EspnetRelPositionalEncoding(output_size, positional_dropout_rate),
+        )
+        # Look ahead
+        self.pre_lookahead_layer = PreLookaheadLayer(channels=output_size, pre_lookahead_len=pre_lookahead_len)
+        # Norm
+        self.normalize_before = normalize_before
+        self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
+        # Act
+        activation = torch.nn.SiLU()
+        # Self-attention module definition
+        encoder_selfattn_layer_args = (
+            attention_heads,
+            output_size,
+            attention_dropout_rate,
+            key_bias,
+        )
+        # Feed-forward module definition
+        positionwise_layer_args = (
+            output_size,
+            linear_units,
+            dropout_rate,
+            activation,
+        )
+        # 1st Conformer
+        self.encoders = torch.nn.ModuleList([
+            ConformerEncoderLayer(
+                output_size,
+                # Self-attn
+                RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args),
+                # FFN
+                PositionwiseFeedForward(*positionwise_layer_args),
+                dropout_rate,
+                normalize_before,
+            ) for _ in range(num_blocks)
+        ])
+        # Upsample 
+        self.up_layer = Upsample1D(channels=output_size, out_channels=output_size, stride=up_stride)
+        # Input embedding2
+        self.up_embed = LinearNoSubsampling(
+            input_size,
+            output_size,
+            dropout_rate,
+            # Positional encoding
+            EspnetRelPositionalEncoding(output_size, positional_dropout_rate),
+        )
+        # 2nd Conformer
+        self.up_encoders = torch.nn.ModuleList([
+            ConformerEncoderLayer(
+                output_size,
+                # Self-attn
+                RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args),
+                # FFN
+                PositionwiseFeedForward(*positionwise_layer_args),
+                dropout_rate,
+                normalize_before,
+            ) for _ in range(num_up_blocks)
+        ])
+
+    """For non-streaming inference.
+    """
+    def forward(
+        self,
+        xs: torch.Tensor,
+        xs_lens: torch.Tensor,
+        # attention mask BEFORE upsample
+        attn_mask1: torch.Tensor=None,
+        # attention mask AFTER upsample
+        attn_mask2: torch.Tensor=None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            xs: shape (b, t, c)
+            xs_lens: shape (b,)
+            attn_mask1: (token level) shape (b, t, t)
+            attn_mask2: (mel level) shape (b, 2t, 2t)
+        """
+        # Input & lookahead
+        xs, pos_emb = self.embed(xs)
+        xs = self.pre_lookahead_layer(xs)
+        
+        # 1st Conformer
+        for block in self.encoders:
+            xs, _ = block(xs, mask=attn_mask1, pos_emb=pos_emb)
+        
+        # Upsample to mel-level
+        xs = xs.transpose(1, 2).contiguous()
+        xs, xs_lens = self.up_layer(xs, xs_lens)
+        xs = xs.transpose(1, 2).contiguous()
+        # Input
+        xs, pos_emb = self.up_embed(xs)
+        
+        # 2nd Conformer
+        for block in self.up_encoders:
+            xs, _ = block(xs, mask=attn_mask2, pos_emb=pos_emb)
+
+        if self.normalize_before:
+            xs = self.after_norm(xs)
+        return xs

fireredtts/modules/semantic_llm/llm_gpt2.py

@@ -0,0 +1,608 @@
+import math
+import random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence, unpad_sequence
+import functools
+from transformers import GPT2PreTrainedModel, GPT2Model, GPT2Config
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+
+
+# GPT2 NROMAL INFERENCE MODE
+class GPT2InferenceModel(GPT2PreTrainedModel):
+    """Override GPT2LMHeadModel to allow for prefix conditioning."""
+
+    def __init__(self, config, gpt, pos_emb, embeddings, norm, linear, kv_cache):
+        super().__init__(config)
+        self.transformer = gpt
+        self.pos_embedding = pos_emb
+        self.embeddings = embeddings
+        self.final_norm = norm
+        self.lm_head = nn.Sequential(norm, linear)
+        self.kv_cache = kv_cache
+
+    def store_prefix_emb(self, prefix_emb):
+        self.cached_prefix_emb = prefix_emb
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)  # usually None
+        if not self.kv_cache:
+            past_key_values = None
+
+        # only last token for inputs_ids if past is defined in kwargs
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values is not None:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        assert self.cached_prefix_emb is not None
+        assert inputs_embeds is None  # Not supported by this inference model.
+        assert labels is None  # Training not supported by this inference model.
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # assert len(past_key_values) + len(input_ids) == attention_mask.shape[1]
+
+        # Create embedding
+        prefix_len = self.cached_prefix_emb.shape[1]
+        if input_ids.shape[1] != 1:
+            gen_inputs = input_ids[:, prefix_len:]
+            gen_emb = self.embeddings(gen_inputs)
+            gen_emb = gen_emb + self.pos_embedding(gen_emb)
+            if self.cached_prefix_emb.shape[0] != gen_emb.shape[0]:
+                prefix_emb = self.cached_prefix_emb.repeat_interleave(
+                    gen_emb.shape[0] // self.cached_prefix_emb.shape[0], 0
+                )
+            else:
+                prefix_emb = self.cached_prefix_emb.to(gen_emb.dtype)
+            emb = torch.cat([prefix_emb, gen_emb], dim=1)
+        else:
+            emb = self.embeddings(input_ids)
+            emb = emb + self.pos_embedding.get_fixed_embedding(
+                attention_mask.shape[1] - (prefix_len + 1), attention_mask.device
+            )
+        transformer_outputs = self.transformer(
+            inputs_embeds=emb,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+
+        if not return_dict:
+            return (lm_logits,) + transformer_outputs[1:]
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(
+                past_state.index_select(0, beam_idx.to(past_state.device))
+                for past_state in layer_past
+            )
+            for layer_past in past
+        )
+
+
+# GPT2 INDEX-CONTEXT INFERENCE MODE
+class GPT2ICInferenceModel(GPT2PreTrainedModel):
+    """Override GPT2LMHeadModel to allow for prefix conditioning."""
+
+    def __init__(self, config, gpt, pos_emb, embeddings, norm, linear, kv_cache):
+        super().__init__(config)
+        self.transformer = gpt
+        self.pos_embedding = pos_emb
+        self.embeddings = embeddings
+        self.final_norm = norm
+        self.lm_head = nn.Sequential(norm, linear)
+        self.kv_cache = kv_cache
+
+    def store_prefix_emb(self, prefix_emb):
+        self.cached_prefix_emb = prefix_emb
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)  # usually None
+        if not self.kv_cache:
+            past_key_values = None
+
+        # only last token for inputs_ids if past is defined in kwargs
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values is not None:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        assert self.cached_prefix_emb is not None
+        assert inputs_embeds is None  # Not supported by this inference model.
+        assert labels is None  # Training not supported by this inference model.
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # assert len(past_key_values) + len(input_ids) == attention_mask.shape[1]
+
+        # Create embedding
+        prefix_len = self.cached_prefix_emb.shape[1]
+        if input_ids.shape[1] != 1:
+            # gen_inputs = input_ids[:, prefix_len:]
+            # gen_emb = self.embeddings(gen_inputs)
+            # gen_emb = gen_emb + self.pos_embedding(gen_emb)
+            gen_emb = self.cached_prefix_emb
+            if self.cached_prefix_emb.shape[0] != gen_emb.shape[0]:
+                prefix_emb = self.cached_prefix_emb.repeat_interleave(
+                    gen_emb.shape[0] // self.cached_prefix_emb.shape[0], 0
+                )
+            else:
+                prefix_emb = self.cached_prefix_emb.to(gen_emb.dtype)
+            # emb = torch.cat([prefix_emb, gen_emb], dim=1)
+            emb = gen_emb
+        else:
+            emb = self.embeddings(input_ids)
+            emb = emb + self.pos_embedding.get_fixed_embedding(
+                attention_mask.shape[1] - (prefix_len + 1), attention_mask.device
+            )
+        transformer_outputs = self.transformer(
+            inputs_embeds=emb,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+
+        if not return_dict:
+            return (lm_logits,) + transformer_outputs[1:]
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(
+                past_state.index_select(0, beam_idx.to(past_state.device))
+                for past_state in layer_past
+            )
+            for layer_past in past
+        )
+
+
+def null_position_embeddings(range, dim):
+    return torch.zeros((range.shape[0], range.shape[1], dim), device=range.device)
+
+
+class LearnedPositionEmbeddings(nn.Module):
+    def __init__(self, seq_len, model_dim, init=0.02, relative=False):
+        super().__init__()
+        # nn.Embedding
+        self.emb = torch.nn.Embedding(seq_len, model_dim)
+        # Initializing this way is standard for GPT-2
+        self.emb.weight.data.normal_(mean=0.0, std=init)
+        self.relative = relative
+        self.seq_len = seq_len
+
+    def forward(self, x):
+        sl = x.shape[1]
+        if self.relative:
+            start = random.randint(sl, self.seq_len) - sl
+            return self.emb(torch.arange(start, start + sl, device=x.device))
+        else:
+            return self.emb(torch.arange(0, sl, device=x.device))
+
+    def get_fixed_embedding(self, ind, dev):
+        return self.emb(torch.tensor([ind], device=dev)).unsqueeze(0)
+
+
+def build_hf_gpt_transformer(
+    layers,
+    model_dim,
+    heads,
+    max_mel_seq_len,
+    max_text_seq_len,
+    max_prompt_len,
+    checkpointing,
+):
+    """
+    GPT-2 implemented by the HuggingFace library.
+    """
+
+    gpt_config = GPT2Config(
+        vocab_size=256,  # Unused.
+        n_positions=max_mel_seq_len + max_text_seq_len + max_prompt_len,
+        n_ctx=max_mel_seq_len + max_text_seq_len + max_prompt_len,
+        n_embd=model_dim,
+        n_layer=layers,
+        n_head=heads,
+        gradient_checkpointing=checkpointing,
+        use_cache=not checkpointing,
+    )
+    gpt = GPT2Model(gpt_config)
+    # Override the built in positional embeddings
+    del gpt.wpe
+    gpt.wpe = functools.partial(null_position_embeddings, dim=model_dim)
+    # Built-in token embeddings are unused.
+    del gpt.wte
+
+    mel_pos_emb = (
+        LearnedPositionEmbeddings(max_mel_seq_len, model_dim)
+        if max_mel_seq_len != -1
+        else functools.partial(null_position_embeddings, dim=model_dim)
+    )
+    text_pos_emb = (
+        LearnedPositionEmbeddings(max_text_seq_len, model_dim)
+        if max_mel_seq_len != -1
+        else functools.partial(null_position_embeddings, dim=model_dim)
+    )
+    # gpt = torch.compile(gpt, mode="reduce-overhead", fullgraph=True)
+    return gpt, mel_pos_emb, text_pos_emb, None, None
+
+
+class Speech_LLM_GPT2(nn.Module):
+    def __init__(
+        self,
+        start_text_token,
+        stop_text_token,
+        num_text_tokens,
+        start_audio_token,
+        stop_audio_token,
+        num_audio_tokens,
+        llm_hidden_size,
+        llm_intermediate_size,
+        llm_num_layers,
+        llm_num_heads,
+        llm_max_audio_seq_len,
+        llm_max_text_seq_len,
+        llm_max_prompt_len,
+        code_stride_len=640,
+        max_conditioning_inputs=1,
+        label_smoothing=0.0,
+        checkpointing=False,
+    ):
+        """
+        Args:
+
+        """
+        super().__init__()
+
+        self.label_smoothing = label_smoothing
+        # text token config
+        self.start_text_token = start_text_token
+        self.stop_text_token = stop_text_token
+        self.num_text_tokens = num_text_tokens
+
+        # audio token config
+        self.start_audio_token = start_audio_token
+        self.stop_audio_token = stop_audio_token
+        self.num_audio_tokens = num_audio_tokens
+
+        # prompts token config
+        self.start_prompt_token = start_audio_token
+        self.stop_prompt_token = stop_audio_token
+
+        # other config
+        self.max_conditioning_inputs = max_conditioning_inputs
+
+        # length configs
+        self.max_text_len = llm_max_text_seq_len + 2  # add <bos> <eos>
+        self.max_prompt_len = llm_max_prompt_len
+        self.max_audio_len = llm_max_audio_seq_len + 2 + self.max_conditioning_inputs
+        self.max_gen_audio_tokens = (
+            llm_max_audio_seq_len - self.max_conditioning_inputs - 2
+        )
+        self.code_stride_len = code_stride_len
+
+        # model config
+        self.llm_hidden_size = llm_hidden_size
+        self.llm_intermediate_size = llm_intermediate_size
+        self.llm_num_layers = llm_num_layers
+        self.llm_num_heads = llm_num_heads
+
+        # text embedding and audio embeddings
+        self.text_embedding = nn.Embedding(self.num_text_tokens, self.llm_hidden_size)
+        self.audio_embedding = nn.Embedding(self.num_audio_tokens, self.llm_hidden_size)
+
+        # low-level llm model
+        self.gpt2, self.audio_pos_embedding, self.text_pos_embedding, _, _ = (
+            build_hf_gpt_transformer(
+                layers=self.llm_num_layers,
+                model_dim=self.llm_hidden_size,
+                heads=self.llm_num_heads,
+                max_mel_seq_len=self.max_audio_len,
+                max_text_seq_len=self.max_text_len,
+                max_prompt_len=self.max_prompt_len,
+                checkpointing=checkpointing,
+            )
+        )
+
+        # text and audio linear
+        self.final_norm = nn.LayerNorm(self.llm_hidden_size)
+        self.text_head = nn.Linear(self.llm_hidden_size, self.num_text_tokens)
+        self.audio_head = nn.Linear(self.llm_hidden_size, self.num_audio_tokens)
+
+        # speaker特征变换
+        self.reference_embedding = nn.Sequential(
+            nn.Linear(512, 256),
+            nn.Tanh(),
+            nn.Linear(256, self.llm_hidden_size),
+        )
+
+    def init_gpt_for_inference(self, kv_cache=True, use_deepspeed=False):
+        """_summary_
+
+        Args:
+            kv_cache (bool, optional): _description_. Defaults to True.
+            use_deepspeed (bool, optional): _description_. Defaults to False.
+        """
+        seq_length = self.max_audio_len + self.max_text_len + self.max_prompt_len + 1
+
+        gpt_config = GPT2Config(
+            vocab_size=self.num_audio_tokens,
+            n_positions=seq_length,
+            n_ctx=seq_length,
+            n_embd=self.llm_hidden_size,
+            n_layer=self.llm_num_layers,
+            n_head=self.llm_num_heads,
+            gradient_checkpointing=False,
+            use_cache=True,
+        )
+
+        # normal inference model
+        self.gpt_inference = GPT2InferenceModel(
+            config=gpt_config,
+            gpt=self.gpt2,
+            pos_emb=self.audio_pos_embedding,
+            embeddings=self.audio_embedding,
+            norm=self.final_norm,
+            linear=self.audio_head,
+            kv_cache=kv_cache,
+        )
+
+        # in-context inference model
+        self.gpt_inference_ic = GPT2ICInferenceModel(
+            config=gpt_config,
+            gpt=self.gpt2,
+            pos_emb=self.audio_pos_embedding,
+            embeddings=self.audio_embedding,
+            norm=self.final_norm,
+            linear=self.audio_head,
+            kv_cache=kv_cache,
+        )
+
+        self.gpt2.wte = self.audio_embedding
+
+    # --------------------------- normal inference ---------------------------
+    def inference(self, cond_latents, text_inputs, **hf_generate_kwargs):
+        self.compute_embeddings(cond_latents, text_inputs)
+        return self.generate(cond_latents, text_inputs, **hf_generate_kwargs)
+
+    def compute_embeddings(self, cond_latents, text_inputs):
+
+        text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token)
+        text_inputs = F.pad(text_inputs, (1, 0), value=self.start_text_token)
+        emb = self.text_embedding(text_inputs) + self.text_pos_embedding(text_inputs)
+        emb = torch.cat([cond_latents, emb], dim=1)
+        self.gpt_inference.store_prefix_emb(emb)
+        gpt_inputs = torch.full(
+            (
+                emb.shape[0],
+                emb.shape[1] + 1,  # +1 for the start_audio_token
+            ),
+            fill_value=1,
+            dtype=torch.long,
+            device=text_inputs.device,
+        )
+        gpt_inputs[:, -1] = self.start_audio_token
+        return gpt_inputs
+
+    def generate(self, cond_latents, text_inputs, **hf_generate_kwargs):
+        gpt_inputs = self.compute_embeddings(cond_latents, text_inputs)
+        gen = self.gpt_inference.generate(
+            gpt_inputs,
+            bos_token_id=self.start_audio_token,
+            pad_token_id=self.stop_audio_token,
+            eos_token_id=self.stop_audio_token,
+            max_length=self.max_gen_audio_tokens + gpt_inputs.shape[-1],
+            **hf_generate_kwargs,
+        )
+        if "return_dict_in_generate" in hf_generate_kwargs:
+            return gen.sequences[:, gpt_inputs.shape[1] :], gen
+        return gen[:, gpt_inputs.shape[1] :]
+
+    # --------------------------- normal inference --------------------------
+
+    # --------------------------- IC inference ---------------------------
+    def compute_embeddings_ic(self, cond_latents, text_inputs, prompt_tokens):
+        """_summary_
+
+        Args:
+            cond_latents (_type_): speaker embedding
+            text_inputs (_type_): text tokens
+            prompt_tokens (_type_): prompts_tokens
+
+        Returns:
+            _type_: _description_
+        """
+
+        # text embeddings
+        text_inputs = F.pad(text_inputs, (1, 0), value=self.start_text_token)
+        text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token)
+
+        text_emb = self.text_embedding(text_inputs) + self.text_pos_embedding(
+            text_inputs
+        )
+
+        # prompt_tokens
+
+        prompt_tokens = F.pad(prompt_tokens, (1, 0), value=self.start_audio_token)
+        audio_emb = self.audio_embedding(prompt_tokens) + self.audio_pos_embedding(
+            prompt_tokens
+        )
+
+        emb = torch.cat([cond_latents, text_emb, audio_emb], dim=1)
+
+        self.gpt_inference_ic.store_prefix_emb(emb)
+        gpt_inputs = torch.full(
+            (emb.shape[0], emb.shape[1]),
+            fill_value=1,
+            dtype=torch.long,
+            device=text_inputs.device,
+        )
+        return gpt_inputs
+
+    def generate_ic(
+        self, cond_latents, text_inputs, prompt_tokens, **hf_generate_kwargs
+    ):
+        """_summary_
+
+        Args:
+            cond_latents (_type_): _description_
+            text_inputs (_type_): _description_
+            prompt_tokens (_type_): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        gpt_inputs = self.compute_embeddings_ic(
+            cond_latents, text_inputs, prompt_tokens
+        )
+        gen = self.gpt_inference_ic.generate(
+            gpt_inputs,
+            bos_token_id=self.start_audio_token,
+            pad_token_id=self.stop_audio_token,
+            eos_token_id=self.stop_audio_token,
+            max_length=self.max_gen_audio_tokens + gpt_inputs.shape[-1],
+            **hf_generate_kwargs,
+        )
+        if "return_dict_in_generate" in hf_generate_kwargs:
+            return gen.sequences[:, gpt_inputs.shape[1] :], gen
+
+        return gen[:, gpt_inputs.shape[1] :]
+
+    # --------------------------- IC inference ---------------------------
+
+    def get_generator(self, fake_inputs, **hf_generate_kwargs):
+        return self.gpt_inference.generate_stream(
+            fake_inputs,
+            bos_token_id=self.start_audio_token,
+            pad_token_id=self.stop_audio_token,
+            eos_token_id=self.stop_audio_token,
+            max_length=self.max_gen_mel_tokens + fake_inputs.shape[-1],
+            do_stream=True,
+            **hf_generate_kwargs,
+        )

fireredtts/modules/semantic_tokenizer/__init__.py

@@ -0,0 +1,36 @@
+import os
+import torch
+
+from .hubert import HuBERT
+from .semantic_tokenizer import SemanticVQVAE
+
+
+class SemanticTokenizer:
+
+    def __init__(self, config, path):
+        self.model = SemanticVQVAE(**config)
+        self.model.load_state_dict(
+            torch.load(os.path.join(path, "codec.bin"), map_location="cpu"), strict=True
+        )
+
+        hubert = HuBERT(os.path.join(path, "hubert.pt"))
+        for name, param in hubert.named_parameters():
+            param.requires_grad = False
+        self.model.ssl_extractor = hubert
+
+        if torch.cuda.is_available():
+            self.model = self.model.cuda()
+        self.model.eval()
+
+    def __call__(self, wavs, wav_lengths):
+        tokens, token_lengths, spk_embeddings = self.extract(wavs, wav_lengths)
+        return tokens, token_lengths, spk_embeddings
+
+    def extract(self, wavs, wav_lengths):
+        saved_features = self.model.extract_speech_tokens(wavs, wav_lengths)
+
+        tokens = saved_features["token"]
+        token_lengths = saved_features["token_length"]
+        spk_embeddings = saved_features["spk"]
+
+        return tokens, token_lengths, spk_embeddings

fireredtts/modules/semantic_tokenizer/audio.py

@@ -0,0 +1,138 @@
+from librosa.filters import mel as librosa_mel_fn
+from torch import nn
+from torch.nn import functional as F
+
+import math
+import numpy as np
+import torch
+import torchaudio
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    return np.exp(x) / C
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+    output = dynamic_range_decompression_torch(magnitudes)
+    return output
+
+
+class TorchMelSpectrogram(nn.Module):
+
+    def __init__(
+        self,
+        filter_length=1024,
+        hop_length=200,
+        win_length=800,
+        n_mel_channels=80,
+        mel_fmin=0,
+        mel_fmax=8000,
+        sampling_rate=16000,
+        sampling_rate_org=None,
+        normalize=False,
+        mel_norm_file=None,
+        scale=1.0,
+        padding="center",
+        style="Tortoise",
+    ):
+
+        super().__init__()
+        self.style = style
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_mel_channels = n_mel_channels
+        self.mel_fmin = mel_fmin
+        self.mel_fmax = mel_fmax
+        self.sampling_rate = sampling_rate
+        self.sampling_rate_org = (
+            sampling_rate_org if sampling_rate_org is not None else sampling_rate
+        )
+
+        self.mel_basis = {}
+        self.hann_window = {}
+
+        self.scale = scale
+
+    def forward(self, inp, length=None):
+        if len(inp.shape) == 3:
+            inp = inp.squeeze(1) if inp.shape[1] == 1 else inp.squeeze(2)
+        assert len(inp.shape) == 2
+
+        if self.sampling_rate_org != self.sampling_rate:
+            inp = torchaudio.functional.resample(
+                inp, self.sampling_rate_org, self.sampling_rate
+            )
+
+        y = inp
+        if len(list(self.mel_basis.keys())) == 0:
+            mel = librosa_mel_fn(
+                sr=self.sampling_rate,
+                n_fft=self.filter_length,
+                n_mels=self.n_mel_channels,
+                fmin=self.mel_fmin,
+                fmax=self.mel_fmax,
+            )
+            self.mel_basis[str(self.mel_fmax) + "_" + str(y.device)] = (
+                torch.from_numpy(mel).float().to(y.device)
+            )
+            self.hann_window[str(y.device)] = torch.hann_window(self.win_length).to(
+                y.device
+            )
+
+        y = torch.nn.functional.pad(
+            y.unsqueeze(1),
+            (
+                int((self.filter_length - self.hop_length) / 2),
+                int((self.filter_length - self.hop_length) / 2),
+            ),
+            mode="reflect",
+        )
+        y = y.squeeze(1)
+
+        # complex tensor as default, then use view_as_real for future pytorch compatibility
+        spec = torch.stft(
+            y,
+            self.filter_length,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=self.hann_window[str(y.device)],
+            center=False,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=True,
+        )
+        spec = torch.view_as_real(spec)
+        spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+        spec = torch.matmul(
+            self.mel_basis[str(self.mel_fmax) + "_" + str(y.device)], spec
+        )
+        spec = spectral_normalize_torch(spec)
+
+        max_mel_length = math.ceil(y.shape[-1] / self.hop_length)
+        spec = spec[..., :max_mel_length].transpose(1, 2)
+
+        if length is None:
+            return spec
+        else:
+            spec_len = torch.ceil(length / self.hop_length).clamp(max=spec.shape[1])
+            return spec, spec_len

fireredtts/modules/semantic_tokenizer/ecapa_tdnn.py

@@ -0,0 +1,931 @@
+"""A popular speaker recognition and diarization model.
+
+Authors
+ * Hwidong Na 2020
+"""
+
+import math
+import os
+import torch  # noqa: F401
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def length_to_mask(length, max_len=None, dtype=None, device=None):
+    """Creates a binary mask for each sequence.
+
+    Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3
+
+    Arguments
+    ---------
+    length : torch.LongTensor
+        Containing the length of each sequence in the batch. Must be 1D.
+    max_len : int
+        Max length for the mask, also the size of the second dimension.
+    dtype : torch.dtype, default: None
+        The dtype of the generated mask.
+    device: torch.device, default: None
+        The device to put the mask variable.
+
+    Returns
+    -------
+    mask : tensor
+        The binary mask.
+
+    Example
+    -------
+    >>> length=torch.Tensor([1,2,3])
+    >>> mask=length_to_mask(length)
+    >>> mask
+    tensor([[1., 0., 0.],
+            [1., 1., 0.],
+            [1., 1., 1.]])
+    """
+    assert len(length.shape) == 1
+
+    if max_len is None:
+        max_len = length.max().long().item()  # using arange to generate mask
+    mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand(
+        len(length), max_len
+    ) < length.unsqueeze(1)
+
+    if dtype is None:
+        dtype = length.dtype
+
+    if device is None:
+        device = length.device
+
+    mask = torch.as_tensor(mask, dtype=dtype, device=device)
+    return mask
+
+
+def get_padding_elem(L_in: int, stride: int, kernel_size: int, dilation: int):
+    """This function computes the number of elements to add for zero-padding.
+
+    Arguments
+    ---------
+    L_in : int
+    stride: int
+    kernel_size : int
+    dilation : int
+    """
+    if stride > 1:
+        n_steps = math.ceil(((L_in - kernel_size * dilation) / stride) + 1)
+        L_out = stride * (n_steps - 1) + kernel_size * dilation
+        padding = [kernel_size // 2, kernel_size // 2]
+
+    else:
+        L_out = (L_in - dilation * (kernel_size - 1) - 1) // stride + 1
+
+        padding = [(L_in - L_out) // 2, (L_in - L_out) // 2]
+    return padding
+
+
+class Conv1d(nn.Module):
+    """This function implements 1d convolution.
+
+    Arguments
+    ---------
+    out_channels : int
+        It is the number of output channels.
+    kernel_size : int
+        Kernel size of the convolutional filters.
+    input_shape : tuple
+        The shape of the input. Alternatively use ``in_channels``.
+    in_channels : int
+        The number of input channels. Alternatively use ``input_shape``.
+    stride : int
+        Stride factor of the convolutional filters. When the stride factor > 1,
+        a decimation in time is performed.
+    dilation : int
+        Dilation factor of the convolutional filters.
+    padding : str
+        (same, valid, causal). If "valid", no padding is performed.
+        If "same" and stride is 1, output shape is the same as the input shape.
+        "causal" results in causal (dilated) convolutions.
+    padding_mode : str
+        This flag specifies the type of padding. See torch.nn documentation
+        for more information.
+    skip_transpose : bool
+        If False, uses batch x time x channel convention of speechbrain.
+        If True, uses batch x channel x time convention.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([10, 40, 16])
+    >>> cnn_1d = Conv1d(
+    ...     input_shape=inp_tensor.shape, out_channels=8, kernel_size=5
+    ... )
+    >>> out_tensor = cnn_1d(inp_tensor)
+    >>> out_tensor.shape
+    torch.Size([10, 40, 8])
+    """
+
+    def __init__(
+        self,
+        out_channels,
+        kernel_size,
+        input_shape=None,
+        in_channels=None,
+        stride=1,
+        dilation=1,
+        padding="same",
+        groups=1,
+        bias=True,
+        padding_mode="reflect",
+        skip_transpose=True,
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.padding = padding
+        self.padding_mode = padding_mode
+        self.unsqueeze = False
+        self.skip_transpose = skip_transpose
+
+        if input_shape is None and in_channels is None:
+            raise ValueError("Must provide one of input_shape or in_channels")
+
+        if in_channels is None:
+            in_channels = self._check_input_shape(input_shape)
+
+        self.conv = nn.Conv1d(
+            in_channels,
+            out_channels,
+            self.kernel_size,
+            stride=self.stride,
+            dilation=self.dilation,
+            padding=0,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        """Returns the output of the convolution.
+
+        Arguments
+        ---------
+        x : torch.Tensor (batch, time, channel)
+            input to convolve. 2d or 4d tensors are expected.
+        """
+
+        if not self.skip_transpose:
+            x = x.transpose(1, -1)
+
+        if self.unsqueeze:
+            x = x.unsqueeze(1)
+
+        if self.padding == "same":
+            x = self._manage_padding(x, self.kernel_size, self.dilation, self.stride)
+
+        elif self.padding == "causal":
+            num_pad = (self.kernel_size - 1) * self.dilation
+            x = F.pad(x, (num_pad, 0))
+
+        elif self.padding == "valid":
+            pass
+
+        else:
+            raise ValueError(
+                "Padding must be 'same', 'valid' or 'causal'. Got " + self.padding
+            )
+
+        wx = self.conv(x)
+
+        if self.unsqueeze:
+            wx = wx.squeeze(1)
+
+        if not self.skip_transpose:
+            wx = wx.transpose(1, -1)
+
+        return wx
+
+    def _manage_padding(
+        self,
+        x,
+        kernel_size: int,
+        dilation: int,
+        stride: int,
+    ):
+        """This function performs zero-padding on the time axis
+        such that their lengths is unchanged after the convolution.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Input tensor.
+        kernel_size : int
+            Size of kernel.
+        dilation : int
+            Dilation used.
+        stride : int
+            Stride.
+        """
+
+        # Detecting input shape
+        L_in = x.shape[-1]
+
+        # Time padding
+        padding = get_padding_elem(L_in, stride, kernel_size, dilation)
+
+        # Applying padding
+        x = F.pad(x, padding, mode=self.padding_mode)
+
+        return x
+
+    def _check_input_shape(self, shape):
+        """Checks the input shape and returns the number of input channels."""
+
+        if len(shape) == 2:
+            self.unsqueeze = True
+            in_channels = 1
+        elif self.skip_transpose:
+            in_channels = shape[1]
+        elif len(shape) == 3:
+            in_channels = shape[2]
+        else:
+            raise ValueError("conv1d expects 2d, 3d inputs. Got " + str(len(shape)))
+
+        # Kernel size must be odd
+        if self.kernel_size % 2 == 0:
+            raise ValueError(
+                "The field kernel size must be an odd number. Got %s."
+                % (self.kernel_size)
+            )
+        return in_channels
+
+
+class Fp32BatchNorm(nn.Module):
+    def __init__(self, sync=True, *args, **kwargs):
+        super().__init__()
+
+        if (
+            not torch.distributed.is_initialized()
+            or torch.distributed.get_world_size() == 1
+        ):
+            sync = False
+
+        if sync:
+            self.bn = nn.SyncBatchNorm(*args, **kwargs)
+        else:
+            self.bn = nn.BatchNorm1d(*args, **kwargs)
+
+        self.sync = sync
+
+    def forward(self, input):
+        if self.bn.running_mean.dtype != torch.float:
+            if self.sync:
+                self.bn.running_mean = self.bn.running_mean.float()
+                self.bn.running_var = self.bn.running_var.float()
+                if self.bn.affine:
+                    try:
+                        self.bn.weight = self.bn.weight.float()
+                        self.bn.bias = self.bn.bias.float()
+                    except:
+                        self.bn.float()
+            else:
+                self.bn.float()
+
+        output = self.bn(input.float())
+        return output.type_as(input)
+
+
+class BatchNorm1d(nn.Module):
+    """Applies 1d batch normalization to the input tensor.
+
+    Arguments
+    ---------
+    input_shape : tuple
+        The expected shape of the input. Alternatively, use ``input_size``.
+    input_size : int
+        The expected size of the input. Alternatively, use ``input_shape``.
+    eps : float
+        This value is added to std deviation estimation to improve the numerical
+        stability.
+    momentum : float
+        It is a value used for the running_mean and running_var computation.
+    affine : bool
+        When set to True, the affine parameters are learned.
+    track_running_stats : bool
+        When set to True, this module tracks the running mean and variance,
+        and when set to False, this module does not track such statistics.
+    combine_batch_time : bool
+        When true, it combines batch an time axis.
+
+
+    Example
+    -------
+    >>> input = torch.randn(100, 10)
+    >>> norm = BatchNorm1d(input_shape=input.shape)
+    >>> output = norm(input)
+    >>> output.shape
+    torch.Size([100, 10])
+    """
+
+    def __init__(
+        self,
+        input_shape=None,
+        input_size=None,
+        eps=1e-05,
+        momentum=0.1,
+        affine=True,
+        track_running_stats=True,
+        combine_batch_time=False,
+        skip_transpose=True,
+        enabled=True,
+    ):
+        super().__init__()
+        self.combine_batch_time = combine_batch_time
+        self.skip_transpose = skip_transpose
+
+        if input_size is None and skip_transpose:
+            input_size = input_shape[1]
+        elif input_size is None:
+            input_size = input_shape[-1]
+
+        if enabled:
+            self.norm = Fp32BatchNorm(
+                num_features=input_size,
+                eps=eps,
+                momentum=momentum,
+                affine=affine,
+                track_running_stats=track_running_stats,
+            )
+        else:
+            self.norm = nn.Identity()
+
+    def forward(self, x):
+        """Returns the normalized input tensor.
+
+        Arguments
+        ---------
+        x : torch.Tensor (batch, time, [channels])
+            input to normalize. 2d or 3d tensors are expected in input
+            4d tensors can be used when combine_dims=True.
+        """
+        shape_or = x.shape
+        if self.combine_batch_time:
+            if x.ndim == 3:
+                x = x.reshape(shape_or[0] * shape_or[1], shape_or[2])
+            else:
+                x = x.reshape(shape_or[0] * shape_or[1], shape_or[3], shape_or[2])
+
+        elif not self.skip_transpose:
+            x = x.transpose(-1, 1)
+
+        x_n = self.norm(x)
+
+        if self.combine_batch_time:
+            x_n = x_n.reshape(shape_or)
+        elif not self.skip_transpose:
+            x_n = x_n.transpose(1, -1)
+
+        return x_n
+
+
+class Linear(torch.nn.Module):
+    """Computes a linear transformation y = wx + b.
+
+    Arguments
+    ---------
+    n_neurons : int
+        It is the number of output neurons (i.e, the dimensionality of the
+        output).
+    bias : bool
+        If True, the additive bias b is adopted.
+    combine_dims : bool
+        If True and the input is 4D, combine 3rd and 4th dimensions of input.
+
+    Example
+    -------
+    >>> inputs = torch.rand(10, 50, 40)
+    >>> lin_t = Linear(input_shape=(10, 50, 40), n_neurons=100)
+    >>> output = lin_t(inputs)
+    >>> output.shape
+    torch.Size([10, 50, 100])
+    """
+
+    def __init__(
+        self,
+        n_neurons,
+        input_shape=None,
+        input_size=None,
+        bias=True,
+        combine_dims=False,
+    ):
+        super().__init__()
+        self.combine_dims = combine_dims
+
+        if input_shape is None and input_size is None:
+            raise ValueError("Expected one of input_shape or input_size")
+
+        if input_size is None:
+            input_size = input_shape[-1]
+            if len(input_shape) == 4 and self.combine_dims:
+                input_size = input_shape[2] * input_shape[3]
+
+        # Weights are initialized following pytorch approach
+        self.w = nn.Linear(input_size, n_neurons, bias=bias)
+
+    def forward(self, x):
+        """Returns the linear transformation of input tensor.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Input to transform linearly.
+        """
+        if x.ndim == 4 and self.combine_dims:
+            x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3])
+
+        wx = self.w(x)
+
+        return wx
+
+
+class TDNNBlock(nn.Module):
+    """An implementation of TDNN.
+
+    Arguments
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        The number of output channels.
+    kernel_size : int
+        The kernel size of the TDNN blocks.
+    dilation : int
+        The dilation of the Res2Net block.
+    activation : torch class
+        A class for constructing the activation layers.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2)
+    >>> layer = TDNNBlock(64, 64, kernel_size=3, dilation=1)
+    >>> out_tensor = layer(inp_tensor).transpose(1, 2)
+    >>> out_tensor.shape
+    torch.Size([8, 120, 64])
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        dilation,
+        activation=nn.ReLU,
+        batch_norm=True,
+    ):
+        super(TDNNBlock, self).__init__()
+        self.conv = Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            dilation=dilation,
+        )
+        self.activation = activation()
+        self.norm = BatchNorm1d(input_size=out_channels, enabled=batch_norm)
+
+    def forward(self, x):
+        return self.norm(self.activation(self.conv(x)))
+
+
+class Res2NetBlock(torch.nn.Module):
+    """An implementation of Res2NetBlock w/ dilation.
+
+    Arguments
+    ---------
+    in_channels : int
+        The number of channels expected in the input.
+    out_channels : int
+        The number of output channels.
+    scale : int
+        The scale of the Res2Net block.
+    kernel_size: int
+        The kernel size of the Res2Net block.
+    dilation : int
+        The dilation of the Res2Net block.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2)
+    >>> layer = Res2NetBlock(64, 64, scale=4, dilation=3)
+    >>> out_tensor = layer(inp_tensor).transpose(1, 2)
+    >>> out_tensor.shape
+    torch.Size([8, 120, 64])
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        scale=8,
+        kernel_size=3,
+        dilation=1,
+        batch_norm=True,
+    ):
+        super(Res2NetBlock, self).__init__()
+        assert in_channels % scale == 0
+        assert out_channels % scale == 0
+
+        in_channel = in_channels // scale
+        hidden_channel = out_channels // scale
+
+        self.blocks = nn.ModuleList(
+            [
+                TDNNBlock(
+                    in_channel,
+                    hidden_channel,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    batch_norm=batch_norm,
+                )
+                for i in range(scale - 1)
+            ]
+        )
+        self.scale = scale
+
+    def forward(self, x):
+        y = []
+        for i, x_i in enumerate(torch.chunk(x, self.scale, dim=1)):
+            if i == 0:
+                y_i = x_i
+            elif i == 1:
+                y_i = self.blocks[i - 1](x_i)
+            else:
+                y_i = self.blocks[i - 1](x_i + y_i)
+            y.append(y_i)
+        y = torch.cat(y, dim=1)
+        return y
+
+
+class SEBlock(nn.Module):
+    """An implementation of squeeze-and-excitation block.
+
+    Arguments
+    ---------
+    in_channels : int
+        The number of input channels.
+    se_channels : int
+        The number of output channels after squeeze.
+    out_channels : int
+        The number of output channels.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2)
+    >>> se_layer = SEBlock(64, 16, 64)
+    >>> lengths = torch.rand((8,))
+    >>> out_tensor = se_layer(inp_tensor, lengths).transpose(1, 2)
+    >>> out_tensor.shape
+    torch.Size([8, 120, 64])
+    """
+
+    def __init__(self, in_channels, se_channels, out_channels):
+        super(SEBlock, self).__init__()
+
+        self.conv1 = Conv1d(
+            in_channels=in_channels, out_channels=se_channels, kernel_size=1
+        )
+        self.relu = torch.nn.ReLU(inplace=True)
+        self.conv2 = Conv1d(
+            in_channels=se_channels, out_channels=out_channels, kernel_size=1
+        )
+        self.sigmoid = torch.nn.Sigmoid()
+
+    def forward(self, x, lengths=None):
+        L = x.shape[-1]
+        if lengths is not None:
+            mask = length_to_mask(lengths * L, max_len=L, device=x.device)
+            mask = mask.unsqueeze(1)
+            total = mask.sum(dim=2, keepdim=True)
+            s = (x * mask).sum(dim=2, keepdim=True) / total
+        else:
+            s = x.mean(dim=2, keepdim=True)
+
+        s = self.relu(self.conv1(s))
+        s = self.sigmoid(self.conv2(s))
+
+        return s * x
+
+
+class AttentiveStatisticsPooling(nn.Module):
+    """This class implements an attentive statistic pooling layer for each channel.
+    It returns the concatenated mean and std of the input tensor.
+
+    Arguments
+    ---------
+    channels: int
+        The number of input channels.
+    attention_channels: int
+        The number of attention channels.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2)
+    >>> asp_layer = AttentiveStatisticsPooling(64)
+    >>> lengths = torch.rand((8,))
+    >>> out_tensor = asp_layer(inp_tensor, lengths).transpose(1, 2)
+    >>> out_tensor.shape
+    torch.Size([8, 1, 128])
+    """
+
+    def __init__(
+        self, channels, attention_channels=128, global_context=True, batch_norm=True
+    ):
+        super().__init__()
+
+        self.eps = 1e-12
+        self.global_context = global_context
+        if global_context:
+            self.tdnn = TDNNBlock(
+                channels * 3, attention_channels, 1, 1, batch_norm=batch_norm
+            )
+        else:
+            self.tdnn = TDNNBlock(
+                channels, attention_channels, 1, 1, batch_norm, batch_norm
+            )
+        self.tanh = nn.Tanh()
+        self.conv = Conv1d(
+            in_channels=attention_channels, out_channels=channels, kernel_size=1
+        )
+
+    def forward(self, x, lengths=None):
+        """Calculates mean and std for a batch (input tensor).
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Tensor of shape [N, C, L].
+        """
+        L = x.shape[-1]
+
+        def _compute_statistics(x, m, dim=2, eps=self.eps):
+            mean = (m * x).sum(dim)
+            std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(eps))
+            return mean, std
+
+        if lengths is None:
+            lengths = torch.ones(x.shape[0], device=x.device)
+
+        # Make binary mask of shape [N, 1, L]
+        mask = length_to_mask(lengths * L, max_len=L, device=x.device)
+        mask = mask.unsqueeze(1)
+
+        # Expand the temporal context of the pooling layer by allowing the
+        # self-attention to look at global properties of the utterance.
+        if self.global_context:
+            # torch.std is unstable for backward computation
+            # https://github.com/pytorch/pytorch/issues/4320
+            total = mask.sum(dim=2, keepdim=True).float()
+            mean, std = _compute_statistics(x, mask / total)
+            mean = mean.unsqueeze(2).repeat(1, 1, L)
+            std = std.unsqueeze(2).repeat(1, 1, L)
+            attn = torch.cat([x, mean, std], dim=1)
+        else:
+            attn = x
+
+        # Apply layers
+        attn = self.conv(self.tanh(self.tdnn(attn)))
+
+        # Filter out zero-paddings
+        attn = attn.masked_fill(mask == 0, float("-inf"))
+
+        attn = F.softmax(attn, dim=2)
+        mean, std = _compute_statistics(x, attn)
+        # Append mean and std of the batch
+        pooled_stats = torch.cat((mean, std), dim=1)
+        pooled_stats = pooled_stats.unsqueeze(2)
+
+        return pooled_stats
+
+
+class SERes2NetBlock(nn.Module):
+    """An implementation of building block in ECAPA-TDNN, i.e.,
+    TDNN-Res2Net-TDNN-SEBlock.
+
+    Arguments
+    ----------
+    out_channels: int
+        The number of output channels.
+    res2net_scale: int
+        The scale of the Res2Net block.
+    kernel_size: int
+        The kernel size of the TDNN blocks.
+    dilation: int
+        The dilation of the Res2Net block.
+    activation : torch class
+        A class for constructing the activation layers.
+
+    Example
+    -------
+    >>> x = torch.rand(8, 120, 64).transpose(1, 2)
+    >>> conv = SERes2NetBlock(64, 64, res2net_scale=4)
+    >>> out = conv(x).transpose(1, 2)
+    >>> out.shape
+    torch.Size([8, 120, 64])
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        res2net_scale=8,
+        se_channels=128,
+        kernel_size=1,
+        dilation=1,
+        activation=torch.nn.ReLU,
+        batch_norm=True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.tdnn1 = TDNNBlock(
+            in_channels,
+            out_channels,
+            kernel_size=1,
+            dilation=1,
+            activation=activation,
+            batch_norm=batch_norm,
+        )
+        self.res2net_block = Res2NetBlock(
+            out_channels,
+            out_channels,
+            res2net_scale,
+            kernel_size,
+            dilation,
+            batch_norm=batch_norm,
+        )
+        self.tdnn2 = TDNNBlock(
+            out_channels,
+            out_channels,
+            kernel_size=1,
+            dilation=1,
+            activation=activation,
+            batch_norm=batch_norm,
+        )
+        self.se_block = SEBlock(out_channels, se_channels, out_channels)
+
+        self.shortcut = None
+        if in_channels != out_channels:
+            self.shortcut = Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+            )
+
+    def forward(self, x, lengths=None):
+        residual = x
+        if self.shortcut:
+            residual = self.shortcut(x)
+
+        x = self.tdnn1(x)
+        x = self.res2net_block(x)
+        x = self.tdnn2(x)
+        x = self.se_block(x, lengths)
+
+        return x + residual
+
+
+class ECAPA_TDNN(torch.nn.Module):
+    """An implementation of the speaker embedding model in a paper.
+    "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in
+    TDNN Based Speaker Verification" (https://arxiv.org/abs/2005.07143).
+
+    Arguments
+    ---------
+    device : str
+        Device used, e.g., "cpu" or "cuda".
+    activation : torch class
+        A class for constructing the activation layers.
+    channels : list of ints
+        Output channels for TDNN/SERes2Net layer.
+    kernel_sizes : list of ints
+        List of kernel sizes for each layer.
+    dilations : list of ints
+        List of dilations for kernels in each layer.
+    lin_neurons : int
+        Number of neurons in linear layers.
+
+    Example
+    -------
+    >>> input_feats = torch.rand([5, 120, 80])
+    >>> compute_embedding = ECAPA_TDNN(80, lin_neurons=192)
+    >>> outputs = compute_embedding(input_feats)
+    >>> outputs.shape
+    torch.Size([5, 1, 192])
+    """
+
+    def __init__(
+        self,
+        input_size,
+        lin_neurons=192,
+        activation=torch.nn.ReLU,
+        channels=[512, 512, 512, 512, 1536],
+        kernel_sizes=[5, 3, 3, 3, 1],
+        dilations=[1, 2, 3, 4, 1],
+        attention_channels=128,
+        res2net_scale=8,
+        se_channels=128,
+        global_context=True,
+        batch_norm=True,
+    ):
+
+        super().__init__()
+        assert len(channels) == len(kernel_sizes)
+        assert len(channels) == len(dilations)
+        self.channels = channels
+        self.blocks = nn.ModuleList()
+
+        # The initial TDNN layer
+        self.blocks.append(
+            TDNNBlock(
+                input_size,
+                channels[0],
+                kernel_sizes[0],
+                dilations[0],
+                activation,
+                batch_norm=batch_norm,
+            )
+        )
+
+        # SE-Res2Net layers
+        for i in range(1, len(channels) - 1):
+            self.blocks.append(
+                SERes2NetBlock(
+                    channels[i - 1],
+                    channels[i],
+                    res2net_scale=res2net_scale,
+                    se_channels=se_channels,
+                    kernel_size=kernel_sizes[i],
+                    dilation=dilations[i],
+                    activation=activation,
+                    batch_norm=batch_norm,
+                )
+            )
+
+        # Multi-layer feature aggregation
+        self.mfa = TDNNBlock(
+            channels[-1],
+            channels[-1],
+            kernel_sizes[-1],
+            dilations[-1],
+            activation,
+            batch_norm=batch_norm,
+        )
+
+        # Attentive Statistical Pooling
+        self.asp = AttentiveStatisticsPooling(
+            channels[-1],
+            attention_channels=attention_channels,
+            global_context=global_context,
+            batch_norm=batch_norm,
+        )
+        self.asp_bn = BatchNorm1d(input_size=channels[-1] * 2, enabled=batch_norm)
+
+        # Final linear transformation
+        self.fc = Conv1d(
+            in_channels=channels[-1] * 2,
+            out_channels=lin_neurons,
+            kernel_size=1,
+        )
+
+    # @torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
+    def forward(self, x, lengths=None):
+        """Returns the embedding vector.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Tensor of shape (batch, time, channel).
+        """
+        # Minimize transpose for efficiency
+        x = x.transpose(1, 2)
+
+        xl = []
+        for layer in self.blocks:
+            try:
+                x = layer(x, lengths=lengths)
+            except TypeError:
+                x = layer(x)
+            xl.append(x)
+
+        # Multi-layer feature aggregation
+        x = torch.cat(xl[1:], dim=1)
+        x = self.mfa(x)
+
+        # Attentive Statistical Pooling
+        x = self.asp(x, lengths=lengths)
+        x = self.asp_bn(x)
+
+        # Final linear transformation
+        x = self.fc(x)
+
+        x = x.squeeze(-1)
+        return x
+
+
+if __name__ == "__main__":
+    model = ECAPA_TDNN(128, batch_norm=False)
+    # print(model)

fireredtts/modules/semantic_tokenizer/hubert.py

@@ -0,0 +1,108 @@
+from fairseq import checkpoint_utils
+from torch.nn.utils.rnn import pad_sequence
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_mask_from_lengths(lengths, max_len=None):
+    max_len = torch.max(lengths).item() if max_len is None else max_len
+    ids = torch.arange(0, max_len).to(lengths.device)
+    mask = ~(ids < lengths.unsqueeze(1)).bool()
+    return mask
+
+
+class HuBERT(nn.Module):
+
+    def __init__(self, model_path, sampling_rate=16000):
+
+        super().__init__()
+
+        models, saved_cfg, _ = checkpoint_utils.load_model_ensemble_and_task(
+            [model_path],
+            suffix="",
+        )
+
+        model = models[0]
+        model = model.half()
+        model.eval()
+        self.model = model
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+        self.sampling_rate = sampling_rate
+        self.normalize = saved_cfg.task.normalize
+
+    @torch.no_grad()
+    @torch.cuda.amp.autocast(enabled=False, dtype=torch.float16)
+    def forward(self, inp, length=None, split=True, split_size=4):
+        self.model.eval()
+        if self.training and split:
+            split_size = int(math.ceil(inp.shape[0] / 4))
+            outs, out_lens = [], []
+            for i in range(0, inp.shape[0], split_size):
+                inp_, length_ = inp[i : i + split_size], length[i : i + split_size]
+                out_, out_len_ = self._extract(inp_, length_)
+                outs.append(out_)
+                out_lens.append(out_len_)
+            max_length = max([max(ols) for ols in out_lens])
+
+            return torch.cat(
+                [F.pad(o, (0, 0, 0, max_length - o.shape[1]), value=0) for o in outs],
+                dim=0,
+            ), torch.cat(out_lens, dim=0)
+        else:
+            return self._extract(inp, length)
+
+    @torch.no_grad()
+    def _extract(self, inp, length):
+        frame_samples = int(self.sampling_rate * 0.02)
+        device = inp.device
+
+        if len(inp.shape) == 3:
+            inp = inp.squeeze(1) if inp.shape[1] == 1 else inp.squeeze(2)
+        assert len(inp.shape) == 2
+        assert self.sampling_rate == 16000
+
+        feats = inp
+
+        # Padding with 0
+        padding_size = 3200  # Longer to cover receptive field
+        feats = F.pad(feats, (0, padding_size), mode="constant", value=0)
+
+        # Norm volume using LN
+        feats = self._postprocess(
+            feats, length + padding_size, normalize=self.normalize
+        )
+
+        if length is None:
+            padding_mask = torch.BoolTensor(feats.shape).fill_(False)
+        else:
+            length = torch.ceil(length / 320).int()
+            padding_mask = get_mask_from_lengths(length).bool()
+            padding_mask = F.pad(padding_mask, (0, 9), value=True)
+
+        inputs = {
+            "source": feats.half().to(device),
+            "padding_mask": padding_mask.to(device),
+            "mask": False,
+        }
+        logits, _ = self.model.extract_features(**inputs)
+        logits = logits[:, : length.max()].float()
+
+        return logits, length
+
+    def _postprocess(self, feats, lengths, normalize=False):
+        assert feats.dim() == 2, feats.dim()
+
+        if normalize:
+            with torch.no_grad():
+                feats = [
+                    F.layer_norm(feat[:length], feat[:length].shape)
+                    for feat, length in zip(feats, lengths)
+                ]
+                feats = pad_sequence(feats, batch_first=True, padding_value=0)
+        return feats

fireredtts/modules/semantic_tokenizer/semantic_tokenizer.py

@@ -0,0 +1,877 @@
+from functools import reduce
+from tokenize import Triple
+from torch import nn
+from torch.autograd import Function
+from torch.nn import functional as F
+from torch.nn.utils import spectral_norm, weight_norm
+from torch.utils.checkpoint import checkpoint
+
+import einops
+import math
+import numpy as np
+import os
+import random
+import torch
+import torchaudio
+import typing as tp
+import warnings
+
+from .audio import TorchMelSpectrogram
+from .ecapa_tdnn import ECAPA_TDNN
+from .hubert import HuBERT
+from ..acoustic_codec.vector_quantization import VectorQuantization
+
+
+CONV_NORMALIZATIONS = frozenset(
+    [
+        "none",
+        "weight_norm",
+        "spectral_norm",
+        "time_layer_norm",
+        "layer_norm",
+        "time_group_norm",
+    ]
+)
+NORM = "weight_norm"
+
+
+def get_mask_from_lengths(lengths, max_len=None):
+    max_len = torch.max(lengths).item() if max_len is None else max_len
+    ids = torch.arange(0, max_len).to(lengths.device)
+    mask = ~(ids < lengths.unsqueeze(1)).bool()
+    return mask
+
+
+class ConvLayerNorm(nn.LayerNorm):
+    def __init__(
+        self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs
+    ):
+        super().__init__(normalized_shape, **kwargs)
+
+    def forward(self, x):
+        x = einops.rearrange(x, "b ... t -> b t ...")
+        x = super().forward(x)
+        x = einops.rearrange(x, "b t ... -> b ... t")
+        return
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == "weight_norm":
+        return weight_norm(module)
+    elif norm == "spectral_norm":
+        return spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(
+    module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs
+) -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == "layer_norm":
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return ConvLayerNorm(module.out_channels, **norm_kwargs)
+    elif norm == "time_group_norm":
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(
+    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
+) -> int:
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(
+    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
+):
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+def pad1d(
+    x: torch.Tensor,
+    paddings: tp.Tuple[int, int],
+    mode: str = "zero",
+    value: float = 0.0,
+):
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == "reflect":
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left:end]
+
+
+class NormConv1d(nn.Module):
+
+    def __init__(
+        self,
+        *args,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConv2d(nn.Module):
+
+    def __init__(
+        self,
+        *args,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+
+    def __init__(
+        self,
+        *args,
+        causal: bool = False,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(
+            nn.ConvTranspose1d(*args, **kwargs), norm
+        )
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose2d(nn.Module):
+
+    def __init__(
+        self,
+        *args,
+        norm: str = "none",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        **kwargs,
+    ):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(
+            nn.ConvTranspose2d(*args, **kwargs), norm
+        )
+        self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class SConv1d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causal: bool = False,
+        norm: str = "weight_norm",
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+        pad_mode: str = "reflect",
+    ):
+        super().__init__()
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn(
+                "SConv1d has been initialized with stride > 1 and dilation > 1"
+                f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
+            )
+        self.conv = NormConv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            causal=causal,
+            norm=norm,
+            norm_kwargs=norm_kwargs,
+        )
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        B, C, T = x.shape
+        kernel_size = self.conv.conv.kernel_size[0]
+        stride = self.conv.conv.stride[0]
+        dilation = self.conv.conv.dilation[0]
+        kernel_size = (
+            kernel_size - 1
+        ) * dilation + 1  # effective kernel size with dilations
+        padding_total = kernel_size - stride
+        extra_padding = get_extra_padding_for_conv1d(
+            x, kernel_size, stride, padding_total
+        )
+        if self.causal:
+            # Left padding for causal
+            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(
+                x, (padding_left, padding_right + extra_padding), mode=self.pad_mode
+            )
+        return self.conv(x)
+
+
+class SConvTranspose1d(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        causal: bool = False,
+        norm: str = "weight_norm",
+        trim_right_ratio: float = 1.0,
+        norm_kwargs: tp.Dict[str, tp.Any] = {},
+    ):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            causal=causal,
+            norm=norm,
+            norm_kwargs=norm_kwargs,
+        )
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert (
+            self.causal or self.trim_right_ratio == 1.0
+        ), "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0.0 and self.trim_right_ratio <= 1.0
+
+    def forward(self, x):
+        kernel_size = self.convtr.convtr.kernel_size[0]
+        stride = self.convtr.convtr.stride[0]
+        padding_total = kernel_size - stride
+
+        y = self.convtr(x)
+
+        # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+        # removed at the very end, when keeping only the right length for the output,
+        # as removing it here would require also passing the length at the matching layer
+        # in the encoder.
+        if self.causal:
+            # Trim the padding on the right according to the specified ratio
+            # if trim_right_ratio = 1.0, trim everything from right
+            padding_right = math.ceil(padding_total * self.trim_right_ratio)
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        return y
+
+
+class SLSTM(nn.Module):
+
+    def __init__(
+        self,
+        dimension: int,
+        num_layers: int = 2,
+        bidirectional: bool = False,
+        skip: bool = True,
+    ):
+        super().__init__()
+        self.bidirectional = bidirectional
+        self.skip = skip
+        if bidirectional:
+            self.lstm = nn.LSTM(
+                dimension, dimension // 2, num_layers, bidirectional=bidirectional
+            )
+        else:
+            self.lstm = nn.LSTM(dimension, dimension, num_layers)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        y, _ = self.lstm(x)
+        if self.skip:
+            y = y + x
+        y = y.permute(1, 2, 0)
+        return y
+
+
+class Swish(nn.Module):
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, groups=1):
+        super().__init__()
+
+        self.layers = nn.Sequential(
+            SConv1d(
+                in_channels=in_channels,
+                out_channels=out_channels // 2,
+                kernel_size=kernel_size,
+                groups=groups,
+                norm=NORM,
+            ),
+            Swish(),
+            SConv1d(
+                in_channels=out_channels // 2,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                groups=groups,
+                norm=NORM,
+            ),
+        )
+
+    def forward(self, x):
+        return x + self.layers(x)
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, out_channels, stride):
+        super().__init__()
+
+        self.layers = nn.Sequential(
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels),
+            Swish(),
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels),
+            Swish(),
+            SConv1d(
+                in_channels=out_channels,
+                out_channels=out_channels,
+                kernel_size=2 * stride,
+                stride=stride,
+                norm=NORM,
+            ),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, in_channels, stride):
+        super().__init__()
+        out_channels = in_channels
+        self.layers = nn.Sequential(
+            SConvTranspose1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=2 * stride,
+                stride=stride,
+                norm=NORM,
+            ),
+            Swish(),
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels),
+            Swish(),
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class Encoder(nn.Module):
+    def __init__(self, C, D, strides=[2, 2], checkpointing=True):
+        super().__init__()
+        self.checkpointing = checkpointing
+
+        self.downsample_scale = np.cumprod(np.asarray(strides))[-1]
+        self.layers = [
+            SConv1d(in_channels=C, out_channels=D, kernel_size=3, norm=NORM),
+            Swish(),
+        ]
+        for stride in strides:
+            self.layers += [
+                EncoderBlock(out_channels=D, stride=stride),
+                Swish(),
+            ]
+        self.layers += [
+            SConv1d(in_channels=D, out_channels=D, kernel_size=3, norm=NORM),
+            SLSTM(D, num_layers=1, bidirectional=True),
+        ]
+        self.layers = nn.Sequential(*self.layers)
+
+    def forward(self, x):
+        if self.checkpointing:
+            x = checkpoint(
+                self.layers, x.transpose(1, 2), use_reentrant=False
+            ).transpose(1, 2)
+        else:
+            x = self.layers(x.transpose(1, 2)).transpose(1, 2)
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(self, C, D, H, strides=[2, 2], checkpointing=True):
+        super().__init__()
+        self.checkpointing = checkpointing
+
+        self.in_layer = nn.Sequential(
+            SConv1d(in_channels=D, out_channels=H, kernel_size=3, norm=NORM),
+            SLSTM(H, num_layers=1, bidirectional=True),
+        )
+        self.layers = nn.ModuleList()
+        for stride in strides:
+            self.layers.append(
+                nn.Sequential(DecoderBlock(in_channels=H, stride=stride), Swish())
+            )
+        self.out_layer = SConv1d(
+            in_channels=H, out_channels=C, kernel_size=3, norm=NORM
+        )
+
+    def forward(self, x, g=None):
+        if self.checkpointing:
+            y = checkpoint(self._forward, x, g, use_reentrant=False)
+        else:
+            y = self._forward(x, g)
+        return y
+
+    def _forward(self, x, g=None):
+        h = self.in_layer(x.transpose(1, 2))
+
+        for layer in self.layers:
+            up_g = g.unsqueeze(-1).repeat(1, 1, h.shape[-1])
+            h = h + up_g
+            h = layer(h)
+
+        y = self.out_layer(h)
+
+        return y.transpose(1, 2), h.transpose(1, 2)
+
+
+class TimeRegulator(nn.Module):
+
+    def __init__(self, in_dim, scale, learnable=False):
+        super().__init__()
+        self.scale = scale
+        self.learnable = learnable
+
+    def forward(self, x, x_len, downsample=True):
+        if downsample:
+            x = self.downsample(x, x_len)
+        else:
+            x = self.upsample(x, x_len)
+        return x
+
+    def downsample(self, x, x_len):
+        x = torch.nn.functional.avg_pool1d(
+            x.transpose(1, 2), self.scale, stride=self.scale, ceil_mode=True
+        ).transpose(1, 2)
+        x_len = (x_len / self.scale).ceil()
+        return x, x_len
+
+    def upsample(self, x, x_len):
+        if self.learnable:
+            x = self.upsampler(x.transpose(1, 2)).transpose(1, 2)
+        else:
+            x = torch.repeat_interleave(x, self.scale, dim=1)
+        return x
+
+
+class TreeVectorQuantization(nn.Module):
+
+    def __init__(
+        self,
+        in_dim,
+        vq_class="VectorQuantization",
+        vq_config={},
+        tree_config={},
+    ):
+        super().__init__()
+        self.vq_config = vq_config
+        self.tree_config = tree_config
+
+        self.quantizers = nn.ModuleList()
+        self.time_regulators = nn.ModuleList()
+        for config in self.tree_config:
+            vq_config = self.vq_config.copy()
+            if not isinstance(vq_config["codebook_size"], (tuple, list)):
+                vq_config["codebook_size"] = [vq_config["codebook_size"]]
+                vq_config["codebook_dim"] = [vq_config["codebook_dim"]]
+            vq_config["codebook_size"] = vq_config["codebook_size"] * config["n_groups"]
+            vq_config["codebook_dim"] = vq_config["codebook_dim"] * config["n_groups"]
+            self.quantizers.append(
+                VectorQuantization(
+                    in_dim,
+                    n_groups=config.get("n_groups", 1),
+                    dropout_rate_per_group=config.get("dropout_rate_per_group", 0),
+                    ordered=config.get("ordered", False),
+                    **vq_config,
+                )
+            )
+            self.time_regulators.append(
+                TimeRegulator(
+                    in_dim,
+                    config["downsample_rate"],
+                    config.get("learnable_time_regulator", False),
+                )
+            )
+
+    def forward(
+        self, inp, inp_len, enable_vq=True, update_codebook=True, return_pre_quant=False
+    ):
+        output, (quants, losses, embed_inds) = self.quantize(
+            inp,
+            inp_len,
+            enable_vq=enable_vq,
+            update_codebook=update_codebook,
+            return_pre_quant=return_pre_quant,
+        )
+        loss = sum(losses) / len(losses)
+        return output, (quants, loss, embed_inds)
+
+    def quantize(
+        self, inp, inp_len, enable_vq=True, update_codebook=True, return_pre_quant=False
+    ):
+        quants, losses, embed_inds = [], [], []
+
+        pre_quant_output, quant_output, residual = 0, 0, inp
+        for tree_config, quantizer, regulator in zip(
+            self.tree_config, self.quantizers, self.time_regulators
+        ):
+            # Downsample
+            x, x_len = regulator(residual, inp_len, True)
+
+            # Quantization
+            q, diff, embed_ind = quantizer(
+                x,
+                x_len,
+                enable_vq=enable_vq,
+                update_codebook=update_codebook,
+                return_pre_quant=return_pre_quant,
+            )
+            if return_pre_quant:
+                pq, q = q
+
+            # Upsample
+            x = regulator(q, x_len, False)[:, : residual.shape[1]]
+
+            residual = residual - x
+            quant_output = quant_output + x
+
+            if return_pre_quant:
+                pq = regulator(pq, x_len, False)[:, : residual.shape[1]]
+                pre_quant_output = pre_quant_output + pq
+
+            quants.append(q)
+            losses.append(diff)
+            embed_inds.append(embed_ind)
+
+        if return_pre_quant:
+            return (pre_quant_output, quant_output), (quants, losses, embed_inds)
+        return quant_output, (quants, losses, embed_inds)
+
+    def decode(self, seqs, seq_lens=None):
+        if not isinstance(seqs, (tuple, list)):
+            tokens, token_lens = self.deserialize(seqs, seq_lens)
+        else:
+            tokens, token_lens = seqs, seq_lens
+
+        quant_output = 0
+        for token, quantizer, regulator in zip(
+            tokens, self.quantizers, self.time_regulators
+        ):
+            x = quantizer.decode(token).transpose(1, 2)
+            x = regulator(x, None, False)
+            if torch.is_tensor(quant_output):
+                x = x[:, : quant_output.size(1)]
+            quant_output = quant_output + x
+
+        return quant_output, token_lens
+
+    def serialize(self, tokens, token_lens):
+        assert len(tokens) <= 2, "we only support 1 or 2-scale sequences now..."
+
+        scale = self.tree_config[0]["downsample_rate"]
+        token_lens = ((token_lens.float() / scale).ceil() * scale).int()
+
+        seq1 = tokens[0].unsqueeze(-1)
+
+        if len(tokens) == 1:
+            seq_cat = seq1.view(seq1.shape[0], -1)
+            seq_cat_lens = (token_lens / scale * seq1.shape[2]).int()
+        elif len(tokens) == 2:
+            seq2 = F.pad(
+                tokens[1], (0, token_lens.max() - tokens[1].size(1)), "replicate"
+            )
+            seq2 = torch.stack([seq2[:, i::scale] for i in range(scale)], dim=-1)
+            seq_cat = torch.cat((seq1, seq2), dim=-1).view(seq1.shape[0], -1)
+            seq_cat_lens = (token_lens / scale + token_lens).int()
+
+        return seq_cat, seq_cat_lens
+
+    def deserialize(self, seqs, seq_lens):
+        if len(self.tree_config) == 1:
+            return [seqs], seq_lens
+
+        max_scale = max(config["downsample_rate"] for config in self.tree_config)
+        total_scale = sum(config["downsample_rate"] for config in self.tree_config)
+
+        # Cut for aligning
+        if seq_lens is None:
+            seq_lens = torch.full([seqs.shape[0]], seqs.shape[1]).to(seqs.device)
+        seq_lens = (seq_lens / total_scale).int() * total_scale
+        token_lens = (seq_lens / total_scale).int() * max_scale
+        seqs = seqs[:, : seq_lens.max()]
+
+        # Separate
+        tokens = torch.stack(
+            [seqs[:, i::total_scale] for i in range(total_scale)], dim=-1
+        )
+        seq1 = tokens[..., 0]
+        seq2 = tokens[..., 1:].contiguous().view(tokens.shape[0], -1)
+
+        return [seq1, seq2], token_lens
+
+
+class SemanticVQVAE(nn.Module):
+
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        n_model_size,
+        downsample_scales=[1, 2],
+        upsample_scales=[[2, 1], [2, 1]],
+        mel_config={},
+        ssl_config={},
+        # Quantization
+        vq_class="VectorQuantization",
+        vq_config={},
+        tree_config={},
+        # Training
+        checkpointing=True,
+        dual_decoding=False,
+        n_samples_per_token=640,
+        online_extraction=True,
+        ssl_extractor=None,
+    ):
+        super(SemanticVQVAE, self).__init__()
+        self.in_dim = in_dim
+        self.n_model_size = n_model_size
+        self.mel_config = mel_config
+        self.dual_decoding = dual_decoding
+        self.vq_config = vq_config
+        self.tree_config = tree_config
+        self.output_feature = "mel"
+        self.n_samples_per_token = n_samples_per_token
+        self.checkpointing = checkpointing
+
+        self.mel_spectrogram = TorchMelSpectrogram(**mel_config)
+
+        # Speaker encoder
+        self.speaker_encoder = ECAPA_TDNN(
+            out_dim,
+            n_model_size,
+            channels=[512, 512, 512, 512, 1536],
+            kernel_sizes=[5, 3, 3, 3, 1],
+            dilations=[1, 2, 3, 4, 1],
+            attention_channels=128,
+            res2net_scale=4,
+            se_channels=128,
+            global_context=True,
+            batch_norm=True,
+        )
+
+        # Encoder & decoder
+        self.encoder = Encoder(
+            in_dim, n_model_size, downsample_scales, checkpointing=checkpointing
+        )
+
+        # Quantization
+        self.quantizer = TreeVectorQuantization(
+            n_model_size,
+            vq_class=vq_class,
+            vq_config=vq_config,
+            tree_config=tree_config,
+        )
+
+    def forward(
+        self,
+        wav,
+        wav_length,
+        enable_vq=True,
+        decode=True,
+        extract_spk=True,
+        shuffle=False,
+        **kwargs,
+    ):
+        output_dict = {}
+
+        with torch.no_grad():
+            # Pad waveform
+            if wav.shape[1] % self.n_samples_per_token > 0:
+                pad_size = (
+                    self.n_samples_per_token - wav.shape[1] % self.n_samples_per_token
+                )
+                wav = F.pad(wav, (0, pad_size), value=0)
+                wav_length += pad_size
+
+            # Extract mel & sll
+            mel, mel_length = kwargs.get("mel", None), kwargs.get("mel_length", None)
+            if mel is None:
+                mel, mel_length = self.mel_spectrogram(wav, wav_length)
+            output_dict.update({"mel": mel, "mel_length": mel_length})
+
+            ssl, ssl_length = kwargs.get("ssl", None), kwargs.get("ssl_length", None)
+            if ssl is None:
+                ssl, ssl_length = self.ssl_extractor(wav, wav_length)
+            output_dict.update({"ssl": ssl.float(), "ssl_length": ssl_length})
+
+        input, input_length = ssl, ssl_length
+        output, output_length = mel, mel_length
+
+        encoder_outputs = self.encoder(input)
+        quant_length = torch.ceil(input_length / self.encoder.downsample_scale)
+        quant_length = quant_length.clamp(max=encoder_outputs.shape[1])
+
+        quant, (quants, diff, embed_ind) = self.quantizer(
+            encoder_outputs,
+            quant_length,
+            enable_vq=enable_vq,
+            update_codebook=True,
+            return_pre_quant=self.dual_decoding,
+        )
+
+        output_dict.update(
+            {
+                "quants": quants,
+                "token": embed_ind,
+                "token_length": quant_length.int(),
+                "encoder_diffs": diff,
+            }
+        )
+
+        # Speaker
+        if extract_spk:
+            cond, cond_length = output, output_length
+            speaker_embedding = self.speaker_encoder(cond, cond_length)
+            speaker_embedding_1 = speaker_embedding_2 = speaker_embedding
+            output_dict["spk"] = speaker_embedding
+
+        return output_dict
+
+    @torch.no_grad()
+    def extract_speech_tokens(
+        self, wav, wav_length, serialize=True, extract_spk=True, shuffle=False
+    ):
+        output_dict = self.forward(
+            wav, wav_length, True, False, extract_spk=extract_spk, shuffle=shuffle
+        )
+        token_seqs, token_length = output_dict["token"], output_dict["token_length"]
+
+        # Align sequences
+        scale = self.tree_config[0]["downsample_rate"]
+        token_length = (torch.ceil(token_length / scale) * scale).int()
+
+        new_token_seqs, new_token_lens = [], []
+        for i, token_seq in enumerate(token_seqs):
+            # discrete-continuous tokens
+            residual = None
+            if isinstance(token_seq, (tuple, list)):
+                token_seq, residual = token_seq
+
+            scale = self.tree_config[i]["downsample_rate"]
+            new_token_len = token_length // scale
+            pad = int(new_token_len.max()) - token_seq.shape[1]
+            token_seq = F.pad(
+                token_seq,
+                (0, pad) if len(token_seq.shape) == 2 else (0, 0, 0, pad),
+                "replicate",
+            )
+
+            if residual is not None:
+                token_seq = (token_seq, residual)
+            new_token_seqs.append(token_seq)
+            new_token_lens.append(new_token_len)
+
+        if len(new_token_seqs) == 1:
+            new_token_seqs, new_token_lens = new_token_seqs[0], new_token_lens[0]
+        elif serialize:
+            new_token_seqs, new_token_lens = self.quantizer.serialize(
+                new_token_seqs, new_token_lens
+            )
+
+        output_dict.update(
+            {
+                "embed": output_dict["quants"],
+                "token": new_token_seqs,
+                "token_length": new_token_lens,
+            }
+        )
+
+        return output_dict
+
+    @torch.no_grad()
+    def code_to_latent(self, token, mel=None):
+        quant, _ = self.quantizer.decode(token, None)
+        speaker_embedding = self.speaker_encoder(mel)
+        latents = quant + speaker_embedding.unsqueeze(1).repeat(1, quant.shape[1], 1)
+        return {
+            "latents": latents,
+        }

fireredtts/modules/text_normalizer/normalize.py

@@ -0,0 +1,183 @@
+import re
+import regex
+import inflect
+import unicodedata
+from lingua import Language, LanguageDetectorBuilder
+from builtins import str as unicode
+
+from tn.chinese.normalizer import Normalizer as ZhNormalizer
+from tn.english.normalizer import Normalizer as EnNormalizer
+
+from fireredtts.modules.text_normalizer.regex_common import *
+from fireredtts.modules.text_normalizer.utils import *
+
+
+def preprocess_text(sentence):
+    # preprocessing
+    sentence = bytes(sentence, "utf-8").decode("utf-8", "ignore")
+    sentence = regex.sub("[\p{Cf}--[\u200d]]", "", sentence, flags=regex.V1)
+    sentence = regex.sub("\p{Co}", "", sentence)
+    sentence = sentence.replace("\u00a0", " ")
+    sentence = sentence.replace("\ufffd", "")
+    sentence = regex.sub("\p{Zl}", "\n", sentence)
+    sentence = regex.sub("\p{Zp}", "\n", sentence)
+
+    sentence = unicode(sentence)
+    sentence = "".join(
+        char
+        for char in unicodedata.normalize("NFD", sentence)
+        if unicodedata.category(char) != "Mn"
+    )  # Strip accents
+
+    sentence = strip_kaomoji(sentence)
+    # full to half with exemption (to be converted after number TN): 。，：
+    sentence = f2b(sentence, exemption="。，：")
+
+    # clean spaces
+    sentence = sentence.replace("\n", "，")
+    sentence = sentence.replace("\t", "，")
+    sentence = sentence.replace("\r", "，")
+    sentence = re.sub(r"[。.]{3,}", "…", sentence)
+    sentence = re.sub(r"[…⋯]{1,}", "…", sentence)
+    sentence = re.sub(r"[ ]+", " ", sentence)
+    sentence = sentence.strip()
+
+    # punctuation reduction
+    result = ""
+    for idx, char in enumerate(sentence):
+        if char in symbol_reduction:
+            char = symbol_reduction[char]
+
+        if char == " ":
+            if idx == 0:
+                continue
+            if is_chinese(sentence[idx + 1]) and (
+                is_chinese(sentence[idx - 1]) or sentence[idx - 1] in '") '
+            ):
+                result += "，"
+            else:
+                result += " "
+            continue
+
+        if is_valid_char(char):
+            result += char
+    result = re.sub(r"[ ]+", " ", result)
+    return result
+
+
+def rettt(sentence):
+    # handle abbreviations for all languages
+    sentence = sentence.replace("&nd", "and")
+    sentence = sentence.replace("Jan.", "january")
+    sentence = sentence.replace("Feb.", "febrary")
+    sentence = sentence.replace("Mar.", "march")
+    sentence = sentence.replace("Apr.", "april")
+    sentence = sentence.replace("May.", "may")
+    sentence = sentence.replace("Jun.", "june")
+    sentence = sentence.replace("Jul.", "july")
+    sentence = sentence.replace("Aug.", "august")
+    sentence = sentence.replace("Sept.", "september")
+    sentence = sentence.replace("Sep.", "september")
+    sentence = sentence.replace("Oct.", "october")
+    sentence = sentence.replace("Nov.", "november")
+    sentence = sentence.replace("Dec.", "december")
+    sentence = sentence.replace("Mon.", "monday")
+    sentence = sentence.replace("Tues.", "tuesday")
+    sentence = sentence.replace("Wed.", "wednesday")
+    sentence = sentence.replace("Thur.", "thursday")
+    sentence = sentence.replace("Fri.", "friday")
+    sentence = sentence.replace("Sat.", "saturday")
+    if sentence != "Sun.":
+        sentence = sentence.replace("Sun.", "sunday")
+    sentence = re.sub(r" St\. ([A-Z])", r" saint \1", sentence)
+    sentence = re.sub(r" St\.", " street", sentence)
+    sentence = re.sub(r" Rd\.", " road", sentence)
+    sentence = re.sub(r"[Aa]\.[Mm]\.", "A_M", sentence)
+    sentence = re.sub(r"[Pp]\.[Mm]\.", "P_M", sentence)
+    sentence = re.sub(r"[Bb]\.[Cc]\.", "B_C", sentence)
+    sentence = re.sub(r"[Ad]\.[Dd]\.", "A_D", sentence)
+    sentence = sentence.replace("Mr.", "mister")
+    sentence = sentence.replace("Ms.", "miss")
+    sentence = sentence.replace("Mrs.", "misses")
+    sentence = sentence.replace("Ph.D", "P_H_D")
+    sentence = sentence.replace("i.e.", "that is")
+    sentence = sentence.replace("e.g.", "for example")
+    sentence = sentence.replace("btw.", "by the way")
+    sentence = sentence.replace("btw", "by the way")
+    sentence = sentence.replace("b.t.w.", "by the way")
+    sentence = sentence.replace("@", " at ")
+    return sentence
+
+
+class TextNormalizer:
+    def __init__(self):
+        self.language_detector = LanguageDetectorBuilder.from_languages(
+            Language.ENGLISH, Language.CHINESE
+        ).build()
+        self.zh_normalizer = ZhNormalizer()
+        self.en_normalizer = EnNormalizer()
+        self.inflect_parser = inflect.engine()
+        self.lang2token = {Language.ENGLISH: "en", Language.CHINESE: "zh"}
+
+    def tn(self, text):
+        text = preprocess_text(text)
+        text = rettt(text)  # regex replacements
+        # for non chinese languages
+        language = self.language_detector.detect_language_of(text)
+        # enforce chinese if text contains any chinese character
+        if contains_chinese(text):
+            language = Language.CHINESE
+        text_lang = self.lang2token.get(language, "zh")
+
+        if is_upper_eng_and_digit(text):
+            language = Language.CHINESE
+
+        if language == Language.CHINESE:
+            text = self.zh_normalizer.normalize(text)
+            # print("---text after zh_normalizer:", text)
+            text = text.replace("\n", "")
+            text = text.replace(",", "，")
+            text = text.replace(".", "。")
+            text = re.sub(r"[，,]+$", "。", text)
+            # print("---text after zh_normalizer 2:", text)
+        else:
+            text = re.sub(r"[^ 0-9A-Za-z\[\]'.,:?!_\-]", "", text)
+            text = self.en_normalizer.normalize(text)
+            # fallback number normalization
+            pieces = re.split(r"(\d+)", text)
+            text = "".join(
+                [
+                    self.inflect_parser.number_to_words(p) if p.isnumeric() else p
+                    for p in pieces
+                    if len(p) > 0
+                ]
+            )
+
+        # cleanup
+        text = text.replace("_", " ")
+        text = re.sub(r"[ ]+", " ", text)
+
+        # spell caplital words
+        pieces = re.split(r"([A-Z]{2,4}|[ ])", text)
+        for idx, p in enumerate(pieces):
+            if re.match("[A-Z]{2,4}", p):
+                pieces[idx] = " ".join(p)
+        text = " ".join([p for p in pieces if p != " "])
+
+        # post TN full to half
+        # text = text.replace("。", ".")
+        # text = text.replace("，", ",")
+        # text = text.replace("：", ":")
+
+        # model limitations
+        text = text.lower().strip()
+        text = text.replace('"', "")
+        text = text.replace("·", " ")
+        # text = re.sub("[…~！，&*%$#^：；!:;]+", ",", text)
+        text = re.sub("[…~！&*%$#^：；!:;]+", ",", text)
+        text = re.sub("[,]+", ",", text)
+        text = re.sub(r"[,. ]+$", ".", text)
+        if len(text) > 0 and text[-1] not in ".?":
+            text = text + "."
+        text = text.replace("。.", "。")
+        return text, text_lang

fireredtts/modules/text_normalizer/regex_common.py

@@ -0,0 +1,23 @@
+import re
+
+kaomoji_regex = re.compile(
+    r"[oヽwΣ┗╰O︿Ψ凸]?[(|≡*（].{0,4}[Д✿_▽→≧﹏`∩⊙∇☆≡๑〃′エ≦▔＠﹁εヘ•́ω益‿≖ฺ皿•̀艹￣△|ﾟ].{0,5}[|≡*)）][┛ブ凸cｄd︴oOΨ︿w╯ノ]?"
+)
+chinese_regex = re.compile(r"[\u4e00-\u9fa5]")
+digit_regex = re.compile(r"(\\d+)(\\.\\d+)?", re.UNICODE)
+
+chinese_char_regex = re.compile(r"^[\u4e00-\u9fa5]$", re.UNICODE)
+eng_and_digit_char_regex = re.compile(r"^[0-9.,A-Za-z]+$", re.UNICODE)
+upper_eng_and_digit_regex = re.compile(r"^[ 0-9A-Z\"'.,:?!\-]+$", re.UNICODE)
+valid_char_regex = re.compile(
+    r"[\t\r\n ]|"
+    r"[\u4e00-\u9fa5]|"
+    r"\u0080|[\u20a0-\u20bf]|\u00a2|\u00a3|\u00a5|\uffe0|\uffe1|\uffe5|\uffe6|"
+    r"\u3000|\u3002|\u00b7|\u2014|\u2019|\u2026|\uff01|\uff1f|\uff0e|\uff1a|\uff1b|\uff0b|\uff0c|\uff0d|\uff0f|[\ufe10-\ufe16]|[\ufe50-\ufe51]|[\ufe55-\ufe57]|\ufe6a|"
+    r"[\u0030-\u0040]|"
+    r"[\u0391-\u03c9]|"
+    r"[\u00b0-\u00b3]|[\u2015-\u2018]|[\u3000-\u303f]|"
+    r"[\u0022-\u002f\u003a-\u003e\u0040\u005b-\u0060\u007b-\u007e]|"
+    r"[\uff21-\uff3a]|[\uff41-\uff5a]|[\u0041-\u005a]|[\u0061-\u007a]",
+    re.UNICODE,
+)

fireredtts/modules/text_normalizer/utils.py

@@ -0,0 +1,171 @@
+from fireredtts.modules.text_normalizer.regex_common import *
+from sentencex import segment
+import re
+
+
+symbol_reduction = {
+    "「": '"',
+    "」": '"',
+    "｀": '"',
+    "〝": '"',
+    "〞": '"',
+    "‟": '"',
+    "„": '"',
+    "｛": "(",
+    "｝": ")",
+    "【": "(",
+    "】": ")",
+    "〖": "(",
+    "〗": ")",
+    "〔": "(",
+    "〕": ")",
+    "〘": "(",
+    "〙": ")",
+    "《": "(",
+    "》": ")",
+    "｟": "(",
+    "｠": ")",
+    "〚": "(",
+    "〛": ")",
+    "『": '"',
+    "』": '"',
+    "｢": '"',
+    "｣": '"',
+    "{": "(",
+    "}": ")",
+    "〈": "(",
+    "〉": ")",
+    "•": "·",
+    "‧": "·",
+    "〰": "…",
+    "﹏": "…",
+    "〜": "~",
+    "～": "~",
+    "＋": "+",
+    "､": "、",
+    "｡": "。",
+    "︐": "，",
+    "﹐": "，",
+    "︑": "、",
+    "﹑": "、",
+    "︒": "。",
+    "︓": "：",
+    "﹕": "：",
+    "︔": "；",
+    "﹔": "；",
+    "︕": "！",
+    "﹗": "！",
+    "︖": "？",
+    "﹖": "？",
+    "﹙": "(",
+    "﹚": ")",
+    "﹪": "%",
+    "﹠": "&",
+    "＞": ">",
+    "|": "、",
+    "＝": "=",
+    "‐": "-",
+    "‑": "-",
+    "‒": "-",
+    "–": "-",
+    "—": "-",
+    "―": "-",
+    "％": "%",
+    "μ": "u",
+}
+
+
+strong_break = re.compile("([。”;；!！：…?？）\)\]』】」}~\r\n]| \.)", re.UNICODE)
+weak_break = re.compile(
+    "["
+    "\U00002702-\U000027b0\U0001f926-\U0001f937\U00010000-\U0001fbff\U00030000-\U0010ffff"
+    "\u2640-\u2642\u2600-\u2b55\u23cf\u23e9\u231a\ufe0f\u3030"
+    "\t，,. ]",
+    re.UNICODE,
+)
+
+
+def contains_chinese(text):
+    return bool(chinese_regex.search(text))
+
+
+def strip_kaomoji(text):
+    return kaomoji_regex.sub(" ", text)
+
+
+def is_chinese(char):
+    return chinese_char_regex.match(char)
+
+
+def is_eng_and_digit(char):
+    return eng_and_digit_char_regex.match(char)
+
+
+def is_upper_eng_and_digit(text):
+    return upper_eng_and_digit_regex.match(text)
+
+
+def is_valid_char(char):
+    return valid_char_regex.match(char)
+
+
+def is_digit(text):
+    return digit_regex.match(text)
+
+
+def f2b(ustr, exemption="。，："):
+    half = []
+    for u in ustr:
+        num = ord(u)
+        if num == 0x3000:
+            half.append(" ")
+        elif u in exemption:  # exemption
+            half.append(u)
+        elif 0xFF01 <= num <= 0xFF5E:
+            num -= 0xFEE0
+            half.append(chr(num))
+        else:
+            half.append(u)
+    return "".join(half)
+
+
+def zh_text_split(text, length=80):
+    if length == 0:
+        return []
+    if length == 1:
+        return [c for c in length]
+    if len(text) <= length:
+        return [text]
+
+    match_strong = re.search(strong_break, text[:length][::-1])
+    match_weak = re.search(weak_break, text[:length][::-1])
+    end_ind_strong = length - match_strong.start() if match_strong else 0
+    end_ind_weak = length - match_weak.start() if match_weak else 0
+
+    if end_ind_strong < length // 3:
+        if end_ind_weak < length // 3:
+            valid_max = max(end_ind_strong, end_ind_weak)
+            if valid_max >= 3:
+                return [text[:valid_max]] + zh_text_split(text[valid_max:])
+            else:
+                return [text[:length]] + zh_text_split(text[length:])
+        else:
+            return [text[:end_ind_weak]] + zh_text_split(text[end_ind_weak:])
+    else:
+        return [text[:end_ind_strong]] + zh_text_split(text[end_ind_strong:])
+
+
+def text_split(text):
+    if contains_chinese(text):
+        substrings = list(segment("zh", text))
+        new_substrings = []
+        for s in substrings:
+            if len(s) > 50:
+                new_substrings += zh_text_split(s, length=50)
+            else:
+                new_substrings.append(s)
+        substrings = new_substrings
+    else:
+        substrings = list(segment("en", text))
+
+    return substrings

fireredtts/setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup, find_packages
+
+setup(name="fireredtts", version="0.1", packages=find_packages())

fireredtts/utils/spliter.py

@@ -0,0 +1,161 @@
+import re
+import string
+
+SYMBOLS_MAPPING = {
+    "\n": "",
+    "…": ".",
+    "“": "'",
+    "”": "'",
+    "‘": "'",
+    "’": "'",
+    "【": "",
+    "】": "",
+    "[": "",
+    "]": "",
+    "（": "",
+    "）": "",
+    "(": "",
+    ")": "",
+    "・": "",
+    "·": "",
+    "「": "'",
+    "」": "'",
+    "《": "'",
+    "》": "'",
+    "—": "",
+    "～": "",
+    "~": "",
+    "：": ",",
+    "；": ",",
+    ";": ",",
+    ":": ",",
+    '"': "",
+    "！": "。",
+    "!": ".",
+    "————": "，",
+    "——": "，",
+    "—": "，",
+    "……": "，",
+}
+
+REPLACE_SYMBOL_REGEX = re.compile(
+    "|".join(re.escape(p) for p in SYMBOLS_MAPPING.keys())
+)
+
+
+EMOJI_REGEX = re.compile(
+    "["
+    "\U0001f600-\U0001f64f"  # emoticons
+    "\U0001f300-\U0001f5ff"  # symbols & pictographs
+    "\U0001f680-\U0001f6ff"  # transport & map symbols
+    "\U0001f1e0-\U0001f1ff"  # flags (iOS)
+    "]+",
+    flags=re.UNICODE,
+)
+
+
+def clean_text(text):
+    # Clean the text
+    text = text.strip()
+    text = text.replace("\xa0", "")
+
+    # Replace all chinese symbols with their english counterparts
+    text = REPLACE_SYMBOL_REGEX.sub(lambda x: SYMBOLS_MAPPING[x.group()], text)
+
+    # Remove emojis
+    text = EMOJI_REGEX.sub(r"", text)
+
+    # Remove continuous periods (...) and commas (,,,)
+    text = re.sub(r"[.,]{2,}", lambda m: m.group()[0], text)
+
+    return text
+
+
+def utf_8_len(text):
+    return len(text.encode("utf-8"))
+
+
+def break_text(texts, length, splits: set):
+    for text in texts:
+        if utf_8_len(text) <= length:
+            yield text
+            continue
+
+        curr = ""
+        for char in text:
+            curr += char
+
+            if char in splits:
+                yield curr
+                curr = ""
+
+        if curr:
+            yield curr
+
+
+def break_text_by_length(texts, length):
+    for text in texts:
+        if utf_8_len(text) <= length:
+            yield text
+            continue
+
+        curr = ""
+        for char in text:
+            curr += char
+
+            if utf_8_len(curr) >= length:
+                yield curr
+                curr = ""
+
+        if curr:
+            yield curr
+
+
+def add_cleaned(curr, segments):
+    curr = curr.strip()
+    if curr and not all(c.isspace() or c in string.punctuation for c in curr):
+        segments.append(curr)
+
+
+def protect_float(text):
+    # Turns 3.14 into <3_f_14> to prevent splitting
+    return re.sub(r"(\d+)\.(\d+)", r"<\1_f_\2>", text)
+
+
+def unprotect_float(text):
+    # Turns <3_f_14> into 3.14
+    return re.sub(r"<(\d+)_f_(\d+)>", r"\1.\2", text)
+
+
+def split_text(text, length):
+    text = clean_text(text)
+
+    # Break the text into pieces with following rules:
+    # 1. Split the text at ".", "!", "?" if text is NOT a float
+    # 2. If the text is longer than length, split at ","
+    # 3. If the text is still longer than length, split at " "
+    # 4. If the text is still longer than length, split at any character to length
+
+    texts = [text]
+    texts = map(protect_float, texts)
+    texts = break_text(texts, length, {".", "!", "?", "。", "！", "？"})
+    texts = map(unprotect_float, texts)
+    texts = break_text(texts, length, {",", "，"})
+    texts = break_text(texts, length, {" "})
+    texts = list(break_text_by_length(texts, length))
+
+    # Then, merge the texts into segments with length <= length
+    segments = []
+    curr = ""
+
+    for text in texts:
+        if utf_8_len(curr) + utf_8_len(text) <= length:
+            curr += text
+        else:
+            add_cleaned(curr, segments)
+            curr = text
+
+    if curr:
+        add_cleaned(curr, segments)
+
+    return segments

fireredtts/utils/utils.py

@@ -0,0 +1,37 @@
+import os
+from pathlib import Path
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+
+
+def load_audio(audiopath, sampling_rate):
+    """_summary_
+
+    Args:
+        audiopath (_type_): audio_path
+        sampling_rate (_type_): sampling_rate
+
+    Returns:
+        _type_: _description_
+    """
+    audio, lsr = torchaudio.load(audiopath)
+
+    # stereo to mono if needed
+    if audio.size(0) != 1:
+        audio = torch.mean(audio, dim=0, keepdim=True)
+
+    # resample
+    audio_resampled = torchaudio.functional.resample(audio, lsr, sampling_rate)
+    if torch.any(audio > 10) or not torch.any(audio < 0):
+        print(f"Error with {audiopath}. Max={audio.max()} min={audio.min()}")
+
+    if torch.any(audio_resampled > 10) or not torch.any(audio_resampled < 0):
+        print(
+            f"Error with {audiopath}. Max={audio_resampled.max()} min={audio_resampled.min()}"
+        )
+    # clip audio invalid values
+    audio.clip_(-1, 1)
+    audio_resampled.clip_(-1, 1)
+    return audio, lsr, audio_resampled

pre-requirements.txt

@@ -0,0 +1 @@
+pip==24.0

pretrained_models/README.md

@@ -0,0 +1,3 @@
+## Pretrained Models
+
+Download the required model files and place them in the folder `pretrained_models`

requirements.txt

@@ -0,0 +1,12 @@
+torchaudio
+fairseq
+diffusers==0.27.2
+librosa==0.10.2
+soundfile==0.12.1
+einops==0.8.0
+transformers==4.44.2
+tiktoken==0.7.0
+inflect==7.4.0
+lingua-language-detector==2.0.2
+WeTextProcessing==1.0.3
+sentencex==0.6.1