initial deploy

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.pyc
+.DS_Store
+.venv/
+ckpts/
+Vevo2-repo/

README.md

@@ -1,13 +1,19 @@
 ---
 title: Vevo2
-emoji: 💻
-colorFrom: yellow
-colorTo: yellow
+emoji: 🎤
+colorFrom: blue
+colorTo: purple
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 5.28.0
+python_version: '3.11'
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Vevo2
+
+A unified, controllable framework for speech and singing voice generation and conversion. Supports zero-shot voice/singing conversion, text-to-singing, singing editing, singing style conversion, and melody control (hum or play a melody to control a generated vocal).
+
+Paper: [Vevo2: Bridging Controllable Speech and Singing Voice Generation via Unified Prosody Learning](https://arxiv.org/abs/2508.16332)
+Source: [open-mmlab/Amphion](https://github.com/open-mmlab/Amphion/blob/main/models/svc/vevo2/README.md)

SOURCES.md

@@ -0,0 +1,4 @@
+# Sources — Vevo2
+
+- Source repo: https://github.com/open-mmlab/Amphion
+- Paper: https://arxiv.org/abs/2508.16332

app.py

@@ -0,0 +1,312 @@
+try:
+    import spaces
+except ImportError:
+    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
+    class spaces:
+        class GPU:
+            def __init__(self, func=None, duration=60):
+                self.func = func
+
+            def __call__(self, *args, **kwargs):
+                if self.func is not None:
+                    return self.func(*args, **kwargs)
+                func = args[0]
+                return func
+
+
+import sys
+
+sys.stdout.reconfigure(line_buffering=True)
+
+import os
+import re
+import tempfile
+import threading
+
+import torch
+import gradio as gr
+import whisper
+from huggingface_hub import snapshot_download
+from pyharp import ModelCard, build_endpoint
+
+from models.svc.vevo2.vevo2_utils import Vevo2InferencePipeline, save_audio
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+CKPT_DIR = "./ckpts/Vevo2"
+# training-only artifacts (optimizer/scheduler/rng state, and the earlier
+# "pretrained" AR checkpoint superseded by "posttrained") aren't needed for
+# inference and would roughly double the download.
+CKPT_IGNORE_PATTERNS = [
+    "*/optimizer.pt",
+    "*/optimizer.bin",
+    "*/scheduler.pt",
+    "*/scheduler.bin",
+    "*/rng_state*.pth",
+    "*/random_states_*.pkl",
+    "*/trainer_state.json",
+    "*/training_args.bin",
+    "contentstyle_modeling/pretrained/*",
+]
+
+pipeline = None
+download_ready = False
+download_error = None
+
+
+FMT_CONFIG_PATH = os.path.join(
+    CKPT_DIR, "acoustic_modeling/fm_emilia101k_singnet7k_repa/config.json"
+)
+
+
+def _fix_whisper_stats_path():
+    """The checkpoint's config.json hardcodes whisper_stats_path as
+    models/svc/vevosing/config/whisper_stats.pt, relative to the full
+    Amphion repo — a directory this deployment doesn't copy. The identical
+    file ships right next to this config, so point it there instead."""
+    correct_path = os.path.join(
+        CKPT_DIR, "acoustic_modeling/fm_emilia101k_singnet7k_repa/whisper_stats.pt"
+    )
+    with open(FMT_CONFIG_PATH) as f:
+        text = f.read()
+    text = re.sub(
+        r'"whisper_stats_path":\s*"[^"]*"',
+        f'"whisper_stats_path": "{correct_path}"',
+        text,
+    )
+    with open(FMT_CONFIG_PATH, "w") as f:
+        f.write(text)
+
+
+def load_checkpoint():
+    """Download Vevo2's weights and pre-warm the Whisper cache. CPU-only —
+    building the actual pipeline happens lazily on first process_fn call
+    (see get_pipeline) instead of here, since Vevo2InferencePipeline moves
+    every submodule onto `device` inside its own constructor and ZeroGPU
+    only intercepts CUDA calls made inside an @spaces.GPU-decorated call,
+    not from a background thread."""
+    global download_ready, download_error
+    try:
+        snapshot_download(
+            repo_id="RMSnow/Vevo2",
+            local_dir=CKPT_DIR,
+            ignore_patterns=CKPT_IGNORE_PATTERNS,
+        )
+        _fix_whisper_stats_path()
+        whisper.load_model("medium", device="cpu")  # warms ~/.cache/whisper
+        print("Checkpoint download complete.")
+    except Exception as e:
+        download_error = str(e)
+        print(f"Download error: {e}")
+    finally:
+        download_ready = True
+
+
+threading.Thread(target=load_checkpoint, daemon=True).start()
+
+
+def get_pipeline():
+    """Build the inference pipeline on first use, inside the @spaces.GPU
+    call (see load_checkpoint for why this can't happen in the background
+    thread)."""
+    global pipeline
+    if pipeline is None:
+        pipeline = Vevo2InferencePipeline(
+            prosody_tokenizer_ckpt_path=os.path.join(
+                CKPT_DIR, "tokenizer/prosody_fvq512_6.25hz"
+            ),
+            content_style_tokenizer_ckpt_path=os.path.join(
+                CKPT_DIR, "tokenizer/contentstyle_fvq16384_12.5hz"
+            ),
+            ar_cfg_path=os.path.join(
+                CKPT_DIR, "contentstyle_modeling/posttrained/amphion_config.json"
+            ),
+            ar_ckpt_path=os.path.join(CKPT_DIR, "contentstyle_modeling/posttrained"),
+            fmt_cfg_path=os.path.join(
+                CKPT_DIR, "acoustic_modeling/fm_emilia101k_singnet7k_repa/config.json"
+            ),
+            fmt_ckpt_path=os.path.join(
+                CKPT_DIR, "acoustic_modeling/fm_emilia101k_singnet7k_repa"
+            ),
+            vocoder_cfg_path=os.path.join(CKPT_DIR, "vocoder/config.json"),
+            vocoder_ckpt_path=os.path.join(CKPT_DIR, "vocoder"),
+            device=DEVICE,
+        )
+    return pipeline
+
+
+model_card = ModelCard(
+    name="Vevo2",
+    description=(
+        "Zero-shot speech and singing voice generation and conversion: "
+        "voice/singing conversion, text-to-singing, singing editing, "
+        "singing style conversion, and melody control."
+    ),
+    author=(
+        "Xueyao Zhang, Junan Zhang, Yuancheng Wang, Chaoren Wang, "
+        "Yuanzhe Chen, Dongya Jia, Zhuo Chen, Zhizheng Wu"
+    ),
+    tags=["voice-conversion", "singing-synthesis", "text-to-speech"],
+)
+
+TASKS = [
+    "Voice/Singing Conversion",
+    "Text-to-Speech / Text-to-Singing",
+    "Singing Editing",
+    "Singing Style Conversion",
+    "Melody Control (Humming/Instrument to Singing)",
+]
+
+
+@spaces.GPU(duration=120)
+@torch.inference_mode()
+def process_fn(
+    task: str,
+    input_audio_path: str,
+    reference_audio_path: str,
+    target_text: str,
+    reference_text: str,
+    flow_matching_steps: int,
+) -> str:
+    """Runs the Vevo2 task selected in the Task dropdown. Each task calls a
+    different combination of Vevo2InferencePipeline.inference_fm /
+    inference_ar_and_fm, mirroring the task wrapper functions in the
+    original repo's infer_vevo2_fm.py / infer_vevo2_ar.py."""
+    if not download_ready:
+        raise gr.Error("Model is still downloading, please wait a moment and try again.")
+    if download_error is not None:
+        raise gr.Error(f"Model failed to download: {download_error}")
+
+    pipe = get_pipeline()
+    target_text = (target_text or "").strip() or None
+    reference_text = reference_text or ""
+    reference_audio_path = reference_audio_path or None
+
+    if task == "Voice/Singing Conversion":
+        if reference_audio_path is None:
+            raise gr.Error(
+                "Voice/Singing Conversion needs a reference voice (Reference Voice input)."
+            )
+        gen_audio = pipe.inference_fm(
+            src_wav_path=input_audio_path,
+            timbre_ref_wav_path=reference_audio_path,
+            use_pitch_shift=True,
+            flow_matching_steps=flow_matching_steps,
+        )
+    elif task == "Text-to-Speech / Text-to-Singing":
+        if target_text is None:
+            raise gr.Error("Text-to-Speech needs Text / Lyrics.")
+        gen_audio = pipe.inference_ar_and_fm(
+            target_text=target_text,
+            style_ref_wav_path=input_audio_path,
+            style_ref_wav_text=reference_text,
+            timbre_ref_wav_path=reference_audio_path or input_audio_path,
+            use_prosody_code=False,
+            flow_matching_steps=flow_matching_steps,
+        )
+    elif task == "Singing Editing":
+        if target_text is None:
+            raise gr.Error("Singing Editing needs the edited Text / Lyrics.")
+        gen_audio = pipe.inference_ar_and_fm(
+            target_text=target_text,
+            prosody_wav_path=input_audio_path,
+            style_ref_wav_path=input_audio_path,
+            style_ref_wav_text=reference_text,
+            timbre_ref_wav_path=input_audio_path,
+            use_prosody_code=True,
+            flow_matching_steps=flow_matching_steps,
+        )
+    elif task == "Singing Style Conversion":
+        if reference_audio_path is None:
+            raise gr.Error(
+                "Singing Style Conversion needs a style reference (Reference Voice input)."
+            )
+        gen_audio = pipe.inference_ar_and_fm(
+            target_text=target_text,
+            prosody_wav_path=input_audio_path,
+            style_ref_wav_path=reference_audio_path,
+            style_ref_wav_text=reference_text,
+            timbre_ref_wav_path=input_audio_path,
+            use_prosody_code=True,
+            use_pitch_shift=True,
+            flow_matching_steps=flow_matching_steps,
+        )
+    else:  # Melody Control
+        if target_text is None:
+            raise gr.Error("Melody Control needs the Text / Lyrics to sing.")
+        if reference_audio_path is None:
+            raise gr.Error(
+                "Melody Control needs a reference voice (Reference Voice input) — "
+                "the melody input alone (e.g. humming) has no usable vocal timbre."
+            )
+        gen_audio = pipe.inference_ar_and_fm(
+            target_text=target_text,
+            prosody_wav_path=input_audio_path,
+            style_ref_wav_path=reference_audio_path,
+            style_ref_wav_text=reference_text,
+            timbre_ref_wav_path=reference_audio_path,
+            use_prosody_code=True,
+            use_pitch_shift=True,
+            flow_matching_steps=flow_matching_steps,
+        )
+
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+        output_path = f.name
+    save_audio(gen_audio, output_path=output_path)
+    return output_path
+
+
+with gr.Blocks() as demo:
+    input_components = [
+        gr.Dropdown(
+            choices=TASKS,
+            value=TASKS[0],
+            label="Task",
+            info="What to do with the input audio.",
+        ),
+        gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info(
+            "Content/source for Voice/Singing Conversion and Singing Style Conversion; "
+            "style reference for Text-to-Speech; original recording for Singing Editing; "
+            "melody reference (e.g. humming or an instrument) for Melody Control."
+        ),
+        gr.Audio(type="filepath", label="Reference Voice (Timbre)")
+        .harp_required(False)
+        .set_info(
+            "A clip of the target voice. Required for Voice/Singing Conversion, "
+            "Singing Style Conversion, and Melody Control; optional for "
+            "Text-to-Speech (defaults to Input Audio's voice); ignored for Singing Editing."
+        ),
+        gr.Textbox(
+            label="Text / Lyrics",
+            info="Required for Text-to-Speech, Singing Editing, and Melody Control. "
+            "Optional for Singing Style Conversion. Ignored for Voice/Singing Conversion.",
+        ),
+        gr.Textbox(
+            label="Input Audio Transcript (optional)",
+            info="Transcript of Input Audio, if known — improves quality. "
+            "Not used for Voice/Singing Conversion.",
+        ),
+        gr.Slider(
+            minimum=8,
+            maximum=64,
+            step=1,
+            value=32,
+            label="Generation Detail (Steps)",
+            info="More steps trade generation speed for audio detail (default: 32, per repo).",
+        ),
+    ]
+    output_components = [
+        gr.Audio(type="filepath", label="Output Audio").set_info(
+            "Generated speech or singing voice."
+        ),
+    ]
+
+    build_endpoint(
+        model_card=model_card,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn,
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch(pwa=True)

evaluation/metrics/f0/f0_corr.py

@@ -0,0 +1,224 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import librosa
+import numpy as np
+import parselmouth
+
+from utils.f0 import interpolate
+
+
+def get_cents(f0_hz):
+    """
+    F_{cent} = 1200 * log2 (F/440)
+
+    Reference:
+        APSIPA'17, Perceptual Evaluation of Singing Quality
+    """
+    voiced_f0 = f0_hz[f0_hz != 0]
+    return 1200 * np.log2(voiced_f0 / 440)
+
+
+def get_pitch_sub_median(f0_hz):
+    """
+    f0_hz: (,T)
+    """
+    f0_cent = get_cents(f0_hz)
+    return f0_cent - np.median(f0_cent)
+
+
+def get_f0_features_using_parselmouth(audio, cfg, speed=1):
+    """Using parselmouth to extract the f0 feature.
+    Args:
+        audio
+        mel_len
+        hop_length
+        fs
+        f0_min
+        f0_max
+        speed(default=1)
+    Returns:
+        f0: numpy array of shape (frame_len,)
+        pitch_coarse: numpy array of shape (frame_len,)
+    """
+    hop_size = int(np.round(cfg.hop_size * speed))
+
+    # Calculate the time step for pitch extraction
+    time_step = hop_size / cfg.sample_rate * 1000
+
+    f0 = (
+        parselmouth.Sound(audio, cfg.sample_rate)
+        .to_pitch_ac(
+            time_step=time_step / 1000,
+            voicing_threshold=0.6,
+            pitch_floor=cfg.f0_min,
+            pitch_ceiling=cfg.f0_max,
+        )
+        .selected_array["frequency"]
+    )
+    return f0
+
+
+class JsonHParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if type(v) == dict:
+                v = JsonHParams(**v)
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+
+def extract_f0_hz(
+    wav_path,
+    fs=16000,
+    hop_length=256,
+    f0_min=50,
+    f0_max=1100,
+):
+    cfg = JsonHParams()
+    cfg.sample_rate = fs
+    cfg.hop_size = hop_length
+    cfg.f0_min = f0_min
+    cfg.f0_max = f0_max
+    cfg.pitch_bin = 256
+    cfg.pitch_max = f0_max
+    cfg.pitch_min = f0_min
+
+    # Compute f0
+    audio, _ = librosa.load(wav_path, sr=fs)
+    f0 = get_f0_features_using_parselmouth(
+        audio,
+        cfg,
+    )
+    f0, _ = interpolate(f0)
+    return f0
+
+
+def extract_fpc(
+    audio_ref,
+    audio_deg,
+    fs=16000,
+    need_mean=True,
+    hop_length=256,
+    f0_min=50,
+    f0_max=1100,
+    method="dtw",
+):
+    """Compute F0 Pearson Distance (FPC) between the predicted and the ground truth audio.
+    audio_ref: path to the ground truth audio.
+    audio_deg: path to the predicted audio.
+    fs: sampling rate.
+    hop_length: hop length.
+    f0_min: lower limit for f0.
+    f0_max: upper limit for f0.
+    pitch_bin: number of bins for f0 quantization.
+    pitch_max: upper limit for f0 quantization.
+    pitch_min: lower limit for f0 quantization.
+    need_mean: subtract the mean value from f0 if "True".
+    method: "dtw" will use dtw algorithm to align the length of the ground truth and predicted audio.
+            "cut" will cut both audios into a same length according to the one with the shorter length.
+    """
+    # Initialize method
+    from torchmetrics import PearsonCorrCoef
+
+    pearson = PearsonCorrCoef()
+
+    # Load audio
+    if fs != None:
+        audio_ref, _ = librosa.load(audio_ref, sr=fs)
+        audio_deg, _ = librosa.load(audio_deg, sr=fs)
+    else:
+        audio_ref, ref_fs = librosa.load(audio_ref)
+        audio_deg, deg_fs = librosa.load(audio_deg)
+        assert ref_fs == deg_fs
+        fs = ref_fs
+
+    # Initialize config
+    cfg = JsonHParams()
+    cfg.sample_rate = fs
+    cfg.hop_size = hop_length
+    cfg.f0_min = f0_min
+    cfg.f0_max = f0_max
+    cfg.pitch_bin = 256
+    cfg.pitch_max = f0_max
+    cfg.pitch_min = f0_min
+
+    # Compute f0
+    f0_ref = get_f0_features_using_parselmouth(
+        audio_ref,
+        cfg,
+    )
+
+    f0_deg = get_f0_features_using_parselmouth(
+        audio_deg,
+        cfg,
+    )
+
+    # Subtract mean value from f0
+    if need_mean:
+        f0_ref = torch.from_numpy(f0_ref)
+        f0_deg = torch.from_numpy(f0_deg)
+
+        f0_ref = get_pitch_sub_median(f0_ref).numpy()
+        f0_deg = get_pitch_sub_median(f0_deg).numpy()
+
+    # Avoid silence
+    min_length = min(len(f0_ref), len(f0_deg))
+    if min_length <= 1:
+        return 1
+
+    # F0 length alignment
+    if method == "cut":
+        length = min(len(f0_ref), len(f0_deg))
+        f0_ref = f0_ref[:length]
+        f0_deg = f0_deg[:length]
+    elif method == "dtw":
+        _, wp = librosa.sequence.dtw(f0_ref, f0_deg, backtrack=True)
+        f0_gt_new = []
+        f0_pred_new = []
+        for i in range(wp.shape[0]):
+            gt_index = wp[i][0]
+            pred_index = wp[i][1]
+            f0_gt_new.append(f0_ref[gt_index])
+            f0_pred_new.append(f0_deg[pred_index])
+        f0_ref = np.array(f0_gt_new)
+        f0_deg = np.array(f0_pred_new)
+        assert len(f0_ref) == len(f0_deg)
+
+    # Convert to tensor
+    f0_ref = torch.from_numpy(f0_ref)
+    f0_deg = torch.from_numpy(f0_deg)
+
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        f0_ref = f0_ref.to(device)
+        f0_deg = f0_deg.to(device)
+        pearson = pearson.to(device)
+
+    return pearson(f0_ref, f0_deg).detach().cpu().numpy().tolist()

model.json

@@ -0,0 +1,10 @@
+{
+  "name": "Vevo2",
+  "package_dir": ".",
+  "entry_point": "models.svc.vevo2.vevo2_utils.Vevo2InferencePipeline",
+  "checkpoint": {
+    "repo": "RMSnow/Vevo2",
+    "download": "snapshot_download at startup, allow_patterns excluding optimizer/scheduler/rng_state training artifacts",
+    "size_mb": 4400
+  }
+}

models/codec/amphion_codec/quantize/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from models.codec.amphion_codec.quantize.factorized_vector_quantize import (
+    FactorizedVectorQuantize,
+)
+from models.codec.amphion_codec.quantize.vector_quantize import VectorQuantize
+from models.codec.amphion_codec.quantize.lookup_free_quantize import LookupFreeQuantize
+from models.codec.amphion_codec.quantize.residual_vq import ResidualVQ

models/codec/amphion_codec/quantize/factorized_vector_quantize.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class FactorizedVectorQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+        commitment=0.005,
+        codebook_loss_weight=1.0,
+        use_l2_normlize=True,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.commitment = commitment
+        self.codebook_loss_weight = codebook_loss_weight
+        self.use_l2_normlize = use_l2_normlize
+
+        if self.input_dim != self.codebook_dim:
+            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
+            self.out_project = WNConv1d(
+                self.codebook_dim, self.input_dim, kernel_size=1
+            )
+
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+
+        self.codebook = nn.Embedding(self.codebook_size, self.codebook_dim)
+
+    def forward(self, z):
+        """
+        Parameters
+        ----------
+        z: torch.Tensor[B x D x T]
+
+        Returns
+        -------
+        z_q: torch.Tensor[B x D x T]
+            Quantized continuous representation of input
+        commit_loss: Tensor[B]
+            Commitment loss to train encoder to predict vectors closer to codebook entries
+        codebook_loss: Tensor[B]
+            Codebook loss to update the codebook
+        indices: torch.Tensor[B x T]
+            Codebook indices (quantized discrete representation of input)
+        z_e: torch.Tensor[B x D x T]
+            Projected latents (continuous representation of input before quantization)
+        """
+
+        # Factorized codes project input into low-dimensional space if self.input_dim != self.codebook_dim
+        z_e = self.in_project(z)
+        z_q, indices = self.decode_latents(z_e)
+
+        # Compute commitment loss and codebook loss
+        if self.training:
+            commit_loss = (
+                F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+                * self.commitment
+            )
+            codebook_loss = (
+                F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+                * self.codebook_loss_weight
+            )
+        else:
+            commit_loss = torch.zeros(z.shape[0], device=z.device)
+            codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        z_q = z_e + (z_q - z_e).detach()
+
+        z_q = self.out_project(z_q)
+
+        return z_q, commit_loss, codebook_loss, indices, z_e
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight
+
+        # L2 normalize encodings and codebook
+        if self.use_l2_normlize:
+            encodings = F.normalize(encodings)
+            codebook = F.normalize(codebook)
+
+        # Compute euclidean distance between encodings and codebook,
+        # if use_l2_normlize is True, the distance is equal to cosine distance
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+
+        return z_q, indices
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = self.decode_code(vq)
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb
+
+    def latent2dist(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight
+
+        # L2 normalize encodings and codebook
+        if self.use_l2_normlize:
+            encodings = F.normalize(encodings)
+            codebook = F.normalize(codebook)
+
+        # Compute euclidean distance between encodings and codebook,
+        # if use_l2_normlize is True, the distance is equal to cosine distance
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )  # (b*t, k)
+
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        dist = rearrange(dist, "(b t) k -> b t k", b=latents.size(0))
+        z_q = self.decode_code(indices)
+
+        return -dist, indices, z_q

models/codec/amphion_codec/quantize/lookup_free_quantize.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class LookupFreeQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        assert 2**codebook_dim == codebook_size
+
+        if self.input_dim != self.codebook_dim:
+            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
+            self.out_project = WNConv1d(
+                self.codebook_dim, self.input_dim, kernel_size=1
+            )
+
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+
+    def forward(self, z):
+        z_e = self.in_project(z)
+        z_e = F.sigmoid(z_e)
+
+        z_q = z_e + (torch.round(z_e) - z_e).detach()
+
+        z_q = self.out_project(z_q)
+
+        commit_loss = torch.zeros(z.shape[0], device=z.device)
+        codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        bits = (
+            2
+            ** torch.arange(self.codebook_dim, device=z.device)
+            .unsqueeze(0)
+            .unsqueeze(-1)
+            .long()
+        )  # (1, d, 1)
+        indices = (torch.round(z_e.clone().detach()).long() * bits).sum(1).long()
+
+        return z_q, commit_loss, codebook_loss, indices, z_e
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = torch.zeros(
+            vq.shape[0], self.codebook_dim, vq.shape[-1], device=vq.device
+        )  # (B, d, T)
+        for i in range(self.codebook_dim):
+            emb[:, i, :] = (vq % 2).float()
+            vq = vq // 2
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb

models/codec/amphion_codec/quantize/residual_vq.py

@@ -0,0 +1,177 @@
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.utils import weight_norm
+
+from models.codec.amphion_codec.quantize.factorized_vector_quantize import (
+    FactorizedVectorQuantize,
+)
+from models.codec.amphion_codec.quantize.vector_quantize import VectorQuantize
+from models.codec.amphion_codec.quantize.lookup_free_quantize import LookupFreeQuantize
+
+
+class ResidualVQ(nn.Module):
+    """
+    Introduced in SoundStream: An end2end neural audio codec
+    https://arxiv.org/abs/2107.03312
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 256,
+        num_quantizers: int = 8,
+        codebook_size: int = 1024,
+        codebook_dim: int = 256,
+        quantizer_type: str = "vq",  # "vq" or "fvq" or "lfq"
+        quantizer_dropout: float = 0.5,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.num_quantizers = num_quantizers
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.quantizer_type = quantizer_type
+        self.quantizer_dropout = quantizer_dropout
+
+        if quantizer_type == "vq":
+            VQ = VectorQuantize
+        elif quantizer_type == "fvq":
+            VQ = FactorizedVectorQuantize
+        elif quantizer_type == "lfq":
+            VQ = LookupFreeQuantize
+        else:
+            raise ValueError(f"Unknown quantizer type {quantizer_type}")
+
+        self.quantizers = nn.ModuleList(
+            [
+                VQ(
+                    input_dim=input_dim,
+                    codebook_size=codebook_size,
+                    codebook_dim=codebook_dim,
+                    **kwargs,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+
+    def forward(self, z, n_quantizers: int = None):
+        """
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+        n_quantizers : int, optional
+            No. of quantizers to use
+            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
+            Note: if `self.quantizer_dropout` is True, this argument is ignored
+                when in training mode, and a random number of quantizers is used.
+        Returns
+        -------
+        "quantized_out" : Tensor[B x D x T]
+            Quantized continuous representation of input
+        "all_indices" : Tensor[N x B x T]
+            Codebook indices for each codebook
+            (quantized discrete representation of input)
+        "all_commit_losses" : Tensor[N]
+        "all_codebook_losses" : Tensor[N]
+        "all_quantized" : Tensor[N x B x D x T]
+        """
+
+        quantized_out = 0.0
+        residual = z
+
+        all_commit_losses = []
+        all_codebook_losses = []
+        all_indices = []
+        all_quantized = []
+
+        if n_quantizers is None:
+            n_quantizers = self.num_quantizers
+
+        if self.training:
+            n_quantizers = torch.ones((z.shape[0],)) * self.num_quantizers + 1
+            dropout = torch.randint(1, self.num_quantizers + 1, (z.shape[0],))
+            n_dropout = int(z.shape[0] * self.quantizer_dropout)
+            n_quantizers[:n_dropout] = dropout[:n_dropout]
+            n_quantizers = n_quantizers.to(z.device)
+
+        for i, quantizer in enumerate(self.quantizers):
+            if self.training is False and i >= n_quantizers:
+                break
+
+            z_q_i, commit_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+                residual
+            )
+
+            # Create mask to apply quantizer dropout
+            mask = (
+                torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
+            )
+            quantized_out = quantized_out + z_q_i * mask[:, None, None]
+            residual = residual - z_q_i
+
+            commit_loss_i = (commit_loss_i * mask).mean()
+            codebook_loss_i = (codebook_loss_i * mask).mean()
+
+            all_commit_losses.append(commit_loss_i)
+            all_codebook_losses.append(codebook_loss_i)
+            all_indices.append(indices_i)
+            all_quantized.append(z_q_i)
+
+        all_commit_losses, all_codebook_losses, all_indices, all_quantized = map(
+            torch.stack,
+            (all_commit_losses, all_codebook_losses, all_indices, all_quantized),
+        )
+
+        return (
+            quantized_out,
+            all_indices,
+            all_commit_losses,
+            all_codebook_losses,
+            all_quantized,
+        )
+
+    def vq2emb(self, vq, n_quantizers=None):
+        quantized_out = 0.0
+        if n_quantizers is None:
+            n_quantizers = self.num_quantizers
+        for idx, quantizer in enumerate(self.quantizers):
+            if idx >= n_quantizers:
+                break
+            quantized_out += quantizer.vq2emb(vq[idx])
+        return quantized_out
+
+    def latent2dist(self, z, n_quantizers=None):
+        quantized_out = 0.0
+        residual = z
+
+        all_dists = []
+        all_indices = []
+
+        if n_quantizers is None:
+            n_quantizers = self.num_quantizers
+
+        for i, quantizer in enumerate(self.quantizers):
+            if self.training is False and i >= n_quantizers:
+                break
+            dist_i, indices_i, z_q_i = quantizer.latent2dist(residual)
+            all_dists.append(dist_i)
+            all_indices.append(indices_i)
+
+            quantized_out = quantized_out + z_q_i
+            residual = residual - z_q_i
+
+        all_dists = torch.stack(all_dists)
+        all_indices = torch.stack(all_indices)
+
+        return all_dists, all_indices

models/codec/amphion_codec/quantize/vector_quantize.py

@@ -0,0 +1,401 @@
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch.nn.utils import weight_norm
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+def ema_inplace(moving_avg, new, decay):
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def laplace_smoothing(x, n_categories, eps=1e-5):
+    return (x + eps) / (x.sum() + n_categories * eps)
+
+
+def sample_vectors(samples, num):
+    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]
+
+
+def kmeans(samples, num_clusters, num_iters=10, use_cosine_sim=False):
+    dim, dtype, device = samples.shape[-1], samples.dtype, samples.device
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        if use_cosine_sim:
+            dists = samples @ means.t()
+        else:
+            diffs = rearrange(samples, "n d -> n () d") - rearrange(
+                means, "c d -> () c d"
+            )
+            dists = -(diffs**2).sum(dim=-1)
+
+        buckets = dists.max(dim=-1).indices
+        bins = torch.bincount(buckets, minlength=num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
+        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        if use_cosine_sim:
+            new_means = l2norm(new_means)
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+class EuclideanCodebook(nn.Module):
+    def __init__(
+        self,
+        dim,
+        codebook_size,
+        kmeans_init=False,
+        kmeans_iters=10,
+        decay=0.8,
+        eps=1e-5,
+        threshold_ema_dead_code=2,
+        weight_init=False,
+    ):
+        super().__init__()
+
+        self.decay = decay
+        init_fn = torch.randn if not weight_init else torch.zeros
+        embed = init_fn(codebook_size, dim)
+
+        if weight_init:
+            nn.init.uniform_(embed, -1 / codebook_size, 1 / codebook_size)
+
+        self.codebook_size = codebook_size
+        self.kmeans_iters = kmeans_iters
+        self.eps = eps
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+
+        self.register_buffer(
+            "initted", torch.Tensor([not kmeans_init])
+        )  # if kmeans_init is True, then initted is False; otherwise, initted is True
+        self.register_buffer("cluster_size", torch.zeros(codebook_size))
+        self.register_buffer("embed", embed)
+        self.register_buffer("embed_avg", embed.clone())
+
+    def init_embed_(self, data):
+        embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
+        self.embed.data.copy_(embed)
+        self.embed_avg.data.copy_(embed)
+        self.cluster_size.data.copy_(cluster_size)
+        self.initted.data.copy_(torch.Tensor([True]))
+
+    def replace(self, samples, mask):
+        modified_codebook = torch.where(
+            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
+        )
+        self.embed.data.copy_(modified_codebook)
+
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            return
+
+        expired_codes = self.cluster_size < self.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)
+
+    def forward(self, x):
+        shape, dtype = x.shape, x.dtype
+        flatten = rearrange(x, "... d -> (...) d")
+        embed = self.embed.t()  # (codebook_size, dim) -> (dim, codebook_size)
+
+        if not self.initted:
+            self.init_embed_(flatten)
+
+        dist = -(
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+
+        embed_ind = dist.max(dim=-1).indices
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
+        embed_ind = embed_ind.view(*shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed)
+
+        if self.training:
+            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
+            embed_sum = (
+                flatten.t() @ embed_onehot
+            )  # (dim, ...) @ (..., codebook_size) -> (dim, codebook_size)
+            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
+            cluster_size = (
+                laplace_smoothing(self.cluster_size, self.codebook_size, self.eps)
+                * self.cluster_size.sum()
+            )
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embed.data.copy_(embed_normalized)
+            self.expire_codes_(x)
+
+        return quantize, embed_ind
+
+    def vq2emb(self, vq):
+        quantize = F.embedding(vq, self.embed)
+        return quantize
+
+    def latent2dist(self, x):
+        shape, dtype = x.shape, x.dtype
+        flatten = rearrange(x, "... d -> (...) d")
+        embed = self.embed.t()  # (codebook_size, dim) -> (dim, codebook_size)
+
+        if not self.initted:
+            self.init_embed_(flatten)
+
+        dist = -(
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+
+        embed_ind = dist.max(dim=-1).indices
+        embed_ind = embed_ind.view(*shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed)
+
+        dist = dist.view(*shape[:-1], -1)
+
+        return dist, embed_ind, quantize
+
+
+class SimpleCodebook(nn.Module):
+    def __init__(
+        self,
+        dim,
+        codebook_size,
+        use_l2_normlize=False,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.codebook_size = codebook_size
+        self.use_l2_normlize = use_l2_normlize
+
+        self.embed = nn.Embedding(self.codebook_size, self.dim)
+
+    def forward(self, x):
+        shape, dtype = x.shape, x.dtype
+        flatten = rearrange(x, "... d -> (...) d")
+        embed = self.embed.weight.t()  # (codebook_size, dim) -> (dim, codebook_size)
+
+        if self.use_l2_normlize:
+            flatten = F.normalize(flatten)
+            embed = F.normalize(embed)
+
+        dist = -(
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+
+        embed_ind = dist.max(dim=-1).indices
+        embed_ind = embed_ind.view(*shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed)
+
+        return quantize, embed_ind
+
+    def vq2emb(self, vq):
+        quantize = F.embedding(vq, self.embed.weight)
+        return quantize
+
+    def latent2dist(self, x):
+        shape, dtype = x.shape, x.dtype
+        flatten = rearrange(x, "... d -> (...) d")
+        embed = self.embed.weight.t()  # (codebook_size, dim) -> (dim, codebook_size)
+
+        if self.use_l2_normlize:
+            flatten = F.normalize(flatten)
+            embed = F.normalize(embed)
+
+        dist = -(
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+
+        embed_ind = dist.max(dim=-1).indices
+        embed_ind = embed_ind.view(*shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed)
+
+        dist = dist.view(*shape[:-1], -1)
+
+        return dist, embed_ind, quantize
+
+
+class VectorQuantize(nn.Module):
+    """Vector quantization and factorized vecotor quantization implementation
+    Args:
+        input_dim (int): Dimension of input.
+        codebook_size (int): Codebook size.
+        codebook_dim (int): Codebook dimension. We suggest use codebook_dim = input_dim
+            if use codebook_type == "euclidean", otherwise, if you want to use
+            factorized vector quantization, use codebook_dim as small number (e.g. 8 or 32).
+        commitment (float): Weight for commitment loss.
+        use_l2_normlize (bool): Whether to use l2 normlized codes for factorized vecotor quantization,
+            we suggest use it as True if you want to use factorized vector quantization
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+
+    def __init__(
+        self,
+        input_dim,
+        codebook_size,
+        codebook_dim,
+        commitment=0.005,
+        codebook_loss_weight=1.0,
+        use_l2_normlize=False,
+        codebook_type="euclidean",  # "euclidean" or "simple"
+        kmeans_init=False,
+        kmeans_iters=10,
+        decay=0.8,
+        eps=1e-5,
+        threshold_ema_dead_code=2,
+        weight_init=False,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.commitment = commitment
+        self.codebook_loss_weight = codebook_loss_weight
+        self.use_l2_normlize = use_l2_normlize
+        self.codebook_type = codebook_type
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.decay = decay
+        self.eps = eps
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.weight_init = weight_init
+
+        if self.input_dim != self.codebook_dim:
+            self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
+            self.out_project = WNConv1d(
+                self.codebook_dim, self.input_dim, kernel_size=1
+            )
+
+        else:
+            self.in_project = nn.Identity()
+            self.out_project = nn.Identity()
+
+        if self.codebook_type == "euclidean":
+            self.codebook = EuclideanCodebook(
+                self.codebook_dim,
+                codebook_size=self.codebook_size,
+                kmeans_init=self.kmeans_init,
+                kmeans_iters=self.kmeans_iters,
+                decay=self.decay,
+                eps=self.eps,
+                threshold_ema_dead_code=self.threshold_ema_dead_code,
+                weight_init=self.weight_init,
+            )
+        elif self.codebook_type == "simple":
+            self.codebook = SimpleCodebook(
+                self.codebook_dim,
+                codebook_size=self.codebook_size,
+                use_l2_normlize=self.use_l2_normlize,
+            )
+        else:
+            raise NotImplementedError(
+                f"codebook_type {self.codebook_type} is not implemented!"
+            )
+
+    def forward(self, z):
+        """
+        Parameters
+        ----------
+        z: torch.Tensor[B x D x T]
+
+        Returns
+        -------
+        z_q: torch.Tensor[B x D x T]
+            Quantized continuous representation of input
+        commit_loss: Tensor[B]
+            Commitment loss to train encoder to predict vectors closer to codebook entries
+        codebook_loss: Tensor[B]
+            Codebook loss to update the codebook
+        indices: torch.Tensor[B x T]
+            Codebook indices (quantized discrete representation of input)
+        z_e: torch.Tensor[B x D x T]
+            Projected latents (continuous representation of input before quantization)
+        """
+
+        # Factorized codes project input into low-dimensional space if self.input_dim != self.codebook_dim
+        z_e = self.in_project(z)
+        z_q, indices = self.decode_latents(z_e)
+
+        # Compute commitment loss and codebook loss
+        if self.training:
+            commit_loss = (
+                F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+                * self.commitment
+            )
+            codebook_loss = (
+                F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+                * self.codebook_loss_weight
+            )
+        else:
+            commit_loss = torch.zeros(z.shape[0], device=z.device)
+            codebook_loss = torch.zeros(z.shape[0], device=z.device)
+
+        z_q = z_e + (z_q - z_e).detach()
+
+        z_q = self.out_project(z_q)
+
+        return z_q, commit_loss, codebook_loss, indices, z_e
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> b t d")
+        z_q, indices = self.codebook(encodings)
+        z_q = z_q.transpose(1, 2)
+        return z_q, indices
+
+    def vq2emb(self, vq, out_proj=True):
+        emb = self.codebook.vq2emb(vq)
+        emb = emb.transpose(1, 2)
+        if out_proj:
+            emb = self.out_project(emb)
+        return emb
+
+    def latent2dist(self, latents):
+        latents = rearrange(latents, "b d t -> b t d")
+        dist, embed_ind, quantize = self.codebook.latent2dist(latents)
+        return dist, embed_ind, quantize.transpose(1, 2)

models/codec/amphion_codec/vocos.py

@@ -0,0 +1,881 @@
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional, Tuple
+
+import numpy as np
+import scipy
+import torch
+from torch import nn, view_as_real, view_as_complex
+from torch import nn
+from torch.nn.utils import weight_norm, remove_weight_norm
+from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz
+import librosa
+
+
+def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor:
+    """
+    Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values.
+
+    Args:
+        x (Tensor): Input tensor.
+        clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7.
+
+    Returns:
+        Tensor: Element-wise logarithm of the input tensor with clipping applied.
+    """
+    return torch.log(torch.clip(x, min=clip_val))
+
+
+def symlog(x: torch.Tensor) -> torch.Tensor:
+    return torch.sign(x) * torch.log1p(x.abs())
+
+
+def symexp(x: torch.Tensor) -> torch.Tensor:
+    return torch.sign(x) * (torch.exp(x.abs()) - 1)
+
+
+class STFT(nn.Module):
+    def __init__(
+        self,
+        n_fft: int,
+        hop_length: int,
+        win_length: int,
+        center=True,
+    ):
+        super().__init__()
+        self.center = center
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, T * hop_length)
+
+        if not self.center:
+            pad = self.win_length - self.hop_length
+            x = torch.nn.functional.pad(x, (pad // 2, pad // 2), mode="reflect")
+
+        stft_spec = torch.stft(
+            x,
+            self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=self.window,
+            center=self.center,
+            return_complex=False,
+        )  # (B, n_fft // 2 + 1, T, 2)
+
+        rea = stft_spec[:, :, :, 0]  # (B, n_fft // 2 + 1, T, 2)
+        imag = stft_spec[:, :, :, 1]  # (B, n_fft // 2 + 1, T, 2)
+
+        log_mag = torch.log(
+            torch.abs(torch.sqrt(torch.pow(rea, 2) + torch.pow(imag, 2))) + 1e-5
+        )  # (B, n_fft // 2 + 1, T)
+        phase = torch.atan2(imag, rea)  # (B, n_fft // 2 + 1, T)
+
+        return log_mag, phase
+
+
+class ISTFT(nn.Module):
+    """
+    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+    See issue: https://github.com/pytorch/pytorch/issues/62323
+    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+    The NOLA constraint is met as we trim padded samples anyway.
+
+    Args:
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames.
+        win_length (int): The size of window frame and STFT filter.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(
+        self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"
+    ):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+        Args:
+            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+                            N is the number of frequency bins, and T is the number of time frames.
+
+        Returns:
+            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
+        """
+        if self.padding == "center":
+            # Fallback to pytorch native implementation
+            return torch.istft(
+                spec,
+                self.n_fft,
+                self.hop_length,
+                self.win_length,
+                self.window,
+                center=True,
+            )
+        elif self.padding == "same":
+            pad = (self.win_length - self.hop_length) // 2
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        assert spec.dim() == 3, "Expected a 3D tensor as input"
+        B, N, T = spec.shape
+
+        # Inverse FFT
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        # Overlap and Add
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        )[:, 0, 0, pad:-pad]
+
+        # Window envelope
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = torch.nn.functional.fold(
+            window_sq,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        ).squeeze()[pad:-pad]
+
+        # Normalize
+        assert (window_envelope > 1e-11).all()
+        y = y / window_envelope
+
+        return y
+
+
+class MDCT(nn.Module):
+    """
+    Modified Discrete Cosine Transform (MDCT) module.
+
+    Args:
+        frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, frame_len: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.frame_len = frame_len
+        N = frame_len // 2
+        n0 = (N + 1) / 2
+        window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
+        self.register_buffer("window", window)
+
+        pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len)
+        post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N)
+        # view_as_real: NCCL Backend does not support ComplexFloat data type
+        # https://github.com/pytorch/pytorch/issues/71613
+        self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
+        self.register_buffer("post_twiddle", view_as_real(post_twiddle))
+
+    def forward(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the Modified Discrete Cosine Transform (MDCT) to the input audio.
+
+        Args:
+            audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size
+                and T is the length of the audio.
+
+        Returns:
+            Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames
+                and N is the number of frequency bins.
+        """
+        if self.padding == "center":
+            audio = torch.nn.functional.pad(
+                audio, (self.frame_len // 2, self.frame_len // 2)
+            )
+        elif self.padding == "same":
+            # hop_length is 1/2 frame_len
+            audio = torch.nn.functional.pad(
+                audio, (self.frame_len // 4, self.frame_len // 4)
+            )
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        x = audio.unfold(-1, self.frame_len, self.frame_len // 2)
+        N = self.frame_len // 2
+        x = x * self.window.expand(x.shape)
+        X = torch.fft.fft(
+            x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1
+        )[..., :N]
+        res = X * view_as_complex(self.post_twiddle).expand(X.shape) * np.sqrt(1 / N)
+        return torch.real(res) * np.sqrt(2)
+
+
+class IMDCT(nn.Module):
+    """
+    Inverse Modified Discrete Cosine Transform (IMDCT) module.
+
+    Args:
+        frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, frame_len: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.frame_len = frame_len
+        N = frame_len // 2
+        n0 = (N + 1) / 2
+        window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
+        self.register_buffer("window", window)
+
+        pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N)
+        post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2))
+        self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
+        self.register_buffer("post_twiddle", view_as_real(post_twiddle))
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients.
+
+        Args:
+            X (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size,
+                L is the number of frames, and N is the number of frequency bins.
+
+        Returns:
+            Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio.
+        """
+        B, L, N = X.shape
+        Y = torch.zeros((B, L, N * 2), dtype=X.dtype, device=X.device)
+        Y[..., :N] = X
+        Y[..., N:] = -1 * torch.conj(torch.flip(X, dims=(-1,)))
+        y = torch.fft.ifft(
+            Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1
+        )
+        y = (
+            torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape))
+            * np.sqrt(N)
+            * np.sqrt(2)
+        )
+        result = y * self.window.expand(y.shape)
+        output_size = (1, (L + 1) * N)
+        audio = torch.nn.functional.fold(
+            result.transpose(1, 2),
+            output_size=output_size,
+            kernel_size=(1, self.frame_len),
+            stride=(1, self.frame_len // 2),
+        )[:, 0, 0, :]
+
+        if self.padding == "center":
+            pad = self.frame_len // 2
+        elif self.padding == "same":
+            pad = self.frame_len // 4
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        audio = audio[:, pad:-pad]
+        return audio
+
+
+class FourierHead(nn.Module):
+    """Base class for inverse fourier modules."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class ISTFTHead(FourierHead):
+    """
+    ISTFT Head module for predicting STFT complex coefficients.
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames, which should align with
+                          the resolution of the input features.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
+        super().__init__()
+        out_dim = n_fft + 2
+        self.out = torch.nn.Linear(dim, out_dim)
+        self.istft = ISTFT(
+            n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the ISTFTHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x).transpose(1, 2)
+        mag, p = x.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(
+            mag, max=1e2
+        )  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value
+        S = mag * (x + 1j * y)
+        audio = self.istft(S)
+        return audio
+
+
+class IMDCTSymExpHead(FourierHead):
+    """
+    IMDCT Head module for predicting MDCT coefficients with symmetric exponential function
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        mdct_frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+        sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized
+                                     based on perceptual scaling. Defaults to None.
+        clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        mdct_frame_len: int,
+        padding: str = "same",
+        sample_rate: Optional[int] = None,
+        clip_audio: bool = False,
+    ):
+        super().__init__()
+        out_dim = mdct_frame_len // 2
+        self.out = nn.Linear(dim, out_dim)
+        self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
+        self.clip_audio = clip_audio
+
+        if sample_rate is not None:
+            # optionally init the last layer following mel-scale
+            m_max = _hz_to_mel(sample_rate // 2)
+            m_pts = torch.linspace(0, m_max, out_dim)
+            f_pts = _mel_to_hz(m_pts)
+            scale = 1 - (f_pts / f_pts.max())
+
+            with torch.no_grad():
+                self.out.weight.mul_(scale.view(-1, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the IMDCTSymExpHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x)
+        x = symexp(x)
+        x = torch.clip(
+            x, min=-1e2, max=1e2
+        )  # safeguard to prevent excessively large magnitudes
+        audio = self.imdct(x)
+        if self.clip_audio:
+            audio = torch.clip(x, min=-1.0, max=1.0)
+
+        return audio
+
+
+class IMDCTCosHead(FourierHead):
+    """
+    IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p)
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        mdct_frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+        clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        mdct_frame_len: int,
+        padding: str = "same",
+        clip_audio: bool = False,
+    ):
+        super().__init__()
+        self.clip_audio = clip_audio
+        self.out = nn.Linear(dim, mdct_frame_len)
+        self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the IMDCTCosHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x)
+        m, p = x.chunk(2, dim=2)
+        m = torch.exp(m).clip(
+            max=1e2
+        )  # safeguard to prevent excessively large magnitudes
+        audio = self.imdct(m * torch.cos(p))
+        if self.clip_audio:
+            audio = torch.clip(x, min=-1.0, max=1.0)
+        return audio
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+    Args:
+        dim (int): Number of input channels.
+        intermediate_dim (int): Dimensionality of the intermediate layer.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+            None means non-conditional LayerNorm. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        layer_scale_init_value: float,
+        adanorm_num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(
+            dim, dim, kernel_size=7, padding=3, groups=dim
+        )  # depthwise conv
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(
+            dim, intermediate_dim
+        )  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(
+        self, x: torch.Tensor, cond_embedding_id: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+        if self.adanorm:
+            assert cond_embedding_id is not None
+            x = self.norm(x, cond_embedding_id)
+        else:
+            x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+
+        x = residual + x
+        return x
+
+
+class AdaLayerNorm(nn.Module):
+    """
+    Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes
+
+    Args:
+        num_embeddings (int): Number of embeddings.
+        embedding_dim (int): Dimension of the embeddings.
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.dim = embedding_dim
+        self.scale = nn.Embedding(
+            num_embeddings=num_embeddings, embedding_dim=embedding_dim
+        )
+        self.shift = nn.Embedding(
+            num_embeddings=num_embeddings, embedding_dim=embedding_dim
+        )
+        torch.nn.init.ones_(self.scale.weight)
+        torch.nn.init.zeros_(self.shift.weight)
+
+    def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor:
+        scale = self.scale(cond_embedding_id)
+        shift = self.shift(cond_embedding_id)
+        x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps)
+        x = x * scale + shift
+        return x
+
+
+class ResBlock1(nn.Module):
+    """
+    ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions,
+    but without upsampling layers.
+
+    Args:
+        dim (int): Number of input channels.
+        kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3.
+        dilation (tuple[int], optional): Dilation factors for the dilated convolutions.
+            Defaults to (1, 3, 5).
+        lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function.
+            Defaults to 0.1.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        kernel_size: int = 3,
+        dilation: Tuple[int, int, int] = (1, 3, 5),
+        lrelu_slope: float = 0.1,
+        layer_scale_init_value: Optional[float] = None,
+    ):
+        super().__init__()
+        self.lrelu_slope = lrelu_slope
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=self.get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=self.get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=self.get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=self.get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=self.get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=self.get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+
+        self.gamma = nn.ParameterList(
+            [
+                (
+                    nn.Parameter(
+                        layer_scale_init_value * torch.ones(dim, 1), requires_grad=True
+                    )
+                    if layer_scale_init_value is not None
+                    else None
+                ),
+                (
+                    nn.Parameter(
+                        layer_scale_init_value * torch.ones(dim, 1), requires_grad=True
+                    )
+                    if layer_scale_init_value is not None
+                    else None
+                ),
+                (
+                    nn.Parameter(
+                        layer_scale_init_value * torch.ones(dim, 1), requires_grad=True
+                    )
+                    if layer_scale_init_value is not None
+                    else None
+                ),
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma):
+            xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope)
+            xt = c1(xt)
+            xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope)
+            xt = c2(xt)
+            if gamma is not None:
+                xt = gamma * 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)
+
+    @staticmethod
+    def get_padding(kernel_size: int, dilation: int = 1) -> int:
+        return int((kernel_size * dilation - dilation) / 2)
+
+
+class Backbone(nn.Module):
+    """Base class for the generator's backbone. It preserves the same temporal resolution across all layers."""
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
+                        C denotes output features, and L is the sequence length.
+
+        Returns:
+            Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
+                    and H denotes the model dimension.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class VocosBackbone(Backbone):
+    """
+    Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+        num_layers (int): Number of ConvNeXtBlock layers.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+                                                None means non-conditional model. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        dim: int,
+        intermediate_dim: int,
+        num_layers: int,
+        layer_scale_init_value: Optional[float] = None,
+        adanorm_num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        layer_scale_init_value = layer_scale_init_value or 1 / num_layers
+        self.convnext = nn.ModuleList(
+            [
+                ConvNeXtBlock(
+                    dim=dim,
+                    intermediate_dim=intermediate_dim,
+                    layer_scale_init_value=layer_scale_init_value,
+                    adanorm_num_embeddings=adanorm_num_embeddings,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        bandwidth_id = kwargs.get("bandwidth_id", None)
+        x = self.embed(x)
+        if self.adanorm:
+            assert bandwidth_id is not None
+            x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id)
+        else:
+            x = self.norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        for conv_block in self.convnext:
+            x = conv_block(x, cond_embedding_id=bandwidth_id)
+        x = self.final_layer_norm(x.transpose(1, 2))
+        return x
+
+
+class VocosResNetBackbone(Backbone):
+    """
+    Vocos backbone module built with ResBlocks.
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        num_blocks (int): Number of ResBlock1 blocks.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        input_channels,
+        dim,
+        num_blocks,
+        layer_scale_init_value=None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.embed = weight_norm(
+            nn.Conv1d(input_channels, dim, kernel_size=3, padding=1)
+        )
+        layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3
+        self.resnet = nn.Sequential(
+            *[
+                ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value)
+                for _ in range(num_blocks)
+            ]
+        )
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        x = self.embed(x)
+        x = self.resnet(x)
+        x = x.transpose(1, 2)
+        return x
+
+
+class Vocos(nn.Module):
+    def __init__(
+        self,
+        input_channels: int = 256,
+        dim: int = 384,
+        intermediate_dim: int = 1152,
+        num_layers: int = 8,
+        n_fft: int = 800,
+        hop_size: int = 200,
+        padding: str = "same",
+        adanorm_num_embeddings=None,
+        cfg=None,
+    ):
+        super().__init__()
+
+        input_channels = (
+            cfg.input_channels
+            if cfg is not None and hasattr(cfg, "input_channels")
+            else input_channels
+        )
+        dim = cfg.dim if cfg is not None and hasattr(cfg, "dim") else dim
+        intermediate_dim = (
+            cfg.intermediate_dim
+            if cfg is not None and hasattr(cfg, "intermediate_dim")
+            else intermediate_dim
+        )
+        num_layers = (
+            cfg.num_layers
+            if cfg is not None and hasattr(cfg, "num_layers")
+            else num_layers
+        )
+        adanorm_num_embeddings = (
+            cfg.adanorm_num_embeddings
+            if cfg is not None and hasattr(cfg, "adanorm_num_embeddings")
+            else adanorm_num_embeddings
+        )
+        n_fft = cfg.n_fft if cfg is not None and hasattr(cfg, "n_fft") else n_fft
+        hop_size = (
+            cfg.hop_size if cfg is not None and hasattr(cfg, "hop_size") else hop_size
+        )
+        padding = (
+            cfg.padding if cfg is not None and hasattr(cfg, "padding") else padding
+        )
+
+        self.backbone = VocosBackbone(
+            input_channels=input_channels,
+            dim=dim,
+            intermediate_dim=intermediate_dim,
+            num_layers=num_layers,
+            adanorm_num_embeddings=adanorm_num_embeddings,
+        )
+        self.head = ISTFTHead(dim, n_fft, hop_size, padding)
+
+    def forward(self, x):
+        x = self.backbone(x)
+        x = self.head(x)
+
+        return x[:, None, :]

models/codec/coco/rep_coco_model.py

@@ -0,0 +1,441 @@
+# Copyright (c) 2024 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from concurrent.futures import ALL_COMPLETED
+import numpy as np
+import torch
+import torch.nn as nn
+
+from torch.nn import functional as F
+
+from models.codec.amphion_codec.quantize import ResidualVQ
+from models.codec.amphion_codec.vocos import VocosBackbone
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+    if isinstance(m, nn.Linear):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+def compute_codebook_perplexity(indices, codebook_size):
+    indices = indices.flatten()
+    prob = torch.bincount(indices, minlength=codebook_size).float() / indices.size(0)
+    perp = torch.exp(-torch.sum(prob * torch.log(prob + 1e-10)))
+    return perp
+
+
+class CocoContentStyle(nn.Module):
+    def __init__(
+        self,
+        codebook_size=8192,
+        hidden_size=1024,
+        codebook_dim=8,
+        num_quantizers=1,
+        quantizer_type="fvq",
+        use_whisper=True,
+        use_chromagram=True,
+        construct_only_for_quantizer=False,
+        cfg=None,
+    ):
+        super().__init__()
+
+        assert cfg is not None
+        self.cfg = cfg
+
+        codebook_size = getattr(cfg, "codebook_size", codebook_size)
+        hidden_size = getattr(cfg, "hidden_size", hidden_size)
+        codebook_dim = getattr(cfg, "codebook_dim", codebook_dim)
+        num_quantizers = getattr(cfg, "num_quantizers", num_quantizers)
+        quantizer_type = getattr(cfg, "quantizer_type", quantizer_type)
+
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.hidden_size = hidden_size
+        self.num_quantizers = num_quantizers
+        self.quantizer_type = quantizer_type
+
+        if use_whisper:
+            self.whisper_input_layer = nn.Linear(self.cfg.whisper_dim, hidden_size)
+        if use_chromagram:
+            self.chromagram_input_layer = nn.Linear(
+                self.cfg.chromagram_dim, hidden_size
+            )
+
+        downsample_rate = getattr(cfg, "downsample_rate", 1)
+        if downsample_rate > 1:
+            self.do_downsample = True
+            assert np.log2(downsample_rate).is_integer()
+
+            down_layers = []
+            up_layers = []
+            for _ in range(int(np.log2(downsample_rate))):
+                down_layers.extend(
+                    [
+                        nn.Conv1d(
+                            hidden_size,
+                            hidden_size,
+                            kernel_size=3,
+                            stride=2,
+                            padding=1,
+                        ),
+                        nn.GELU(),
+                    ]
+                )
+                up_layers.extend(
+                    [
+                        nn.ConvTranspose1d(
+                            hidden_size, hidden_size, kernel_size=4, stride=2, padding=1
+                        ),
+                        nn.GELU(),
+                    ]
+                )
+            self.downsample_layers = nn.Sequential(*down_layers)
+            self.upsample_layers = nn.Sequential(*up_layers)
+
+        else:
+            self.do_downsample = False
+
+        self.encoder = nn.Sequential(
+            VocosBackbone(
+                input_channels=self.hidden_size,
+                dim=self.cfg.encoder.vocos_dim,
+                intermediate_dim=self.cfg.encoder.vocos_intermediate_dim,
+                num_layers=self.cfg.encoder.vocos_num_layers,
+                adanorm_num_embeddings=None,
+            ),
+            nn.Linear(self.cfg.encoder.vocos_dim, self.hidden_size),
+        )
+
+        self.quantizer = ResidualVQ(
+            input_dim=hidden_size,
+            num_quantizers=num_quantizers,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            quantizer_type=quantizer_type,
+            quantizer_dropout=0.0,
+            commitment=0.15,
+            codebook_loss_weight=1.0,
+            use_l2_normlize=True,
+        )
+
+        if not construct_only_for_quantizer:
+            self.decoder = nn.Sequential(
+                VocosBackbone(
+                    input_channels=self.hidden_size,
+                    dim=self.cfg.decoder.vocos_dim,
+                    intermediate_dim=self.cfg.decoder.vocos_intermediate_dim,
+                    num_layers=self.cfg.decoder.vocos_num_layers,
+                    adanorm_num_embeddings=None,
+                ),
+                nn.Linear(self.cfg.decoder.vocos_dim, self.hidden_size),
+            )
+
+            if use_whisper:
+                self.whisper_output_layer = nn.Linear(
+                    self.hidden_size, self.cfg.whisper_dim
+                )
+            if use_chromagram:
+                self.chromagram_output_layer = nn.Linear(
+                    self.hidden_size, self.cfg.chromagram_dim
+                )
+
+        self.reset_parameters()
+
+    def forward(
+        self,
+        whisper_feats,
+        chromagram_feats,
+        return_for_quantizer=False,
+    ):
+        """
+        Args:
+            whisper_feats: [B, T, 1024]
+            chromagram_feats: [B, T, 24]
+        Returns:
+            whisper_rec: [B, T, 1024]
+            chromagram_rec: [B, T, 24]
+            codebook_loss: float
+            all_indices: [N, B, T] or [B, T] if num_of_quantizers == 1
+        """
+        T = whisper_feats.shape[1]
+
+        # [B, T, D]
+        x = self.whisper_input_layer(whisper_feats) + self.chromagram_input_layer(
+            chromagram_feats
+        )
+        # print("Before downsample:", x.shape)
+
+        # ====== Downsample ======
+        if self.do_downsample:
+            x = self.downsample_layers(x.transpose(1, 2)).transpose(1, 2)
+
+        # print("After downsample:", x.shape)
+
+        # ====== Encoder ======
+        x = self.encoder(x.transpose(1, 2)).transpose(1, 2)  # [B, T, D] -> [B, D, T]
+
+        # ====== Quantizer ======
+        (
+            quantized_out,  # [B, D, T]
+            all_indices,  # [num_of_quantizers, B, T]
+            all_commit_losses,  # [num_of_quantizers]
+            all_codebook_losses,  # [num_of_quantizers]
+            _,
+        ) = self.quantizer(x)
+
+        if return_for_quantizer:
+            if all_indices.shape[0] == 1:
+                return all_indices.squeeze(0), quantized_out.transpose(1, 2)
+            return all_indices, quantized_out.transpose(1, 2)
+
+        # ====== Decoder ======
+        x_rec = self.decoder(quantized_out)  # [B, T, D]
+
+        # ====== Upsample ======
+        if self.do_downsample:
+            x_rec = self.upsample_layers(x_rec.transpose(1, 2)).transpose(1, 2)
+
+        # print("After upsample:", x_rec.shape)
+
+        # Ensure output dimensions match input
+        if x_rec.shape[1] >= T:  # Check time dimension
+            x_rec = x_rec[:, :T, :]
+        else:
+            padding_frames = T - x_rec.shape[1]
+            last_frame = x_rec[:, -1:, :]
+            padding = last_frame.repeat(1, padding_frames, 1)
+            x_rec = torch.cat([x_rec, padding], dim=1)
+
+        # ====== Loss ======
+        whisper_rec = self.whisper_output_layer(x_rec)  # [B, T, 1024]
+        chromagram_rec = self.chromagram_output_layer(x_rec)  # [B, T, 24]
+
+        codebook_loss = (all_codebook_losses + all_commit_losses).mean()
+        all_indices = all_indices
+
+        return whisper_rec, chromagram_rec, codebook_loss, all_indices
+
+    def quantize(self, whisper_feats, chromagram_feats):
+        """
+        Args:
+            whisper_feats: [B, T, 1024]
+            chromagram_feats: [B, T, 24]
+        Returns:
+            all_indices: [N, B, T], or [B, T] if num_of_quantizers == 1
+            quantized_out: [B, D, T]
+        """
+        all_indices, quantized_out = self.forward(
+            whisper_feats,
+            chromagram_feats,
+            return_for_quantizer=True,
+        )
+        return all_indices, quantized_out
+
+    def reset_parameters(self):
+        self.apply(init_weights)
+
+
+class CocoContent(CocoContentStyle):
+    def __init__(
+        self,
+        cfg,
+        use_whisper=True,
+        use_chromagram=False,
+        construct_only_for_quantizer=False,
+    ):
+        super().__init__(
+            cfg=cfg,
+            use_whisper=use_whisper,
+            use_chromagram=use_chromagram,
+            construct_only_for_quantizer=construct_only_for_quantizer,
+        )
+
+    def forward(
+        self,
+        whisper_feats,
+        return_for_quantizer=False,
+    ):
+        """
+        Args:
+            whisper_feats: [B, T, 1024]
+        Returns:
+            whisper_rec: [B, T, 1024]
+            codebook_loss: float
+            all_indices: [N, B, T]
+        """
+        T = whisper_feats.shape[1]
+
+        # [B, T, D]
+        x = self.whisper_input_layer(whisper_feats)
+
+        # ====== Downsample ======
+        if self.do_downsample:
+            x = self.downsample_layers(x.transpose(1, 2)).transpose(1, 2)
+
+        # ====== Encoder ======
+        x = self.encoder(x.transpose(1, 2)).transpose(1, 2)  # [B, T, D] -> [B, D, T]
+
+        # ====== Quantizer ======
+        (
+            quantized_out,  # [B, D, T]
+            all_indices,  # [num_of_quantizers, B, T]
+            all_commit_losses,  # [num_of_quantizers]
+            all_codebook_losses,  # [num_of_quantizers]
+            _,
+        ) = self.quantizer(x)
+
+        if return_for_quantizer:
+            if all_indices.shape[0] == 1:
+                return all_indices.squeeze(0), quantized_out.transpose(1, 2)
+            return all_indices, quantized_out.transpose(1, 2)
+
+        # ====== Decoder ======
+        x_rec = self.decoder(quantized_out)  # [B, T, D]
+
+        # ====== Upsample ======
+        if self.do_downsample:
+            x_rec = self.upsample_layers(x_rec.transpose(1, 2)).transpose(1, 2)
+
+        # Ensure output dimensions match input
+        if x_rec.shape[1] >= T:  # Check time dimension
+            x_rec = x_rec[:, :T, :]
+        else:
+            padding_frames = T - x_rec.shape[1]
+            last_frame = x_rec[:, -1:, :]
+            padding = last_frame.repeat(1, padding_frames, 1)
+            x_rec = torch.cat([x_rec, padding], dim=1)
+
+        # ====== Loss ======
+        whisper_rec = self.whisper_output_layer(x_rec)  # [B, T, 1024]
+
+        codebook_loss = (all_codebook_losses + all_commit_losses).mean()
+        all_indices = all_indices
+
+        return whisper_rec, codebook_loss, all_indices
+
+    def quantize(self, whisper_feats):
+        all_indices, quantized_out = self.forward(
+            whisper_feats, return_for_quantizer=True
+        )
+        return all_indices, quantized_out
+
+
+class CocoStyle(CocoContentStyle):
+    def __init__(
+        self,
+        cfg,
+        use_whisper=False,
+        use_chromagram=True,
+        construct_only_for_quantizer=False,
+    ):
+        super().__init__(
+            cfg=cfg,
+            use_whisper=use_whisper,
+            use_chromagram=use_chromagram,
+            construct_only_for_quantizer=construct_only_for_quantizer,
+        )
+
+    def forward(
+        self,
+        chromagram_feats,
+        return_for_quantizer=False,
+    ):
+        """
+        Args:
+            chromagram_feats: [B, T, 24]
+        Returns:
+            chromagram_rec: [B, T, 24]
+            codebook_loss: float
+            all_indices: [N, B, T]
+        """
+        T = chromagram_feats.shape[1]
+
+        # [B, T, D]
+        x = self.chromagram_input_layer(chromagram_feats)
+
+        # ====== Downsample ======
+        if self.do_downsample:
+            x = self.downsample_layers(x.transpose(1, 2)).transpose(1, 2)
+
+        # ====== Encoder ======
+        x = self.encoder(x.transpose(1, 2)).transpose(1, 2)  # [B, T, D] -> [B, D, T]
+
+        # ====== Quantizer ======
+        (
+            quantized_out,  # [B, D, T]
+            all_indices,  # [num_of_quantizers, B, T]
+            all_commit_losses,  # [num_of_quantizers]
+            all_codebook_losses,  # [num_of_quantizers]
+            _,
+        ) = self.quantizer(x)
+
+        if return_for_quantizer:
+            if all_indices.shape[0] == 1:
+                return all_indices.squeeze(0), quantized_out.transpose(1, 2)
+            return all_indices, quantized_out.transpose(1, 2)
+
+        # ====== Decoder ======
+        x_rec = self.decoder(quantized_out)  # [B, T, D]
+
+        # ====== Upsample ======
+        if self.do_downsample:
+            x_rec = self.upsample_layers(x_rec.transpose(1, 2)).transpose(1, 2)
+
+        # Ensure output dimensions match input
+        if x_rec.shape[1] >= T:  # Check time dimension
+            x_rec = x_rec[:, :T, :]
+        else:
+            padding_frames = T - x_rec.shape[1]
+            last_frame = x_rec[:, -1:, :]
+            padding = last_frame.repeat(1, padding_frames, 1)
+            x_rec = torch.cat([x_rec, padding], dim=1)
+
+        # ====== Loss ======
+        chromagram_rec = self.chromagram_output_layer(x_rec)  # [B, T, 24]
+
+        codebook_loss = (all_codebook_losses + all_commit_losses).mean()
+        all_indices = all_indices
+
+        return chromagram_rec, codebook_loss, all_indices
+
+    def quantize(self, chromagram_feats):
+        all_indices, quantized_out = self.forward(
+            chromagram_feats, return_for_quantizer=True
+        )
+        return all_indices, quantized_out
+
+
+# if __name__ == "__main__":
+#     from utils.util import JsonHParams
+
+#     cfg = JsonHParams(
+#         **{
+#             "whisper_dim": 1024,
+#             "chromagram_dim": 24,
+#             "global_speaker_encoder": {
+#                 "input_dim": 128,  # Eg: n_mels
+#                 "hidden_size": 512,  # 768 for emilia298k
+#                 "num_hidden_layers": 4,  # 6 for emilia298k
+#                 "num_attention_heads": 8,
+#             },
+#         }
+#     )
+#     model = Coco(cfg=cfg)
+
+#     x = torch.randn(2, 150, 1024)
+#     tone_height = torch.randn(2)
+#     mels = torch.randn(2, 150, 128)
+#     mel_mask = torch.ones(2, 150)
+
+#     x_rec, codebook_loss, all_indices, auxillary_pred_outputs = model(
+#         x, tone_height, mels, mel_mask
+#     )
+#     print(x_rec.shape, codebook_loss, all_indices.shape)
+#     for k, v in auxillary_pred_outputs.items():
+#         print(k, v.shape)

models/codec/melvqgan/melspec.py

@@ -0,0 +1,108 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import pyworld as pw
+import numpy as np
+import soundfile as sf
+import os
+from torchaudio.functional import pitch_shift
+import librosa
+from librosa.filters import mel as librosa_mel_fn
+import torch.nn as nn
+import torch.nn.functional as F
+import tqdm
+
+
+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 MelSpectrogram(nn.Module):
+    def __init__(
+        self,
+        n_fft,
+        num_mels,
+        sampling_rate,
+        hop_size,
+        win_size,
+        fmin,
+        fmax,
+        center=False,
+    ):
+        super(MelSpectrogram, self).__init__()
+        self.n_fft = n_fft
+        self.hop_size = hop_size
+        self.win_size = win_size
+        self.sampling_rate = sampling_rate
+        self.num_mels = num_mels
+        self.fmin = fmin
+        self.fmax = fmax
+        self.center = center
+
+        mel_basis = {}
+        hann_window = {}
+
+        mel = librosa_mel_fn(
+            sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
+        )
+        mel_basis = torch.from_numpy(mel).float()
+        hann_window = torch.hann_window(win_size)
+
+        self.register_buffer("mel_basis", mel_basis)
+        self.register_buffer("hann_window", hann_window)
+
+    def forward(self, y):
+        y = torch.nn.functional.pad(
+            y.unsqueeze(1),
+            (
+                int((self.n_fft - self.hop_size) / 2),
+                int((self.n_fft - self.hop_size) / 2),
+            ),
+            mode="reflect",
+        )
+        y = y.squeeze(1)
+        spec = torch.stft(
+            y,
+            self.n_fft,
+            hop_length=self.hop_size,
+            win_length=self.win_size,
+            window=self.hann_window,
+            center=self.center,
+            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, spec)
+        spec = spectral_normalize_torch(spec)
+
+        return spec

models/svc/flow_matching_transformer/fmt_model.py

@@ -0,0 +1,328 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import numpy as np
+import torch.nn as nn
+import math
+from einops import rearrange
+from models.vc.flow_matching_transformer.llama_nar import DiffLlama
+import torch.nn.functional as F
+
+
+class FlowMatchingTransformer(nn.Module):
+    def __init__(
+        self,
+        mel_dim=100,
+        hidden_size=1024,
+        num_layers=12,
+        num_heads=16,
+        cfg_scale=0.2,
+        cond_codebook_size=1024,
+        cond_scale_factor=1,
+        sigma=1e-5,
+        time_scheduler="linear",
+        cfg=None,
+    ):
+        super().__init__()
+        self.cfg = cfg
+
+        if cfg is not None:
+            mel_dim = getattr(cfg, "mel_dim", mel_dim)
+            hidden_size = getattr(cfg, "hidden_size", hidden_size)
+            num_layers = getattr(cfg, "num_layers", num_layers)
+            num_heads = getattr(cfg, "num_heads", num_heads)
+            cfg_scale = getattr(cfg, "cfg_scale", cfg_scale)
+            cond_codebook_size = getattr(cfg, "cond_codebook_size", cond_codebook_size)
+            time_scheduler = getattr(cfg, "time_scheduler", time_scheduler)
+            sigma = getattr(cfg, "sigma", sigma)
+            cond_scale_factor = getattr(cfg, "cond_scale_factor", cond_scale_factor)
+
+        self.mel_dim = mel_dim
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.cfg_scale = cfg_scale
+        self.cond_codebook_size = cond_codebook_size
+        self.time_scheduler = time_scheduler
+        self.sigma = sigma
+        self.cond_scale_factor = cond_scale_factor
+
+        self.cond_emb = nn.Embedding(cond_codebook_size, self.hidden_size)
+
+        if cond_scale_factor != 1:
+            self.do_resampling = True
+            assert np.log2(cond_scale_factor).is_integer()
+
+            up_layers = []
+            for _ in range(int(np.log2(cond_scale_factor))):
+                up_layers.extend(
+                    [
+                        nn.ConvTranspose1d(
+                            hidden_size, hidden_size, kernel_size=4, stride=2, padding=1
+                        ),
+                        nn.GELU(),
+                    ]
+                )
+            self.resampling_layers = nn.Sequential(*up_layers)
+        else:
+            self.do_resampling = False
+
+        ### REPA: Use the Wav2Vec2Bert features to align. ###
+        self.use_repa = "repa" in cfg
+        self.repa_layer_index = None
+        if self.use_repa:
+            self.repa_layer_index = cfg.repa.layer_index
+
+            self.repa_mlp_layer = nn.Sequential(
+                nn.Linear(hidden_size, hidden_size * 4),
+                nn.SiLU(),
+                nn.Linear(hidden_size * 4, cfg.repa.output_dim),
+            )
+
+        ### CTC: Use the ASR loss ###
+        self.use_ctc = "ctc" in cfg
+        self.ctc_layer_index = None
+        if self.use_ctc:
+            self.ctc_layer_index = cfg.ctc.layer_index
+
+            self.ctc_mlp_layer = nn.Sequential(
+                nn.Linear(hidden_size, hidden_size * 4),
+                nn.SiLU(),
+                nn.Linear(hidden_size * 4, cfg.ctc.output_dim),
+            )
+
+        self.reset_parameters()
+
+        self.diff_estimator = DiffLlama(
+            mel_dim=mel_dim,
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            num_layers=num_layers,
+        )
+
+        self.sigma = sigma
+
+    @torch.no_grad()
+    def forward_diffusion(self, x, t):
+        """
+        x: (B, T, mel_dim)
+        t: (B,)
+        """
+        new_t = t
+        t = t.unsqueeze(-1).unsqueeze(-1)
+        z = torch.randn(
+            x.shape, dtype=x.dtype, device=x.device, requires_grad=False
+        )  # (B, T, mel_dim)
+
+        cfg_scale = self.cfg_scale
+
+        # get prompt len
+        if torch.rand(1) > cfg_scale:
+            prompt_len = torch.randint(
+                min(x.shape[1] // 4, 5), int(x.shape[1] * 0.4), (x.shape[0],)
+            ).to(
+                x.device
+            )  # (B,)
+        else:
+            prompt_len = torch.zeros(x.shape[0]).to(x)  # (B,)
+
+        # get is prompt
+        is_prompt = torch.zeros_like(x[:, :, 0])  # (B, T)
+        col_indices = (
+            torch.arange(is_prompt.shape[1])
+            .repeat(is_prompt.shape[0], 1)
+            .to(prompt_len)
+        )  # (B, T)
+        is_prompt[col_indices < prompt_len.unsqueeze(1)] = 1  # (B, T) 1 if prompt
+
+        mask = torch.ones_like(x[:, :, 0])  # mask if 1, not mask if 0
+        mask[is_prompt.bool()] = 0
+        mask = mask[:, :, None]
+
+        # flow matching: xt = (1 - (1 - sigma) * t) * x0 + t * x; where x0 ~ N(0, 1), x is a sample
+        # flow gt: x - (1 - sigma) * x0 = x - (1 - sigma) * noise
+        xt = ((1 - (1 - self.sigma) * t) * z + t * x) * mask + x * (1 - mask)
+
+        return xt, z, new_t, prompt_len, mask
+
+    def loss_t(
+        self,
+        x,
+        x_mask,
+        t,
+        cond=None,
+    ):
+        xt, z, new_t, prompt_len, mask = self.forward_diffusion(x, t)
+
+        noise = z
+
+        # drop all condition for cfg, so if prompt_len is 0, we also drop cond
+        if cond is not None:
+            cond = cond * torch.where(
+                prompt_len > 0,
+                torch.ones_like(prompt_len),
+                torch.zeros_like(prompt_len),
+            ).to(cond.device).unsqueeze(-1).unsqueeze(-1)
+
+        dit_output = self.diff_estimator(xt, new_t, cond, x_mask, return_dict=True)
+        flow_pred = dit_output["output"]  # (B, T, mel_dim)
+
+        # final mask used for loss calculation
+        final_mask = mask * x_mask[..., None]  # (B, T, 1)
+
+        results = {"output": (noise, x, flow_pred, final_mask, prompt_len)}
+
+        if self.use_repa:
+            repa_hidden_states = dit_output["hidden_states"][
+                self.repa_layer_index
+            ]  # (B, T, hidden_size)
+
+            repa_pred = self.repa_mlp_layer(repa_hidden_states)  # (B, T, repa_dim)
+            results["repa"] = repa_pred
+
+        if self.use_ctc:
+            ctc_hidden_states = dit_output["hidden_states"][
+                self.ctc_layer_index
+            ]  # (B, T, hidden_size)
+            ctc_pred = self.ctc_mlp_layer(ctc_hidden_states)  # (B, T, ctc_dim)
+            results["ctc"] = ctc_pred
+
+        return results
+
+    def compute_loss(self, x, x_mask, cond=None):
+        # x0: (B, T, num_quantizer)
+        # x_mask: (B, T) mask is 0 for padding
+        t = torch.rand(x.shape[0], device=x.device, requires_grad=False)
+        t = torch.clamp(t, 1e-5, 1.0)
+        # from CosyVoice: considering the generation process at the beginning is harder than follows, we involve a cosine scheduler for the timestep t
+        if self.time_scheduler == "cos":
+            t = 1 - torch.cos(t * math.pi * 0.5)
+        else:
+            pass
+        return self.loss_t(x, x_mask, t, cond)
+
+    def reset_parameters(self):
+        def _reset_parameters(m):
+            if isinstance(m, nn.MultiheadAttention):
+                if m._qkv_same_embed_dim:
+                    nn.init.normal_(m.in_proj_weight, std=0.02)
+                else:
+                    nn.init.normal_(m.q_proj_weight, std=0.02)
+                    nn.init.normal_(m.k_proj_weight, std=0.02)
+                    nn.init.normal_(m.v_proj_weight, std=0.02)
+
+                if m.in_proj_bias is not None:
+                    nn.init.constant_(m.in_proj_bias, 0.0)
+                    nn.init.constant_(m.out_proj.bias, 0.0)
+                if m.bias_k is not None:
+                    nn.init.xavier_normal_(m.bias_k)
+                if m.bias_v is not None:
+                    nn.init.xavier_normal_(m.bias_v)
+
+            elif (
+                isinstance(m, nn.Conv1d)
+                or isinstance(m, nn.ConvTranspose1d)
+                or isinstance(m, nn.Conv2d)
+                or isinstance(m, nn.ConvTranspose2d)
+            ):
+                m.weight.data.normal_(0.0, 0.02)
+
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(mean=0.0, std=0.02)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+
+            elif isinstance(m, nn.Embedding):
+                m.weight.data.normal_(mean=0.0, std=0.02)
+                if m.padding_idx is not None:
+                    m.weight.data[m.padding_idx].zero_()
+
+        self.apply(_reset_parameters)
+
+    @torch.no_grad()
+    def reverse_diffusion(
+        self,
+        cond,
+        prompt,
+        x_mask=None,
+        prompt_mask=None,
+        text_embedding=None,
+        n_timesteps=10,
+        cfg=1.0,
+        rescale_cfg=0.75,
+    ):
+        h = 1.0 / n_timesteps
+        prompt_len = prompt.shape[1]
+        target_len = cond.shape[1] - prompt_len
+
+        if x_mask == None:
+            x_mask = torch.ones(cond.shape[0], target_len).to(cond.device)  # (B, T)
+        if prompt_mask == None:
+            prompt_mask = torch.ones(cond.shape[0], prompt_len).to(
+                cond.device
+            )  # (B, prompt_len)
+        xt_mask = torch.cat([prompt_mask, x_mask], dim=1)
+        z = torch.randn(
+            (cond.shape[0], target_len, self.mel_dim),
+            dtype=cond.dtype,
+            device=cond.device,
+            requires_grad=False,
+        )
+        xt = z
+
+        # t from 0 to 1: x0 = z ~ N(0, 1)
+        for i in range(n_timesteps):
+            xt_input = torch.cat([prompt, xt], dim=1)
+            t = (0 + (i + 0.5) * h) * torch.ones(
+                z.shape[0], dtype=z.dtype, device=z.device
+            )
+            flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask)
+            flow_pred = flow_pred[:, prompt_len:, :]
+
+            # cfg
+            if cfg > 0:
+                uncond_flow_pred = self.diff_estimator(
+                    xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask
+                )
+                pos_flow_pred_std = flow_pred.std()
+                flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred)
+                rescale_flow_pred = (
+                    flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std()
+                )
+                flow_pred = (
+                    rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg
+                )
+
+            dxt = flow_pred * h
+            xt = xt + dxt
+
+        return xt
+
+    def forward(self, x, x_mask, cond_code):
+        """
+        Args:
+            x: (B, T, mel_dim)
+            x_mask: (B, T)
+            cond_code: (B, T), Note that cond_code might be not at 50Hz!
+        """
+        T = x.shape[1]
+
+        cond = self.cond_emb(cond_code)  # (B, T, hidden_size)
+        if self.do_resampling:
+            # Align to the frame rate of Mels
+            cond = self.resampling_layers(cond.transpose(1, 2)).transpose(1, 2)
+
+        # print("cond_code: {}, after resampling: {}".format(cond_code.shape, cond.shape))
+
+        if cond.shape[1] >= T:  # Check time dimension
+            cond = cond[:, :T, :]
+        else:
+            padding_frames = T - cond.shape[1]
+            last_frame = cond[:, -1:, :]
+            padding = last_frame.repeat(1, padding_frames, 1)
+            cond = torch.cat([cond, padding], dim=1)
+
+        return self.compute_loss(x, x_mask, cond)

models/svc/vevo2/qwen_utils.py

@@ -0,0 +1,105 @@
+import re
+import torch
+
+
+def format_chat_prompt(messages, add_assistant_token):
+    """
+    Convert the messages list into the Qwen chat template format.
+
+    Args:
+        messages: A list of messages containing role and content.
+        add_assistant_token: Whether to add assistant token at the end.
+
+    Returns:
+        str: The formatted prompt string.
+    """
+    prompt = ""
+    for msg in messages:
+        role = msg["role"]
+        content = msg["content"]
+        # Add start and end tags for all messages except the last assistant message
+        if msg != messages[-1] or role != "assistant":
+            prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
+        else:
+            # For the last assistant message, only add the start tag and content
+            prompt += f"<|im_start|>{role}\n{content}"
+
+    # If the last message is not from assistant and add_assistant_token is True
+    if messages[-1]["role"] != "assistant" and add_assistant_token:
+        prompt += f"<|im_start|>assistant\n"
+
+    return prompt
+
+
+def gen_chat_prompt(text, add_assistant_token, follow_prosody_instruction):
+    """
+    Args:
+        text (str): The text to be spoken.
+        add_assistant_token (bool): Whether to add assistant token at the end. For pre-training, False. For Inference, True.
+        follow_prosody_instruction (bool): Whether to follow the prosody instruction. When prosody_ids is not None, True. Otherwise, False.
+    """
+    if follow_prosody_instruction:
+        synthesis_instruction = "User will provide you with a text. Please first generate a good prosodic instruction, then vocalize the text based on it."
+    else:
+        synthesis_instruction = "User will provide you with a text. Please vocalize it with natural expression."
+
+    template = [
+        {
+            "role": "system",
+            "content": synthesis_instruction,
+        },
+        {
+            "role": "user",
+            "content": text,
+        },
+    ]
+    return format_chat_prompt(template, add_assistant_token)
+
+
+def gen_chat_response(prosody_ids, content_style_ids, is_full_response=True):
+    """
+    Args:
+        prosody_ids (list): The prosody ids of the text.
+        content_style_ids (list): The content style ids of the text.
+    """
+    if prosody_ids is not None:
+        prosody_text = "".join(["<|prosody_{}|>".format(int(i)) for i in prosody_ids])
+        prosody_text = "<|prosody_start|>" + prosody_text + "<|prosody_end|>"
+    else:
+        prosody_text = ""
+
+    if content_style_ids is not None:
+        content_style_text = "".join(
+            ["<|content_style_{}|>".format(int(i)) for i in content_style_ids]
+        )
+    else:
+        content_style_text = ""
+
+    if is_full_response:
+        return (
+            prosody_text
+            + "<|content_style_start|>"
+            + content_style_text
+            + "<|content_style_end|>"
+            + "<|im_end|>"
+        )
+    else:
+        return prosody_text + "<|content_style_start|>" + content_style_text
+
+
+def extract_content_style_ids(text):
+    """
+    Extract the content_style IDs from the text
+
+    Args:
+        text (str): A string containing content_style tags
+
+    Returns:
+        torch.Tensor: [T]
+    """
+    # Use regex to match all <|content_style_数字|> patterns
+    pattern = r"<\|content_style_(\d+)\|>"
+    # Find all matches and extract the numeric part
+    matches = re.findall(pattern, text)
+    # Convert string numbers to integers
+    return torch.tensor([int(match) for match in matches], dtype=torch.long)

models/svc/vevo2/vevo2_utils.py

@@ -0,0 +1,961 @@
+import math
+import json
+import librosa
+import soundfile as sf
+import torch
+import torchaudio
+import accelerate
+import safetensors
+import numpy as np
+import os
+import yaml
+from IPython.display import display, Audio
+
+import torchvision
+import random
+import numpy as np
+import whisper
+from librosa.feature import chroma_stft
+from librosa.effects import pitch_shift
+
+from models.codec.coco.rep_coco_model import CocoContentStyle, CocoContent, CocoStyle
+from models.svc.flow_matching_transformer.fmt_model import FlowMatchingTransformer
+from models.codec.melvqgan.melspec import MelSpectrogram
+from models.codec.amphion_codec.vocos import Vocos
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from utils.util import load_config
+from models.svc.vevo2.qwen_utils import gen_chat_prompt
+from evaluation.metrics.f0.f0_corr import extract_f0_hz
+
+from transformers.utils import is_flash_attn_2_available
+
+supported_flash_attn = False
+if not torch.cuda.is_available():
+    print("No CUDA available")
+    supported_flash_attn = False
+
+# To check if flash attention is supported
+if is_flash_attn_2_available():
+    supported_flash_attn = True
+    print("Flash Attention is supported")
+else:
+    print("Flash Attention is not supported")
+
+
+# Coco Tokenizer
+def build_coco_model(coco_cfg, device, loading_decoder=False):
+    coco_model_type = getattr(coco_cfg, "coco_type", "content_style")
+    if coco_model_type == "content_style":
+        model = CocoContentStyle(
+            cfg=coco_cfg, construct_only_for_quantizer=not loading_decoder
+        )
+    elif coco_model_type == "content":
+        model = CocoContent(
+            cfg=coco_cfg, construct_only_for_quantizer=not loading_decoder
+        )
+    elif coco_model_type == "style":
+        model = CocoStyle(
+            cfg=coco_cfg, construct_only_for_quantizer=not loading_decoder
+        )
+    else:
+        raise ValueError(f"Unknown coco type: {coco_model_type}")
+
+    model.eval()
+    model.to(device)
+    return model
+
+
+# Flow Matching Transformer
+def build_fmt_model(cfg, device):
+    model = FlowMatchingTransformer(cfg=cfg.model.flow_matching_transformer)
+    model.eval()
+    model.to(device)
+    return model
+
+
+# Autoregressive Transformer
+def build_ar_model(ckpt_path, device):
+    model_kwargs = {
+        "device_map": device,
+        "torch_dtype": "auto",
+        "trust_remote_code": True,
+    }
+
+    # Only add flash attention parameter if supported
+    if supported_flash_attn:
+        model_kwargs["attn_implementation"] = "flash_attention_2"
+
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     cfg.model.pretrained_model_path, **model_kwargs
+    # )
+    model = AutoModelForCausalLM.from_pretrained(ckpt_path, **model_kwargs)
+
+    model.eval()
+    return model
+
+
+# Melspectrogram Extractor
+def build_mel_model(cfg, device):
+    mel_model = MelSpectrogram(
+        sampling_rate=cfg.preprocess.sample_rate,
+        n_fft=cfg.preprocess.n_fft,
+        num_mels=cfg.preprocess.num_mels,
+        hop_size=cfg.preprocess.hop_size,
+        win_size=cfg.preprocess.win_size,
+        fmin=cfg.preprocess.fmin,
+        fmax=cfg.preprocess.fmax,
+    )
+    mel_model.eval()
+    mel_model.to(device)
+    return mel_model
+
+
+# Vocoder
+def build_vocoder_model(cfg, device):
+    vocoder_model = Vocos(cfg=cfg.model.vocos)
+    vocoder_model.eval()
+    vocoder_model.to(device)
+    return vocoder_model
+
+
+def load_checkpoint(build_model_func, cfg, ckpt_path, device):
+    model = build_model_func(cfg, device)
+    accelerate.load_checkpoint_and_dispatch(model, ckpt_path)
+    return model
+
+
+def count_parameters(model):
+    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    if total_params < 1e6:
+        return f"{total_params} params"  # Parameters
+    elif total_params < 1e9:
+        return f"{total_params / 1e6:.2f} M"  # Millions
+    else:
+        return f"{total_params / 1e9:.2f} B"  # Billions
+
+
+def load_wav(wav_path, device, used_duration=None):
+    if wav_path is None:
+        speech = np.zeros((0))  # [T]
+        speech_tensor = torch.zeros(1, 0).to(device)  # [1, T]
+        speech16k = torch.zeros(1, 0).to(device)  # [1, T']
+    else:
+        speech = librosa.load(wav_path, sr=24000)[0]  # [T]
+
+        if used_duration is not None:
+            speech = speech[: int(used_duration * 24000)]
+
+        speech_tensor = torch.tensor(speech).unsqueeze(0).to(device)  # [1, T]
+        speech16k = torchaudio.functional.resample(
+            speech_tensor, 24000, 16000
+        )  # [1, T']
+
+    return speech, speech_tensor, speech16k
+
+
+def display_audio_in_notebook(wav, rate=24000):
+    display(Audio(wav, rate=rate))
+
+
+def save_audio(
+    waveform, sr=24000, output_path=None, target_sample_rate=None, target_db=-25.0
+):
+    """
+    waveform: [1, T]
+    """
+    if target_sample_rate is not None and sr != target_sample_rate:
+        resampler = torchaudio.transforms.Resample(
+            orig_freq=sr, new_freq=target_sample_rate
+        )
+        waveform = resampler(waveform)
+    else:
+        target_sample_rate = sr
+
+    rms = torch.sqrt(torch.mean(waveform**2))
+    current_db = 20 * torch.log10(rms + 1e-9)
+
+    gain = target_db - current_db
+    normalized_waveform = waveform * (10 ** (gain / 20))
+
+    # deviation: torchaudio.save() known to cause dep issues; using soundfile instead
+    sf.write(
+        output_path,
+        normalized_waveform.squeeze(0).cpu().numpy(),
+        target_sample_rate,
+    )
+    return output_path
+
+
+def extract_special_ids(text, prefix="content_style"):
+    """
+    Extract all audio IDs from a string containing <|content_style_X|> or <|prosody_X|> tags
+
+    Args:
+        text (str): A string containing audio tags
+        prefix (str): The prefix of the audio tags
+
+    Returns:
+        list: A list of all audio IDs
+    """
+    import re
+
+    if prefix == "content_style":
+        # Use regex to match all <|content_style_X|> format tags
+        pattern = r"<\|content_style_(\d+)\|>"
+    elif prefix == "prosody":
+        pattern = r"<\|prosody_(\d+)\|>"
+    else:
+        raise ValueError(f"Unknown prefix: {prefix}")
+
+    # Find all matches and extract the numeric part
+    audio_ids = re.findall(pattern, text)
+
+    # Convert string IDs to integers
+    audio_ids = [int(id) for id in audio_ids]
+
+    return audio_ids
+
+
+class Vevo2InferencePipeline:
+    def __init__(
+        self,
+        prosody_tokenizer_ckpt_path=None,
+        content_style_tokenizer_ckpt_path=None,
+        ar_cfg_path=None,
+        ar_ckpt_path=None,
+        fmt_cfg_path=None,
+        fmt_ckpt_path=None,
+        vocoder_cfg_path=None,
+        vocoder_ckpt_path=None,
+        device=None,
+        use_vllm=False,
+    ):
+        self.device = device
+        self.use_vllm = use_vllm
+
+        self.prosody_tokenizer_ckpt_path = prosody_tokenizer_ckpt_path
+        self.content_style_tokenizer_ckpt_path = content_style_tokenizer_ckpt_path
+
+        if ar_cfg_path is not None and ar_ckpt_path is not None:
+            self.ar_cfg = load_config(ar_cfg_path)
+
+            assert not use_vllm, "VLLM is not supported yet"
+
+            if use_vllm:
+                pass
+            else:
+                # self.ar_model = load_checkpoint(
+                #     build_ar_model, self.ar_cfg, ar_ckpt_path, device
+                # )
+
+                self.ar_model = build_ar_model(ar_ckpt_path, device)
+                print(f"#Params of AR model: {count_parameters(self.ar_model)}")
+
+            # self.ar_tokenizer = AutoTokenizer.from_pretrained(
+            #     self.ar_cfg.preprocess.tokenizer_path, local_files_only=True
+            # )
+            self.ar_tokenizer = AutoTokenizer.from_pretrained(
+                ar_ckpt_path, local_files_only=True
+            )
+        else:
+            self.ar_cfg = None
+            self.ar_model = None
+
+        if fmt_cfg_path is not None and fmt_ckpt_path is not None:
+            self.fmt_cfg = load_config(fmt_cfg_path)
+            self.fmt_model = load_checkpoint(
+                build_fmt_model, self.fmt_cfg, fmt_ckpt_path, device
+            )
+            print(f"#Params of Flow Matching model: {count_parameters(self.fmt_model)}")
+
+            if getattr(
+                self.fmt_cfg.model.flow_matching_transformer,
+                "use_text_as_condition",
+                False,
+            ):
+                self.fmt_use_text_as_condition = True
+                self.fmt_text_tokenizer = AutoTokenizer.from_pretrained(
+                    self.fmt_cfg.preprocess.tokenizer_path
+                )
+            else:
+                self.fmt_use_text_as_condition = False
+
+            self.init_coco_tokenizer()
+
+        if vocoder_cfg_path is not None and vocoder_ckpt_path is not None:
+            self.vocoder_cfg = load_config(vocoder_cfg_path)
+            self.mel_model = build_mel_model(self.vocoder_cfg, device)
+            self.vocoder_model = load_checkpoint(
+                build_vocoder_model, self.vocoder_cfg, vocoder_ckpt_path, device
+            )
+            print(f"#Params of Vocoder model: {count_parameters(self.vocoder_model)}")
+
+    def init_coco_tokenizer(self):
+        ## Whisper ##
+        self.whisper_model = whisper.load_model("medium", self.device)  # 1024 dim
+        self.whisper_model.eval()
+
+        self.use_normed_whisper = getattr(
+            self.fmt_cfg.model.coco, "use_normed_whisper", False
+        )
+        if self.use_normed_whisper:
+            whisper_stats = torch.load(
+                self.fmt_cfg.model.coco.whisper_stats_path,
+                map_location=self.device,
+            )
+            self.whisper_mean = whisper_stats["mean"]  # (1024,)
+            self.whisper_std = whisper_stats["std"]  # (1024,)
+
+        ## Prosody Tokenizer ##
+        if self.ar_model is not None:
+            self.style_tokenizer = load_checkpoint(
+                build_coco_model,
+                self.ar_cfg.model.coco_style,
+                self.prosody_tokenizer_ckpt_path,
+                self.device,
+            )
+            print(
+                f"#Params of CocoStyle model: {count_parameters(self.style_tokenizer)}"
+            )
+
+        ## Content-Style Tokenizer ##
+        self.content_style_tokenizer = load_checkpoint(
+            build_coco_model,
+            self.fmt_cfg.model.coco,
+            self.content_style_tokenizer_ckpt_path,
+            self.device,
+        )
+        print(
+            f"#Params of CocoContentStyle model: {count_parameters(self.content_style_tokenizer)}"
+        )
+
+    @torch.no_grad()
+    def extract_mel_feature(self, speech):
+        mel_feature = self.mel_model(speech)  # (B, d, T)
+        mel_feature = mel_feature.transpose(1, 2)
+        mel_feature = (mel_feature - self.vocoder_cfg.preprocess.mel_mean) / math.sqrt(
+            self.vocoder_cfg.preprocess.mel_var
+        )
+        return mel_feature
+
+    def spec_augment(self, mel, height):
+        """
+        Args:
+            mel: tensor (..., n_mels, frames)
+            height: int 68-92 for default 80 mels
+        """
+        tgt = torchvision.transforms.functional.resize(mel, (height, mel.shape[-1]))
+        if height >= mel.shape[-2]:
+            return tgt[:, : mel.shape[-2], :]
+        else:
+            silence = tgt[:, -1:, :].repeat(1, mel.shape[-2] - height, 1)
+            silence += torch.randn_like(silence) / 10
+            return torch.cat((tgt, silence), 1)
+
+    @torch.no_grad()
+    def extract_whisper_features(self, wavs, frame_lens, spec_perturb=False):
+        """
+        Args:
+            wavs: (B, T) at 16khz. Note that the max duration should be 30s
+            frame_lens: (B,)
+        Returns:
+            features: (B, T, D)
+        """
+        # wavs: (batch, max_len)
+        wavs = whisper.pad_or_trim(wavs)
+        # batch_mel: (batch, 80, 3000)
+        batch_mel = whisper.log_mel_spectrogram(wavs, device=self.device)
+
+        if spec_perturb:
+            height = random.randint(68, 92)
+            batch_mel = self.spec_augment(batch_mel, height)
+
+        with torch.no_grad():
+            # (batch, 1500, 1024)
+            features = self.whisper_model.embed_audio(batch_mel)
+
+        max_len = int(frame_lens.max().item())
+        mask = torch.arange(features.size(1), device=features.device).expand(
+            len(frame_lens), -1
+        ) < frame_lens.unsqueeze(1)
+        features = torch.where(mask.unsqueeze(-1), features, torch.zeros_like(features))
+
+        if features.shape[1] >= max_len:
+            features = features[:, :max_len, :]
+        else:
+            padding_frames = max_len - features.shape[1]
+            last_frame = features[:, -1:, :]
+            padding = last_frame.repeat(1, padding_frames, 1)
+            features = torch.cat([features, padding], dim=1)
+
+        if self.use_normed_whisper:
+            features = (features - self.whisper_mean) / self.whisper_std
+
+        return features
+
+    @torch.no_grad()
+    def extract_coco_codec(
+        self,
+        coco_codec_type,
+        wav16k,
+        wav24k_numpy,
+        whisper_spec_perturb=False,
+        frame_len_ratio=1.0,
+        use_shifted_wav_to_extract_chromagram=False,
+        use_shifted_wav_to_extract_whisper=False,
+        pitch_shift_steps=0,
+    ):
+        """
+        Args:
+            coco_codec_type: "content", "style", or "content_style"
+            wav16k: [1, T]
+            wav24k_numpy: [T]
+        Returns:
+            codecs: [1, T]. Note that codecs might be not at 50Hz!
+        """
+        frame_len = len(wav24k_numpy) // self.fmt_cfg.preprocess.hop_size
+
+        if use_shifted_wav_to_extract_chromagram:
+            chromagram_feats = self.get_chromagram(
+                pitch_shift(wav24k_numpy, sr=24000, n_steps=pitch_shift_steps),
+                frame_len,
+            )  # [T, 24]
+        else:
+            chromagram_feats = self.get_chromagram(wav24k_numpy, frame_len)  # [T, 24]
+
+        chromagram_feats = (
+            torch.tensor(chromagram_feats, dtype=torch.float)
+            .unsqueeze(0)
+            .to(self.device)
+        )  # [1, T, 24]
+
+        if frame_len_ratio != 1.0:
+            raw_len = chromagram_feats.shape[1]
+            # Convert [1, T, 24] to [1, 24, T] for interpolation on the last dimension
+            chromagram_feats = chromagram_feats.transpose(1, 2)
+            chromagram_feats = torch.nn.functional.interpolate(
+                chromagram_feats,
+                size=int(
+                    raw_len * frame_len_ratio
+                ),  # Explicitly specify the target length
+                mode="linear",
+                align_corners=False,
+            )  # [1, 24, T']
+            # Convert back to the original shape [1, T', 24]
+            chromagram_feats = chromagram_feats.transpose(1, 2)
+            print(
+                f"Chromagram feats are sampled from {raw_len} to {chromagram_feats.shape[1]}, ratio = {frame_len_ratio}"
+            )
+
+        if use_shifted_wav_to_extract_whisper:
+            wav16k = pitch_shift(
+                wav16k.cpu().numpy()[0], sr=16000, n_steps=pitch_shift_steps
+            )  # [T]
+            wav16k = torch.tensor(wav16k).unsqueeze(0).to(self.device)  # [1, T]
+
+        if coco_codec_type in ["content_style", "content"]:
+            whisper_feats = self.extract_whisper_features(
+                wav16k,
+                torch.tensor([frame_len], dtype=torch.long).to(self.device),
+                spec_perturb=whisper_spec_perturb,
+            )  # [1, T, D]
+
+        if coco_codec_type == "content_style":
+            codecs, _ = self.content_style_tokenizer.quantize(
+                whisper_feats.to(torch.float32), chromagram_feats.to(torch.float32)
+            )
+        elif coco_codec_type == "style":
+            codecs, _ = self.style_tokenizer.quantize(
+                chromagram_feats.to(torch.float32)
+            )
+        else:
+            raise ValueError(f"Unknown coco type: {coco_codec_type}")
+
+        return codecs
+
+    def get_chromagram(self, speech, speech_frames):
+        # [24, T] -> [T, 24]
+        chromagram = chroma_stft(
+            y=speech,
+            sr=self.fmt_cfg.preprocess.sample_rate,
+            n_fft=self.fmt_cfg.preprocess.n_fft,
+            hop_length=self.fmt_cfg.preprocess.hop_size,
+            win_length=self.fmt_cfg.preprocess.win_size,
+            n_chroma=24,
+        ).T
+
+        if chromagram.shape[0] < speech_frames:
+            chromagram = np.pad(
+                chromagram, (0, speech_frames - chromagram.shape[0]), mode="edge"
+            )
+        else:
+            chromagram = chromagram[:speech_frames]
+
+        return chromagram
+
+    def get_shifted_steps(self, src_wav_path, timbre_ref_wav_path):
+        if src_wav_path == timbre_ref_wav_path:
+            return 0
+
+        src_f0 = extract_f0_hz(src_wav_path)
+        timbre_ref_f0 = extract_f0_hz(timbre_ref_wav_path)
+
+        src_f0_median = np.median(src_f0)
+        timbre_ref_f0_median = np.median(timbre_ref_f0)
+
+        src_shifted_steps = 12 * np.log2(timbre_ref_f0_median / src_f0_median)
+        src_shifted_steps = round(src_shifted_steps)
+
+        if src_shifted_steps > 12:
+            src_shifted_steps = src_shifted_steps % 12
+        elif src_shifted_steps < -12:
+            src_shifted_steps = src_shifted_steps % -12
+
+        return src_shifted_steps
+
+    @torch.no_grad()
+    def inference_fm(
+        self,
+        src_wav_path,
+        timbre_ref_wav_path,
+        src_wav_text="",
+        timbre_ref_wav_text="",
+        whisper_spec_perturb=False,
+        use_pitch_shift=False,
+        used_duration_of_timbre_ref_wav_path=None,
+        flow_matching_steps=32,
+        display_audio=False,
+    ):
+        src_speech, src_speech24k, src_speech16k = load_wav(src_wav_path, self.device)
+
+        if display_audio:
+            print("-" * 20)
+            if src_wav_path == timbre_ref_wav_path:
+                print("We want to reconstruct this audio:", src_wav_path)
+                display_audio_in_notebook(src_wav_path, rate=24000)
+            else:
+                print("Source audio:")
+                display_audio_in_notebook(src_speech, rate=24000)
+
+        ## Whether to use shifted src to extract prosody and content-style ##
+        if use_pitch_shift:
+            src_shifted_steps = self.get_shifted_steps(
+                src_wav_path, timbre_ref_wav_path
+            )
+
+            if display_audio:
+                print("-" * 20)
+                print(f"src_shifted_steps: {src_shifted_steps}")
+        else:
+            src_shifted_steps = 0
+
+        ## Diffusion ##
+        src_codecs = self.extract_coco_codec(
+            "content_style",
+            src_speech16k,
+            src_speech,
+            whisper_spec_perturb=whisper_spec_perturb,
+            use_shifted_wav_to_extract_chromagram=use_pitch_shift,
+            pitch_shift_steps=src_shifted_steps,
+        )  # [1, T]
+
+        predict_mel_feat = self.code2mel(
+            src_codecs,
+            timbre_ref_wav_path,
+            prefix_text=timbre_ref_wav_text + " " + src_wav_text,
+            used_duration_of_timbre_ref_wav_path=used_duration_of_timbre_ref_wav_path,
+            flow_matching_steps=flow_matching_steps,
+            logging=display_audio,
+        )  # [1, T, D]
+
+        ## Vocoder and Display ##
+        synthesized_audio = self.mel2audio(
+            predict_mel_feat, logging=display_audio
+        )  # [1, T]
+
+        return synthesized_audio
+
+    def get_llm_prompt_text(self, text, follow_prosody_instruction, logging=False):
+        llm_prompt_text = gen_chat_prompt(
+            text,
+            add_assistant_token=True,
+            follow_prosody_instruction=follow_prosody_instruction,
+        )
+
+        # if logging:
+        #     print("-" * 20)
+        #     print("LLM Prompt Text:\n{}".format(llm_prompt_text))
+
+        return llm_prompt_text
+
+    def get_llm_prompt_prosody(
+        self,
+        use_prosody_code,
+        predict_target_prosody,
+        prosody_wav_path=None,
+        style_ref_wav_path=None,
+        use_pitch_shift=False,
+        prosody_wav_pitch_shift_steps=0,
+        style_ref_wav_pitch_shift_steps=0,
+        target_duration=None,
+        logging=False,
+    ):
+        if not use_prosody_code:
+            return ""
+
+        if not predict_target_prosody:
+            # Just use the ground truth prosody #
+
+            assert prosody_wav_path is not None
+            prosody_speech, prosody_speech24k, prosody_speech16k = load_wav(
+                prosody_wav_path, self.device
+            )
+
+            if target_duration is not None:
+                # Calculate the chromagram frame len ratio
+                prosody_wav_duration = prosody_speech.shape[0] / 24000
+                prosody_wav_chromagram_frame_len_ratio = (
+                    target_duration / prosody_wav_duration
+                )
+            else:
+                prosody_wav_chromagram_frame_len_ratio = 1.0
+
+            prosody_wav_prosody_ids = self.extract_coco_codec(
+                "style",
+                prosody_speech16k,
+                prosody_speech,
+                frame_len_ratio=prosody_wav_chromagram_frame_len_ratio,
+                use_shifted_wav_to_extract_chromagram=use_pitch_shift,
+                pitch_shift_steps=prosody_wav_pitch_shift_steps,
+            )  # [1, T]
+
+            if style_ref_wav_path is not None:
+                style_ref_speech, style_ref_speech24k, style_ref_speech16k = load_wav(
+                    style_ref_wav_path, self.device
+                )
+                style_ref_wav_prosody_ids = self.extract_coco_codec(
+                    "style",
+                    style_ref_speech16k,
+                    style_ref_speech,
+                    use_shifted_wav_to_extract_chromagram=use_pitch_shift,
+                    pitch_shift_steps=style_ref_wav_pitch_shift_steps,
+                )  # [1, T]
+            else:
+                style_ref_wav_prosody_ids = torch.zeros(1, 0).to(self.device)
+
+            prosody_ids = torch.cat(
+                [style_ref_wav_prosody_ids, prosody_wav_prosody_ids], dim=1
+            )  # [1, T]
+
+            prosody_ids = prosody_ids[0].tolist()
+            prosody_text = "".join(["<|prosody_{}|>".format(i) for i in prosody_ids])
+            prosody_text = "<|prosody_start|>" + prosody_text + "<|prosody_end|>"
+
+            if logging:
+                print("-" * 20)
+                print("Prosody (Melody) Audio: ", prosody_wav_path)
+                display_audio_in_notebook(prosody_speech, rate=24000)
+
+        else:
+            raise NotImplementedError("Not implemented yet")
+
+        return prosody_text
+
+    def get_llm_prompt_contentstyle(
+        self,
+        style_ref_wav_path=None,
+        use_pitch_shift=False,
+        pitch_shift_steps=0,
+        logging=False,
+    ):
+        if style_ref_wav_path is None:
+            return "<|content_style_start|>"
+
+        style_ref_speech, style_ref_speech24k, style_ref_speech16k = load_wav(
+            style_ref_wav_path, self.device
+        )
+        if logging:
+            print("-" * 20)
+            print("Style Reference Audio: ", style_ref_wav_path)
+            display_audio_in_notebook(style_ref_speech, rate=24000)
+
+        prompt_output_ids = self.extract_coco_codec(
+            "content_style",
+            style_ref_speech16k,
+            style_ref_speech,
+            use_shifted_wav_to_extract_chromagram=use_pitch_shift,
+            pitch_shift_steps=pitch_shift_steps,
+        )  # [1, T]
+
+        prompt_output_ids = prompt_output_ids[0].tolist()
+        prompt_output_text = "".join(
+            ["<|content_style_{}|>".format(i) for i in prompt_output_ids]
+        )
+        prompt_output_text = "<|content_style_start|>" + prompt_output_text
+        return prompt_output_text
+
+    def parse_llm_generated_ids(self, generated_ids, llm_input_ids, logging=False):
+        """
+        Args:
+            generated_ids: [1, T]
+            llm_input_ids: [1, T]
+        Returns:
+            coco_codecs: [1, T]
+        """
+        input_len = llm_input_ids.shape[1]
+        generated_ids = generated_ids[:, input_len:]
+
+        # Eg: <|content_style_start|> <|content_style_1|> <|content_style_2|> <|content_style_end|> <|im_end|>
+        generated_text = self.ar_tokenizer.decode(
+            generated_ids[0], skip_special_tokens=False
+        )
+
+        content_style_ids = extract_special_ids(generated_text, prefix="content_style")
+        content_style_ids = (
+            torch.tensor(content_style_ids, dtype=torch.long)
+            .to(self.device)
+            .unsqueeze(0)
+        )  # [1, T]
+
+        if logging:
+            print("-" * 20)
+            print("LLM input_ids: ", llm_input_ids.shape)
+            print("Generated content-style ids: ", content_style_ids.shape)
+
+        return content_style_ids
+
+    @torch.no_grad()
+    def code2mel(
+        self,
+        contentstyle_codecs,
+        timbre_ref_wav_path,
+        prefix_text="",
+        used_duration_of_timbre_ref_wav_path=None,
+        flow_matching_steps=32,
+        logging=False,
+    ):
+        timbre_ref_speech, timbre_ref_speech24k, timbre_ref_speech16k = load_wav(
+            timbre_ref_wav_path,
+            self.device,
+            used_duration=used_duration_of_timbre_ref_wav_path,
+        )
+        if logging:
+            print("-" * 20)
+            print("Timbre Reference Audio: ", timbre_ref_wav_path)
+            display_audio_in_notebook(timbre_ref_speech, rate=24000)
+
+        timbre_ref_codecs = self.extract_coco_codec(
+            "content_style",
+            timbre_ref_speech16k,
+            timbre_ref_speech,
+        )  # [1, T]
+
+        diffusion_input_codecs = torch.cat(
+            [timbre_ref_codecs, contentstyle_codecs], dim=1
+        )
+
+        # Prepare the condition for diffusion
+        diffusion_cond = self.fmt_model.cond_emb(diffusion_input_codecs)  # [1, T, D]
+        if self.fmt_model.do_resampling:
+            # Align to the frame rate of Mels
+            diffusion_cond = self.fmt_model.resampling_layers(
+                diffusion_cond.transpose(1, 2)
+            ).transpose(1, 2)
+
+        timbre_ref_mels = self.extract_mel_feature(timbre_ref_speech24k)  # [1, T, D]
+
+        # Text as condition
+        if self.fmt_use_text_as_condition:
+            prefix_text_ids = self.fmt_text_tokenizer.encode(
+                prefix_text, add_special_tokens=False
+            )
+            prefix_text_ids = torch.tensor(prefix_text_ids, dtype=torch.long).to(
+                self.device
+            )  # [T]
+            prefix_text_ids = prefix_text_ids.unsqueeze(0)  # [1, T]
+            prefix_text_embedding = self.fmt_model.text_cond_emb(
+                prefix_text_ids
+            )  # [1, T, D]
+        else:
+            prefix_text_embedding = None
+
+        # [1, T, D]
+        predict_mel_feat = self.fmt_model.reverse_diffusion(
+            cond=diffusion_cond,
+            prompt=timbre_ref_mels,
+            text_embedding=prefix_text_embedding,
+            n_timesteps=flow_matching_steps,
+        )
+        return predict_mel_feat
+
+    @torch.no_grad()
+    def mel2audio(self, predict_mel_feat, logging=False):
+        # [1, 1, T] -> [1, T]
+        synthesized_audio = (
+            self.vocoder_model(predict_mel_feat.transpose(1, 2)).detach().cpu()
+        )[0]
+
+        if logging:
+            print("-" * 20)
+            print("Synthesized Audio:")
+            # [T]
+            audio = synthesized_audio.numpy()[0]
+            display_audio_in_notebook(audio, rate=24000)
+
+        return synthesized_audio
+
+    @torch.no_grad()
+    def inference_ar_and_fm(
+        self,
+        target_text=None,
+        prosody_wav_path=None,
+        style_ref_wav_path=None,
+        style_ref_wav_text="",
+        timbre_ref_wav_path=None,
+        use_prosody_code=True,
+        predict_target_prosody=False,
+        top_k=25,
+        top_p=0.8,
+        temperature=1.0,
+        use_pitch_shift=False,
+        prosody_wav_shifted_steps=0,
+        style_ref_wav_shifted_steps=0,
+        target_duration=None,
+        used_duration_of_timbre_ref_wav_path=None,
+        flow_matching_steps=32,
+        display_audio=False,
+    ):
+        """
+        Based on the style reference wav to conduct the continuation generation:
+            [Style_reference_Text, Target_Text], [Style_reference_Prosody, Target_Prosody], [Style_reference_cscodes, Target_cscodes]
+        """
+        assert self.ar_model is not None
+        assert target_text is not None
+        # assert style_ref_wav_path is not None
+        assert timbre_ref_wav_path is not None
+
+        ## Text Tokens ##
+        if display_audio:
+            print("-" * 20)
+            print("Target text: \n", target_text)
+
+        input_text = style_ref_wav_text + " " + target_text
+        llm_prompt_text = self.get_llm_prompt_text(
+            input_text, use_prosody_code, logging=display_audio
+        )
+
+        ## Whether to use shifted chromagram for timbre_ref_wav ##
+        if use_pitch_shift:
+            if prosody_wav_path is not None and prosody_wav_shifted_steps == 0:
+                prosody_wav_shifted_steps = self.get_shifted_steps(
+                    prosody_wav_path, timbre_ref_wav_path
+                )
+            if style_ref_wav_path is not None and style_ref_wav_shifted_steps == 0:
+                style_ref_wav_shifted_steps = self.get_shifted_steps(
+                    style_ref_wav_path, timbre_ref_wav_path
+                )
+
+            if display_audio:
+                print("-" * 20)
+                print(f"prosody_wav_shifted_steps: {prosody_wav_shifted_steps}")
+                print(f"style_ref_wav_shifted_steps: {style_ref_wav_shifted_steps}")
+
+        ## Prosody Tokens ##
+        llm_prompt_prosody = self.get_llm_prompt_prosody(
+            use_prosody_code,
+            predict_target_prosody,
+            prosody_wav_path=prosody_wav_path,
+            style_ref_wav_path=style_ref_wav_path,
+            use_pitch_shift=use_pitch_shift,
+            prosody_wav_pitch_shift_steps=prosody_wav_shifted_steps,
+            style_ref_wav_pitch_shift_steps=style_ref_wav_shifted_steps,
+            target_duration=target_duration,
+            logging=display_audio,
+        )
+
+        ## Content-Style Tokens ##
+        llm_prompt_contentstyle = self.get_llm_prompt_contentstyle(
+            style_ref_wav_path,
+            use_pitch_shift=use_pitch_shift,
+            pitch_shift_steps=style_ref_wav_shifted_steps,
+            logging=display_audio,
+        )
+
+        ## AR ##
+        llm_prompt = llm_prompt_text + llm_prompt_prosody + llm_prompt_contentstyle
+        llm_input_ids = self.ar_tokenizer.encode(llm_prompt, add_special_tokens=True)
+        llm_input_ids = (
+            torch.tensor(llm_input_ids, dtype=torch.long).to(self.device).unsqueeze(0)
+        )  # [1, T]
+
+        if self.use_vllm:
+            sampling_params = SamplingParams(
+                max_tokens=500,
+                top_k=top_k,
+                top_p=top_p,
+                temperature=temperature,
+                stop_token_ids=[self.ar_tokenizer.eos_token_id],
+                skip_special_tokens=False,
+            )
+            outputs = self.ar_model.generate([llm_prompt], sampling_params)
+
+            assert len(outputs) == 1
+            output = outputs[0]
+            prompt = output.prompt
+            predicted_coco_text = output.outputs[0].text
+            # print("vllm predicted_coco_text: ", predicted_coco_text)
+
+            predicted_coco_codecs = extract_special_ids(
+                predicted_coco_text, prefix="content_style"
+            )
+            predicted_coco_codecs = (
+                torch.tensor(predicted_coco_codecs, dtype=torch.long)
+                .to(self.device)
+                .unsqueeze(0)
+            )  # [1, T]
+        else:
+            generate_ids = self.ar_model.generate(
+                input_ids=llm_input_ids,
+                min_new_tokens=15,
+                max_new_tokens=500,
+                eos_token_id=self.ar_tokenizer.eos_token_id,
+                do_sample=True,
+                top_k=top_k,
+                top_p=top_p,
+                temperature=temperature,
+            )  # [1, T]
+
+            predicted_coco_codecs = self.parse_llm_generated_ids(
+                generate_ids, llm_input_ids, logging=display_audio
+            )  # [1, T]
+
+        ## Diffusion ##
+        predict_mel_feat = self.code2mel(
+            predicted_coco_codecs,
+            timbre_ref_wav_path,
+            used_duration_of_timbre_ref_wav_path=used_duration_of_timbre_ref_wav_path,
+            flow_matching_steps=flow_matching_steps,
+            logging=display_audio,
+        )  # [1, T, D]
+
+        ## Vocoder ##
+        synthesized_audio = self.mel2audio(
+            predict_mel_feat, logging=display_audio
+        )  # [1, T]
+        return synthesized_audio
+
+    @torch.no_grad()
+    def inference_vocoder_resynthesis(self, wav_path, display_audio=False):
+        speech, speech24k, speech16k = load_wav(wav_path, self.device)
+        if display_audio:
+            print("Ground Truth audio:")
+            display_audio_in_notebook(speech, rate=24000)
+
+        mel = self.extract_mel_feature(speech24k)  # [1, T, D]
+        audio = self.vocoder_model(mel.transpose(1, 2)).detach().cpu()[0]
+        if display_audio:
+            print("Resynthesized audio:")
+            display_audio_in_notebook(audio, rate=24000)
+        return audio

models/vc/flow_matching_transformer/llama_nar.py

@@ -0,0 +1,593 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from transformers import LlamaConfig, LlamaModel
+import torch
+import torch.nn as nn
+from typing import List, Optional, Tuple, Union
+import math
+import torch.nn.functional as F
+
+from transformers.models.llama.modeling_llama import (
+    LlamaDecoderLayer,
+    Cache,
+    apply_rotary_pos_emb,
+    repeat_kv,
+    BaseModelOutputWithPast,
+    LlamaRotaryEmbedding,
+)
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+# sinusoidal positional encoding
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :] * 1.0
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class LlamaAdaptiveRMSNorm(nn.Module):
+    def __init__(self, hidden_size=1024, eps=1e-6, dim_cond=1024):
+        super().__init__()
+        self.to_weight = nn.Linear(dim_cond, hidden_size)
+        nn.init.zeros_(self.to_weight.weight)
+        nn.init.ones_(self.to_weight.bias)
+        self.variance_epsilon = eps
+        self._is_hf_initialized = True  # disable automatic init
+
+    def forward(self, hidden_states, cond_embedding):
+        input_dtype = hidden_states.dtype
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        weight = self.to_weight(cond_embedding)
+        if len(weight.shape) == 2:
+            weight = weight.unsqueeze(1)
+
+        return (weight * hidden_states).to(input_dtype)
+
+
+class OldLlamaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+                "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
+        )
+
+        # TODO (joao): remove in v4.46 (RoPE is computed in the model, not in the decoder layers)
+        self.rotary_emb = LlamaRotaryEmbedding(config=self.config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[
+            Tuple[torch.Tensor, torch.Tensor]
+        ] = None,  # will become mandatory in v4.46
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        if self.config.pretraining_tp > 1:
+            key_value_slicing = (
+                self.num_key_value_heads * self.head_dim
+            ) // self.config.pretraining_tp
+            query_slices = self.q_proj.weight.split(
+                (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
+            )
+            key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
+            value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
+
+            query_states = [
+                F.linear(hidden_states, query_slices[i])
+                for i in range(self.config.pretraining_tp)
+            ]
+            query_states = torch.cat(query_states, dim=-1)
+
+            key_states = [
+                F.linear(hidden_states, key_slices[i])
+                for i in range(self.config.pretraining_tp)
+            ]
+            key_states = torch.cat(key_states, dim=-1)
+
+            value_states = [
+                F.linear(hidden_states, value_slices[i])
+                for i in range(self.config.pretraining_tp)
+            ]
+            value_states = torch.cat(value_states, dim=-1)
+
+        else:
+            query_states = self.q_proj(hidden_states)
+            key_states = self.k_proj(hidden_states)
+            value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+
+        if position_embeddings is None:
+            logger.warning_once(
+                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
+                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
+                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
+                "removed and `position_embeddings` will be mandatory."
+            )
+            cos, sin = self.rotary_emb(value_states, position_ids)
+        else:
+            cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        attn_weights = torch.matmul(
+            query_states, key_states.transpose(2, 3)
+        ) / math.sqrt(self.head_dim)
+
+        if attention_mask is not None:  # no matter the length, we just slice it
+            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+            attn_weights = attn_weights + causal_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(
+            attn_weights, dim=-1, dtype=torch.float32
+        ).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(
+            attn_weights, p=self.attention_dropout, training=self.training
+        )
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+
+        if self.config.pretraining_tp > 1:
+            attn_output = attn_output.split(
+                self.hidden_size // self.config.pretraining_tp, dim=2
+            )
+            o_proj_slices = self.o_proj.weight.split(
+                self.hidden_size // self.config.pretraining_tp, dim=1
+            )
+            attn_output = sum(
+                [
+                    F.linear(attn_output[i], o_proj_slices[i])
+                    for i in range(self.config.pretraining_tp)
+                ]
+            )
+        else:
+            attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LlamaNARDecoderLayer(LlamaDecoderLayer):
+    def __init__(self, config: LlamaConfig, layer_idx: int):
+        """Override to adaptive layer norm"""
+        super().__init__(config, layer_idx)  # init attention, mlp, etc.
+        self.input_layernorm = LlamaAdaptiveRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, dim_cond=config.hidden_size
+        )
+        self.post_attention_layernorm = LlamaAdaptiveRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, dim_cond=config.hidden_size
+        )
+
+        # For transformers v4.46 (added by Xueyao)
+        self.self_attn = OldLlamaAttention(config=config, layer_idx=layer_idx)
+
+    # add `cond` in forward function
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cond_embedding: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(
+            hidden_states, cond_embedding=cond_embedding
+        )
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(
+            hidden_states, cond_embedding=cond_embedding
+        )
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+class DiffLlama(LlamaModel):
+    def __init__(
+        self,
+        mel_dim=100,
+        hidden_size=1024,
+        num_heads=16,
+        num_layers=16,
+        dropout=0.1,
+        ffn_dropout=0.1,
+        attention_dropout=0.0,
+        # deviation: newer transformers rejects positional args, so kwargs
+        config=LlamaConfig(
+            vocab_size=0,
+            hidden_size=256,
+            intermediate_size=1024,
+            num_hidden_layers=1,
+            num_attention_heads=1,
+        ),
+    ):
+        super().__init__(config)
+
+        self.layers = nn.ModuleList(
+            [
+                LlamaNARDecoderLayer(
+                    LlamaConfig(
+                        hidden_size=hidden_size,
+                        num_attention_heads=num_heads,
+                        max_position_embeddings=4096,
+                        intermediate_size=hidden_size * 4,
+                    ),
+                    layer_idx=i,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.norm = LlamaAdaptiveRMSNorm(hidden_size, dim_cond=hidden_size)
+
+        self.diff_step_embedding = SinusoidalPosEmb(hidden_size)
+        self.diff_step_mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(hidden_size * 4, hidden_size),
+        )
+
+        self.cond_mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(hidden_size * 4, hidden_size),
+        )
+
+        self.mel_mlp = nn.Sequential(
+            nn.Linear(mel_dim, hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(hidden_size * 4, hidden_size),
+        )
+
+        self.mel_out_mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(hidden_size * 4, mel_dim),
+        )
+
+        for layer in self.layers:
+            layer.input_layernorm = LlamaAdaptiveRMSNorm(
+                hidden_size, dim_cond=hidden_size
+            )
+            layer.post_attention_layernorm = LlamaAdaptiveRMSNorm(
+                hidden_size, dim_cond=hidden_size
+            )
+
+        self.embed_tokens = None
+
+        self.post_init()
+
+        # self.reset_parameters()
+
+    def _prepare_decoder_attention_mask(
+        self, attention_mask, input_shape, inputs_embeds, past_key_values_length
+    ):
+        # create noncausal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+
+        def _expand_mask(
+            mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None
+        ):
+            """
+            Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+            """
+            bsz, src_len = mask.size()
+            tgt_len = tgt_len if tgt_len is not None else src_len
+
+            expanded_mask = (
+                mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+            )
+
+            inverted_mask = 1.0 - expanded_mask
+
+            return inverted_mask.masked_fill(
+                inverted_mask.to(torch.bool), torch.finfo(dtype).min
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(
+                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+            ).to(inputs_embeds.device)
+            combined_attention_mask = (
+                expanded_attn_mask
+                if combined_attention_mask is None
+                else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        x,
+        diffusion_step,
+        cond,
+        x_mask,
+        input_ids: torch.LongTensor = None,  # [num_quant, B, T]
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = False,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+
+        # retrieve some shape info
+        batch_size, seq_length, _ = x.shape
+
+        # condtion mlp
+        cond_embedding = self.cond_mlp(cond)  # (B, T, C)
+
+        # condition mel
+        x = self.mel_mlp(x)
+
+        # diffusion step embedding
+        diffusion_step = self.diff_step_embedding(diffusion_step).to(x.device)
+        diffusion_step = self.diff_step_mlp(diffusion_step)  # (B, C)
+        x = x + cond_embedding
+
+        inputs_embeds = x
+        attention_mask = x_mask
+
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length,
+                seq_length + past_key_values_length,
+                dtype=torch.long,
+                device=device,
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        # embed positions
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past),
+                dtype=torch.bool,
+                device=inputs_embeds.device,
+            )
+        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask,
+            (batch_size, seq_length),
+            inputs_embeds,
+            past_key_values_length,
+        )
+
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        all_layer_hidden_states = []
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = (
+                past_key_values[idx] if past_key_values is not None else None
+            )
+
+            if self.gradient_checkpointing and self.training:
+                raise NotImplementedError
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, output_attentions, None)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    None,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                    cond_embedding=diffusion_step,
+                )
+
+            hidden_states = layer_outputs[0]
+            all_layer_hidden_states.append(hidden_states.clone())
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states, cond_embedding=diffusion_step)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+
+        hidden_states = self.mel_out_mlp(hidden_states)
+
+        # if not return_dict:
+        #     return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        # return BaseModelOutputWithPast(
+        #     last_hidden_state=hidden_states,
+        #     past_key_values=next_cache,
+        #     hidden_states=all_hidden_states,
+        #     attentions=all_self_attns,
+        # )
+        if return_dict:
+            return {
+                "output": hidden_states,
+                "hidden_states": all_layer_hidden_states,
+            }
+
+        return hidden_states

requirements.txt

@@ -0,0 +1,22 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@develop
+# model-specific deps below:
+spaces
+torch==2.8.0
+torchaudio==2.8.0
+accelerate
+# newest 4.x: reads the checkpoint's tokenizer.json but keeps pre-5.0 Llama internals
+transformers==4.57.6
+torchvision
+einops
+librosa
+huggingface_hub
+ruamel.yaml
+pyyaml
+setuptools<81
+pyworld
+praat-parselmouth
+torchcrepe
+openai-whisper
+json5
+ipython
+soundfile

utils/f0.py

@@ -0,0 +1,275 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import librosa
+import numpy as np
+import torch
+import parselmouth
+import torchcrepe
+import pyworld as pw
+
+
+def f0_to_coarse(f0, pitch_bin, f0_min, f0_max):
+    """
+    Convert f0 (Hz) to pitch (mel scale), and then quantize the mel-scale pitch to the
+    range from [1, 2, 3, ..., pitch_bin-1]
+
+    Reference: https://en.wikipedia.org/wiki/Mel_scale
+
+    Args:
+        f0 (array or Tensor): Hz
+        pitch_bin (int): the vocabulary size
+        f0_min (int): the minimum f0 (Hz)
+        f0_max (int): the maximum f0 (Hz)
+
+    Returns:
+        quantized f0 (array or Tensor)
+    """
+    f0_mel_min = 1127 * np.log(1 + f0_min / 700)
+    f0_mel_max = 1127 * np.log(1 + f0_max / 700)
+
+    is_torch = isinstance(f0, torch.Tensor)
+    f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700)
+    f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (pitch_bin - 2) / (
+        f0_mel_max - f0_mel_min
+    ) + 1
+
+    f0_mel[f0_mel <= 1] = 1
+    f0_mel[f0_mel > pitch_bin - 1] = pitch_bin - 1
+    f0_coarse = (f0_mel + 0.5).long() if is_torch else np.rint(f0_mel).astype(np.int32)
+    assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (
+        f0_coarse.max(),
+        f0_coarse.min(),
+    )
+    return f0_coarse
+
+
+def interpolate(f0):
+    """Interpolate the unvoiced part. Thus the f0 can be passed to a subtractive synthesizer.
+    Args:
+        f0: A numpy array of shape (seq_len,)
+    Returns:
+        f0: Interpolated f0 of shape (seq_len,)
+        uv: Unvoiced part of shape (seq_len,)
+    """
+    uv = f0 == 0
+    if len(f0[~uv]) > 0:
+        # interpolate the unvoiced f0
+        f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv])
+        uv = uv.astype("float")
+        uv = np.min(np.array([uv[:-2], uv[1:-1], uv[2:]]), axis=0)
+        uv = np.pad(uv, (1, 1))
+    return f0, uv
+
+
+def get_log_f0(f0):
+    f0[np.where(f0 == 0)] = 1
+    log_f0 = np.log(f0)
+    return log_f0
+
+
+def get_f0_features_using_pyin(audio, cfg):
+    """Using pyin to extract the f0 feature.
+    Args:
+        audio
+        fs
+        win_length
+        hop_length
+        f0_min
+        f0_max
+    Returns:
+        f0: numpy array of shape (frame_len,)
+    """
+    f0, voiced_flag, voiced_probs = librosa.pyin(
+        y=audio,
+        fmin=cfg.f0_min,
+        fmax=cfg.f0_max,
+        sr=cfg.sample_rate,
+        win_length=cfg.win_size,
+        hop_length=cfg.hop_size,
+    )
+    # Set nan to 0
+    f0[voiced_flag == False] = 0
+    return f0
+
+
+def get_f0_features_using_parselmouth(audio, cfg, speed=1):
+    """Using parselmouth to extract the f0 feature.
+    Args:
+        audio
+        mel_len
+        hop_length
+        fs
+        f0_min
+        f0_max
+        speed(default=1)
+    Returns:
+        f0: numpy array of shape (frame_len,)
+        pitch_coarse: numpy array of shape (frame_len,)
+    """
+    hop_size = int(np.round(cfg.hop_size * speed))
+
+    # Calculate the time step for pitch extraction
+    time_step = hop_size / cfg.sample_rate * 1000
+
+    f0 = (
+        parselmouth.Sound(audio, cfg.sample_rate)
+        .to_pitch_ac(
+            time_step=time_step / 1000,
+            voicing_threshold=0.6,
+            pitch_floor=cfg.f0_min,
+            pitch_ceiling=cfg.f0_max,
+        )
+        .selected_array["frequency"]
+    )
+    return f0
+
+
+def get_f0_features_using_dio(audio, cfg):
+    """Using dio to extract the f0 feature.
+    Args:
+        audio
+        mel_len
+        fs
+        hop_length
+        f0_min
+        f0_max
+    Returns:
+        f0: numpy array of shape (frame_len,)
+    """
+    # Get the raw f0
+    _f0, t = pw.dio(
+        audio.astype("double"),
+        cfg.sample_rate,
+        f0_floor=cfg.f0_min,
+        f0_ceil=cfg.f0_max,
+        channels_in_octave=2,
+        frame_period=(1000 * cfg.hop_size / cfg.sample_rate),
+    )
+    # Get the f0
+    f0 = pw.stonemask(audio.astype("double"), _f0, t, cfg.sample_rate)
+    return f0
+
+
+def get_f0_features_using_harvest(audio, mel_len, fs, hop_length, f0_min, f0_max):
+    """Using harvest to extract the f0 feature.
+    Args:
+        audio
+        mel_len
+        fs
+        hop_length
+        f0_min
+        f0_max
+    Returns:
+        f0: numpy array of shape (frame_len,)
+    """
+    f0, _ = pw.harvest(
+        audio.astype("double"),
+        fs,
+        f0_floor=f0_min,
+        f0_ceil=f0_max,
+        frame_period=(1000 * hop_length / fs),
+    )
+    f0 = f0.astype("float")[:mel_len]
+    return f0
+
+
+def get_f0_features_using_crepe(
+    audio, mel_len, fs, hop_length, hop_length_new, f0_min, f0_max, threshold=0.3
+):
+    """Using torchcrepe to extract the f0 feature.
+    Args:
+        audio
+        mel_len
+        fs
+        hop_length
+        hop_length_new
+        f0_min
+        f0_max
+        threshold(default=0.3)
+    Returns:
+        f0: numpy array of shape (frame_len,)
+    """
+    # Currently, crepe only supports 16khz audio
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    audio_16k = librosa.resample(audio, orig_sr=fs, target_sr=16000)
+    audio_16k_torch = torch.FloatTensor(audio_16k).unsqueeze(0).to(device)
+
+    # Get the raw pitch
+    f0, pd = torchcrepe.predict(
+        audio_16k_torch,
+        16000,
+        hop_length_new,
+        f0_min,
+        f0_max,
+        pad=True,
+        model="full",
+        batch_size=1024,
+        device=device,
+        return_periodicity=True,
+    )
+
+    # Filter, de-silence, set up threshold for unvoiced part
+    pd = torchcrepe.filter.median(pd, 3)
+    pd = torchcrepe.threshold.Silence(-60.0)(pd, audio_16k_torch, 16000, hop_length_new)
+    f0 = torchcrepe.threshold.At(threshold)(f0, pd)
+    f0 = torchcrepe.filter.mean(f0, 3)
+
+    # Convert unvoiced part to 0hz
+    f0 = torch.where(torch.isnan(f0), torch.full_like(f0, 0), f0)
+
+    # Interpolate f0
+    nzindex = torch.nonzero(f0[0]).squeeze()
+    f0 = torch.index_select(f0[0], dim=0, index=nzindex).cpu().numpy()
+    time_org = 0.005 * nzindex.cpu().numpy()
+    time_frame = np.arange(mel_len) * hop_length / fs
+    f0 = np.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
+    return f0
+
+
+def get_f0(audio, cfg, use_interpolate=False, return_uv=False):
+    if cfg.pitch_extractor == "dio":
+        f0 = get_f0_features_using_dio(audio, cfg)
+    elif cfg.pitch_extractor == "pyin":
+        f0 = get_f0_features_using_pyin(audio, cfg)
+    elif cfg.pitch_extractor == "parselmouth":
+        f0 = get_f0_features_using_parselmouth(audio, cfg)
+
+    if use_interpolate:
+        f0, uv = interpolate(f0)
+    else:
+        uv = f0 == 0
+
+    if return_uv:
+        return f0, uv
+
+    return f0
+
+
+def get_cents(f0_hz):
+    """
+    F_{cent} = 1200 * log2 (F/440)
+
+    Reference:
+        APSIPA'17, Perceptual Evaluation of Singing Quality
+    """
+    voiced_f0 = f0_hz[f0_hz != 0]
+    return 1200 * np.log2(voiced_f0 / 440)
+
+
+def get_pitch_derivatives(f0_hz):
+    """
+    f0_hz: (,T)
+    """
+    f0_cent = get_cents(f0_hz)
+    return f0_cent[1:] - f0_cent[:-1]
+
+
+def get_pitch_sub_median(f0_hz):
+    """
+    f0_hz: (,T)
+    """
+    f0_cent = get_cents(f0_hz)
+    return f0_cent - np.median(f0_cent)

utils/hparam.py

@@ -0,0 +1,660 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# This code is modified from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/training/python/training/hparam.py  pylint: disable=line-too-long
+"""Hyperparameter values."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+import numbers
+import re
+import six
+
+# Define the regular expression for parsing a single clause of the input
+# (delimited by commas).  A legal clause looks like:
+#   <variable name>[<index>]? = <rhs>
+# where <rhs> is either a single token or [] enclosed list of tokens.
+# For example:  "var[1] = a" or "x = [1,2,3]"
+PARAM_RE = re.compile(
+    r"""
+  (?P<name>[a-zA-Z][\w\.]*)      # variable name: "var" or "x"
+  (\[\s*(?P<index>\d+)\s*\])?  # (optional) index: "1" or None
+  \s*=\s*
+  ((?P<val>[^,\[]*)            # single value: "a" or None
+   |
+   \[(?P<vals>[^\]]*)\])       # list of values: None or "1,2,3"
+  ($|,\s*)""",
+    re.VERBOSE,
+)
+
+
+def _parse_fail(name, var_type, value, values):
+    """Helper function for raising a value error for bad assignment."""
+    raise ValueError(
+        "Could not parse hparam '%s' of type '%s' with value '%s' in %s"
+        % (name, var_type.__name__, value, values)
+    )
+
+
+def _reuse_fail(name, values):
+    """Helper function for raising a value error for reuse of name."""
+    raise ValueError("Multiple assignments to variable '%s' in %s" % (name, values))
+
+
+def _process_scalar_value(name, parse_fn, var_type, m_dict, values, results_dictionary):
+    """Update results_dictionary with a scalar value.
+
+    Used to update the results_dictionary to be returned by parse_values when
+    encountering a clause with a scalar RHS (e.g.  "s=5" or "arr[0]=5".)
+
+    Mutates results_dictionary.
+
+    Args:
+      name: Name of variable in assignment ("s" or "arr").
+      parse_fn: Function for parsing the actual value.
+      var_type: Type of named variable.
+      m_dict: Dictionary constructed from regex parsing.
+        m_dict['val']: RHS value (scalar)
+        m_dict['index']: List index value (or None)
+      values: Full expression being parsed
+      results_dictionary: The dictionary being updated for return by the parsing
+        function.
+
+    Raises:
+      ValueError: If the name has already been used.
+    """
+    try:
+        parsed_value = parse_fn(m_dict["val"])
+    except ValueError:
+        _parse_fail(name, var_type, m_dict["val"], values)
+
+    # If no index is provided
+    if not m_dict["index"]:
+        if name in results_dictionary:
+            _reuse_fail(name, values)
+        results_dictionary[name] = parsed_value
+    else:
+        if name in results_dictionary:
+            # The name has already been used as a scalar, then it
+            # will be in this dictionary and map to a non-dictionary.
+            if not isinstance(results_dictionary.get(name), dict):
+                _reuse_fail(name, values)
+        else:
+            results_dictionary[name] = {}
+
+        index = int(m_dict["index"])
+        # Make sure the index position hasn't already been assigned a value.
+        if index in results_dictionary[name]:
+            _reuse_fail("{}[{}]".format(name, index), values)
+        results_dictionary[name][index] = parsed_value
+
+
+def _process_list_value(name, parse_fn, var_type, m_dict, values, results_dictionary):
+    """Update results_dictionary from a list of values.
+
+    Used to update results_dictionary to be returned by parse_values when
+    encountering a clause with a list RHS (e.g.  "arr=[1,2,3]".)
+
+    Mutates results_dictionary.
+
+    Args:
+      name: Name of variable in assignment ("arr").
+      parse_fn: Function for parsing individual values.
+      var_type: Type of named variable.
+      m_dict: Dictionary constructed from regex parsing.
+        m_dict['val']: RHS value (scalar)
+      values: Full expression being parsed
+      results_dictionary: The dictionary being updated for return by the parsing
+        function.
+
+    Raises:
+      ValueError: If the name has an index or the values cannot be parsed.
+    """
+    if m_dict["index"] is not None:
+        raise ValueError("Assignment of a list to a list index.")
+    elements = filter(None, re.split("[ ,]", m_dict["vals"]))
+    # Make sure the name hasn't already been assigned a value
+    if name in results_dictionary:
+        raise _reuse_fail(name, values)
+    try:
+        results_dictionary[name] = [parse_fn(e) for e in elements]
+    except ValueError:
+        _parse_fail(name, var_type, m_dict["vals"], values)
+
+
+def _cast_to_type_if_compatible(name, param_type, value):
+    """Cast hparam to the provided type, if compatible.
+
+    Args:
+      name: Name of the hparam to be cast.
+      param_type: The type of the hparam.
+      value: The value to be cast, if compatible.
+
+    Returns:
+      The result of casting `value` to `param_type`.
+
+    Raises:
+      ValueError: If the type of `value` is not compatible with param_type.
+        * If `param_type` is a string type, but `value` is not.
+        * If `param_type` is a boolean, but `value` is not, or vice versa.
+        * If `param_type` is an integer type, but `value` is not.
+        * If `param_type` is a float type, but `value` is not a numeric type.
+    """
+    fail_msg = "Could not cast hparam '%s' of type '%s' from value %r" % (
+        name,
+        param_type,
+        value,
+    )
+
+    # Some callers use None, for which we can't do any casting/checking. :(
+    if issubclass(param_type, type(None)):
+        return value
+
+    # Avoid converting a non-string type to a string.
+    if issubclass(param_type, (six.string_types, six.binary_type)) and not isinstance(
+        value, (six.string_types, six.binary_type)
+    ):
+        raise ValueError(fail_msg)
+
+    # Avoid converting a number or string type to a boolean or vice versa.
+    if issubclass(param_type, bool) != isinstance(value, bool):
+        raise ValueError(fail_msg)
+
+    # Avoid converting float to an integer (the reverse is fine).
+    if issubclass(param_type, numbers.Integral) and not isinstance(
+        value, numbers.Integral
+    ):
+        raise ValueError(fail_msg)
+
+    # Avoid converting a non-numeric type to a numeric type.
+    if issubclass(param_type, numbers.Number) and not isinstance(value, numbers.Number):
+        raise ValueError(fail_msg)
+
+    return param_type(value)
+
+
+def parse_values(values, type_map, ignore_unknown=False):
+    """Parses hyperparameter values from a string into a python map.
+
+    `values` is a string containing comma-separated `name=value` pairs.
+    For each pair, the value of the hyperparameter named `name` is set to
+    `value`.
+
+    If a hyperparameter name appears multiple times in `values`, a ValueError
+    is raised (e.g. 'a=1,a=2', 'a[1]=1,a[1]=2').
+
+    If a hyperparameter name in both an index assignment and scalar assignment,
+    a ValueError is raised.  (e.g. 'a=[1,2,3],a[0] = 1').
+
+    The hyperparameter name may contain '.' symbols, which will result in an
+    attribute name that is only accessible through the getattr and setattr
+    functions.  (And must be first explicit added through add_hparam.)
+
+    WARNING: Use of '.' in your variable names is allowed, but is not well
+    supported and not recommended.
+
+    The `value` in `name=value` must follows the syntax according to the
+    type of the parameter:
+
+    *  Scalar integer: A Python-parsable integer point value.  E.g.: 1,
+       100, -12.
+    *  Scalar float: A Python-parsable floating point value.  E.g.: 1.0,
+       -.54e89.
+    *  Boolean: Either true or false.
+    *  Scalar string: A non-empty sequence of characters, excluding comma,
+       spaces, and square brackets.  E.g.: foo, bar_1.
+    *  List: A comma separated list of scalar values of the parameter type
+       enclosed in square brackets.  E.g.: [1,2,3], [1.0,1e-12], [high,low].
+
+    When index assignment is used, the corresponding type_map key should be the
+    list name.  E.g. for "arr[1]=0" the type_map must have the key "arr" (not
+    "arr[1]").
+
+    Args:
+      values: String.  Comma separated list of `name=value` pairs where
+        'value' must follow the syntax described above.
+      type_map: A dictionary mapping hyperparameter names to types.  Note every
+        parameter name in values must be a key in type_map.  The values must
+        conform to the types indicated, where a value V is said to conform to a
+        type T if either V has type T, or V is a list of elements of type T.
+        Hence, for a multidimensional parameter 'x' taking float values,
+        'x=[0.1,0.2]' will parse successfully if type_map['x'] = float.
+      ignore_unknown: Bool. Whether values that are missing a type in type_map
+        should be ignored. If set to True, a ValueError will not be raised for
+        unknown hyperparameter type.
+
+    Returns:
+      A python map mapping each name to either:
+      * A scalar value.
+      * A list of scalar values.
+      * A dictionary mapping index numbers to scalar values.
+      (e.g. "x=5,L=[1,2],arr[1]=3" results in {'x':5,'L':[1,2],'arr':{1:3}}")
+
+    Raises:
+      ValueError: If there is a problem with input.
+      * If `values` cannot be parsed.
+      * If a list is assigned to a list index (e.g. 'a[1] = [1,2,3]').
+      * If the same rvalue is assigned two different values (e.g. 'a=1,a=2',
+        'a[1]=1,a[1]=2', or 'a=1,a=[1]')
+    """
+    results_dictionary = {}
+    pos = 0
+    while pos < len(values):
+        m = PARAM_RE.match(values, pos)
+        if not m:
+            raise ValueError("Malformed hyperparameter value: %s" % values[pos:])
+        # Check that there is a comma between parameters and move past it.
+        pos = m.end()
+        # Parse the values.
+        m_dict = m.groupdict()
+        name = m_dict["name"]
+        if name not in type_map:
+            if ignore_unknown:
+                continue
+            raise ValueError("Unknown hyperparameter type for %s" % name)
+        type_ = type_map[name]
+
+        # Set up correct parsing function (depending on whether type_ is a bool)
+        if type_ == bool:
+
+            def parse_bool(value):
+                if value in ["true", "True"]:
+                    return True
+                elif value in ["false", "False"]:
+                    return False
+                else:
+                    try:
+                        return bool(int(value))
+                    except ValueError:
+                        _parse_fail(name, type_, value, values)
+
+            parse = parse_bool
+        else:
+            parse = type_
+
+        # If a singe value is provided
+        if m_dict["val"] is not None:
+            _process_scalar_value(
+                name, parse, type_, m_dict, values, results_dictionary
+            )
+
+        # If the assigned value is a list:
+        elif m_dict["vals"] is not None:
+            _process_list_value(name, parse, type_, m_dict, values, results_dictionary)
+
+        else:  # Not assigned a list or value
+            _parse_fail(name, type_, "", values)
+
+    return results_dictionary
+
+
+class HParams(object):
+    """Class to hold a set of hyperparameters as name-value pairs.
+
+    A `HParams` object holds hyperparameters used to build and train a model,
+    such as the number of hidden units in a neural net layer or the learning rate
+    to use when training.
+
+    You first create a `HParams` object by specifying the names and values of the
+    hyperparameters.
+
+    To make them easily accessible the parameter names are added as direct
+    attributes of the class.  A typical usage is as follows:
+
+    ```python
+    # Create a HParams object specifying names and values of the model
+    # hyperparameters:
+    hparams = HParams(learning_rate=0.1, num_hidden_units=100)
+
+    # The hyperparameter are available as attributes of the HParams object:
+    hparams.learning_rate ==> 0.1
+    hparams.num_hidden_units ==> 100
+    ```
+
+    Hyperparameters have type, which is inferred from the type of their value
+    passed at construction type.   The currently supported types are: integer,
+    float, boolean, string, and list of integer, float, boolean, or string.
+
+    You can override hyperparameter values by calling the
+    [`parse()`](#HParams.parse) method, passing a string of comma separated
+    `name=value` pairs.  This is intended to make it possible to override
+    any hyperparameter values from a single command-line flag to which
+    the user passes 'hyper-param=value' pairs.  It avoids having to define
+    one flag for each hyperparameter.
+
+    The syntax expected for each value depends on the type of the parameter.
+    See `parse()` for a description of the syntax.
+
+    Example:
+
+    ```python
+    # Define a command line flag to pass name=value pairs.
+    # For example using argparse:
+    import argparse
+    parser = argparse.ArgumentParser(description='Train my model.')
+    parser.add_argument('--hparams', type=str,
+                        help='Comma separated list of "name=value" pairs.')
+    args = parser.parse_args()
+    ...
+    def my_program():
+      # Create a HParams object specifying the names and values of the
+      # model hyperparameters:
+      hparams = tf.HParams(learning_rate=0.1, num_hidden_units=100,
+                           activations=['relu', 'tanh'])
+
+      # Override hyperparameters values by parsing the command line
+      hparams.parse(args.hparams)
+
+      # If the user passed `--hparams=learning_rate=0.3` on the command line
+      # then 'hparams' has the following attributes:
+      hparams.learning_rate ==> 0.3
+      hparams.num_hidden_units ==> 100
+      hparams.activations ==> ['relu', 'tanh']
+
+      # If the hyperparameters are in json format use parse_json:
+      hparams.parse_json('{"learning_rate": 0.3, "activations": "relu"}')
+    ```
+    """
+
+    _HAS_DYNAMIC_ATTRIBUTES = True  # Required for pytype checks.
+
+    def __init__(self, model_structure=None, **kwargs):
+        """Create an instance of `HParams` from keyword arguments.
+
+        The keyword arguments specify name-values pairs for the hyperparameters.
+        The parameter types are inferred from the type of the values passed.
+
+        The parameter names are added as attributes of `HParams` object, so they
+        can be accessed directly with the dot notation `hparams._name_`.
+
+        Example:
+
+        ```python
+        # Define 3 hyperparameters: 'learning_rate' is a float parameter,
+        # 'num_hidden_units' an integer parameter, and 'activation' a string
+        # parameter.
+        hparams = tf.HParams(
+            learning_rate=0.1, num_hidden_units=100, activation='relu')
+
+        hparams.activation ==> 'relu'
+        ```
+
+        Note that a few names are reserved and cannot be used as hyperparameter
+        names.  If you use one of the reserved name the constructor raises a
+        `ValueError`.
+
+        Args:
+          model_structure: An instance of ModelStructure, defining the feature
+            crosses to be used in the Trial.
+          **kwargs: Key-value pairs where the key is the hyperparameter name and
+            the value is the value for the parameter.
+
+        Raises:
+          ValueError: If both `hparam_def` and initialization values are provided,
+            or if one of the arguments is invalid.
+
+        """
+        # Register the hyperparameters and their type in _hparam_types.
+        # This simplifies the implementation of parse().
+        # _hparam_types maps the parameter name to a tuple (type, bool).
+        # The type value is the type of the parameter for scalar hyperparameters,
+        # or the type of the list elements for multidimensional hyperparameters.
+        # The bool value is True if the value is a list, False otherwise.
+        self._hparam_types = {}
+        self._model_structure = model_structure
+        for name, value in six.iteritems(kwargs):
+            self.add_hparam(name, value)
+
+    def add_hparam(self, name, value):
+        """Adds {name, value} pair to hyperparameters.
+
+        Args:
+          name: Name of the hyperparameter.
+          value: Value of the hyperparameter. Can be one of the following types:
+            int, float, string, int list, float list, or string list.
+
+        Raises:
+          ValueError: if one of the arguments is invalid.
+        """
+        # Keys in kwargs are unique, but 'name' could the name of a pre-existing
+        # attribute of this object.  In that case we refuse to use it as a
+        # hyperparameter name.
+        if getattr(self, name, None) is not None:
+            raise ValueError("Hyperparameter name is reserved: %s" % name)
+        if isinstance(value, (list, tuple)):
+            if not value:
+                raise ValueError(
+                    "Multi-valued hyperparameters cannot be empty: %s" % name
+                )
+            self._hparam_types[name] = (type(value[0]), True)
+        else:
+            self._hparam_types[name] = (type(value), False)
+        setattr(self, name, value)
+
+    def set_hparam(self, name, value):
+        """Set the value of an existing hyperparameter.
+
+        This function verifies that the type of the value matches the type of the
+        existing hyperparameter.
+
+        Args:
+          name: Name of the hyperparameter.
+          value: New value of the hyperparameter.
+
+        Raises:
+          KeyError: If the hyperparameter doesn't exist.
+          ValueError: If there is a type mismatch.
+        """
+        param_type, is_list = self._hparam_types[name]
+        if isinstance(value, list):
+            if not is_list:
+                raise ValueError(
+                    "Must not pass a list for single-valued parameter: %s" % name
+                )
+            setattr(
+                self,
+                name,
+                [_cast_to_type_if_compatible(name, param_type, v) for v in value],
+            )
+        else:
+            if is_list:
+                raise ValueError(
+                    "Must pass a list for multi-valued parameter: %s." % name
+                )
+            setattr(self, name, _cast_to_type_if_compatible(name, param_type, value))
+
+    def del_hparam(self, name):
+        """Removes the hyperparameter with key 'name'.
+
+        Does nothing if it isn't present.
+
+        Args:
+          name: Name of the hyperparameter.
+        """
+        if hasattr(self, name):
+            delattr(self, name)
+            del self._hparam_types[name]
+
+    def parse(self, values):
+        """Override existing hyperparameter values, parsing new values from a string.
+
+        See parse_values for more detail on the allowed format for values.
+
+        Args:
+          values: String.  Comma separated list of `name=value` pairs where 'value'
+            must follow the syntax described above.
+
+        Returns:
+          The `HParams` instance.
+
+        Raises:
+          ValueError: If `values` cannot be parsed or a hyperparameter in `values`
+          doesn't exist.
+        """
+        type_map = {}
+        for name, t in self._hparam_types.items():
+            param_type, _ = t
+            type_map[name] = param_type
+
+        values_map = parse_values(values, type_map)
+        return self.override_from_dict(values_map)
+
+    def override_from_dict(self, values_dict):
+        """Override existing hyperparameter values, parsing new values from a dictionary.
+
+        Args:
+          values_dict: Dictionary of name:value pairs.
+
+        Returns:
+          The `HParams` instance.
+
+        Raises:
+          KeyError: If a hyperparameter in `values_dict` doesn't exist.
+          ValueError: If `values_dict` cannot be parsed.
+        """
+        for name, value in values_dict.items():
+            self.set_hparam(name, value)
+        return self
+
+    def set_model_structure(self, model_structure):
+        self._model_structure = model_structure
+
+    def get_model_structure(self):
+        return self._model_structure
+
+    def to_json(self, indent=None, separators=None, sort_keys=False):
+        """Serializes the hyperparameters into JSON.
+
+        Args:
+          indent: If a non-negative integer, JSON array elements and object members
+            will be pretty-printed with that indent level. An indent level of 0, or
+            negative, will only insert newlines. `None` (the default) selects the
+            most compact representation.
+          separators: Optional `(item_separator, key_separator)` tuple. Default is
+            `(', ', ': ')`.
+          sort_keys: If `True`, the output dictionaries will be sorted by key.
+
+        Returns:
+          A JSON string.
+        """
+
+        def remove_callables(x):
+            """Omit callable elements from input with arbitrary nesting."""
+            if isinstance(x, dict):
+                return {
+                    k: remove_callables(v)
+                    for k, v in six.iteritems(x)
+                    if not callable(v)
+                }
+            elif isinstance(x, list):
+                return [remove_callables(i) for i in x if not callable(i)]
+            return x
+
+        return json.dumps(
+            remove_callables(self.values()),
+            indent=indent,
+            separators=separators,
+            sort_keys=sort_keys,
+        )
+
+    def parse_json(self, values_json):
+        """Override existing hyperparameter values, parsing new values from a json object.
+
+        Args:
+          values_json: String containing a json object of name:value pairs.
+
+        Returns:
+          The `HParams` instance.
+
+        Raises:
+          KeyError: If a hyperparameter in `values_json` doesn't exist.
+          ValueError: If `values_json` cannot be parsed.
+        """
+        values_map = json.loads(values_json)
+        return self.override_from_dict(values_map)
+
+    def values(self):
+        """Return the hyperparameter values as a Python dictionary.
+
+        Returns:
+          A dictionary with hyperparameter names as keys.  The values are the
+          hyperparameter values.
+        """
+        return {n: getattr(self, n) for n in self._hparam_types.keys()}
+
+    def get(self, key, default=None):
+        """Returns the value of `key` if it exists, else `default`."""
+        if key in self._hparam_types:
+            # Ensure that default is compatible with the parameter type.
+            if default is not None:
+                param_type, is_param_list = self._hparam_types[key]
+                type_str = "list<%s>" % param_type if is_param_list else str(param_type)
+                fail_msg = (
+                    "Hparam '%s' of type '%s' is incompatible with "
+                    "default=%s" % (key, type_str, default)
+                )
+
+                is_default_list = isinstance(default, list)
+                if is_param_list != is_default_list:
+                    raise ValueError(fail_msg)
+
+                try:
+                    if is_default_list:
+                        for value in default:
+                            _cast_to_type_if_compatible(key, param_type, value)
+                    else:
+                        _cast_to_type_if_compatible(key, param_type, default)
+                except ValueError as e:
+                    raise ValueError("%s. %s" % (fail_msg, e))
+
+            return getattr(self, key)
+
+        return default
+
+    def __contains__(self, key):
+        return key in self._hparam_types
+
+    def __str__(self):
+        return str(sorted(self.values().items()))
+
+    def __repr__(self):
+        return "%s(%s)" % (type(self).__name__, self.__str__())
+
+    @staticmethod
+    def _get_kind_name(param_type, is_list):
+        """Returns the field name given parameter type and is_list.
+
+        Args:
+          param_type: Data type of the hparam.
+          is_list: Whether this is a list.
+
+        Returns:
+          A string representation of the field name.
+
+        Raises:
+          ValueError: If parameter type is not recognized.
+        """
+        if issubclass(param_type, bool):
+            # This check must happen before issubclass(param_type, six.integer_types),
+            # since Python considers bool to be a subclass of int.
+            typename = "bool"
+        elif issubclass(param_type, six.integer_types):
+            # Setting 'int' and 'long' types to be 'int64' to ensure the type is
+            # compatible with both Python2 and Python3.
+            typename = "int64"
+        elif issubclass(param_type, (six.string_types, six.binary_type)):
+            # Setting 'string' and 'bytes' types to be 'bytes' to ensure the type is
+            # compatible with both Python2 and Python3.
+            typename = "bytes"
+        elif issubclass(param_type, float):
+            typename = "float"
+        else:
+            raise ValueError("Unsupported parameter type: %s" % str(param_type))
+
+        suffix = "list" if is_list else "value"
+        return "_".join([typename, suffix])

utils/util.py

@@ -0,0 +1,689 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import collections
+import glob
+import os
+import random
+import time
+import argparse
+from collections import OrderedDict
+
+import json5
+import numpy as np
+import glob
+from torch.nn import functional as F
+
+try:
+    from ruamel.yaml import YAML as yaml
+except:
+    from ruamel_yaml import YAML as yaml
+
+import torch
+
+from utils.hparam import HParams
+import logging
+from logging import handlers
+
+
+def str2bool(v):
+    """Used in argparse.ArgumentParser.add_argument to indicate
+    that a type is a bool type and user can enter
+
+        - yes, true, t, y, 1, to represent True
+        - no, false, f, n, 0, to represent False
+
+    See https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse  # noqa
+    """
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+def find_checkpoint_of_mapper(mapper_ckpt_dir):
+    mapper_ckpts = glob.glob(os.path.join(mapper_ckpt_dir, "ckpts/*.pt"))
+
+    # Select the max steps
+    mapper_ckpts.sort()
+    mapper_weights_file = mapper_ckpts[-1]
+    return mapper_weights_file
+
+
+def pad_f0_to_tensors(f0s, batched=None):
+    # Initialize
+    tensors = []
+
+    if batched == None:
+        # Get the max frame for padding
+        size = -1
+        for f0 in f0s:
+            size = max(size, f0.shape[-1])
+
+        tensor = torch.zeros(len(f0s), size)
+
+        for i, f0 in enumerate(f0s):
+            tensor[i, : f0.shape[-1]] = f0[:]
+
+        tensors.append(tensor)
+    else:
+        start = 0
+        while start + batched - 1 < len(f0s):
+            end = start + batched - 1
+
+            # Get the max frame for padding
+            size = -1
+            for i in range(start, end + 1):
+                size = max(size, f0s[i].shape[-1])
+
+            tensor = torch.zeros(batched, size)
+
+            for i in range(start, end + 1):
+                tensor[i - start, : f0s[i].shape[-1]] = f0s[i][:]
+
+            tensors.append(tensor)
+
+            start = start + batched
+
+        if start != len(f0s):
+            end = len(f0s)
+
+            # Get the max frame for padding
+            size = -1
+            for i in range(start, end):
+                size = max(size, f0s[i].shape[-1])
+
+            tensor = torch.zeros(len(f0s) - start, size)
+
+            for i in range(start, end):
+                tensor[i - start, : f0s[i].shape[-1]] = f0s[i][:]
+
+            tensors.append(tensor)
+
+    return tensors
+
+
+def pad_mels_to_tensors(mels, batched=None):
+    """
+    Args:
+        mels: A list of mel-specs
+    Returns:
+        tensors: A list of tensors containing the batched mel-specs
+        mel_frames: A list of tensors containing the frames of the original mel-specs
+    """
+    # Initialize
+    tensors = []
+    mel_frames = []
+
+    # Split mel-specs into batches to avoid cuda memory exceed
+    if batched == None:
+        # Get the max frame for padding
+        size = -1
+        for mel in mels:
+            size = max(size, mel.shape[-1])
+
+        tensor = torch.zeros(len(mels), mels[0].shape[0], size)
+        mel_frame = torch.zeros(len(mels), dtype=torch.int32)
+
+        for i, mel in enumerate(mels):
+            tensor[i, :, : mel.shape[-1]] = mel[:]
+            mel_frame[i] = mel.shape[-1]
+
+        tensors.append(tensor)
+        mel_frames.append(mel_frame)
+    else:
+        start = 0
+        while start + batched - 1 < len(mels):
+            end = start + batched - 1
+
+            # Get the max frame for padding
+            size = -1
+            for i in range(start, end + 1):
+                size = max(size, mels[i].shape[-1])
+
+            tensor = torch.zeros(batched, mels[0].shape[0], size)
+            mel_frame = torch.zeros(batched, dtype=torch.int32)
+
+            for i in range(start, end + 1):
+                tensor[i - start, :, : mels[i].shape[-1]] = mels[i][:]
+                mel_frame[i - start] = mels[i].shape[-1]
+
+            tensors.append(tensor)
+            mel_frames.append(mel_frame)
+
+            start = start + batched
+
+        if start != len(mels):
+            end = len(mels)
+
+            # Get the max frame for padding
+            size = -1
+            for i in range(start, end):
+                size = max(size, mels[i].shape[-1])
+
+            tensor = torch.zeros(len(mels) - start, mels[0].shape[0], size)
+            mel_frame = torch.zeros(len(mels) - start, dtype=torch.int32)
+
+            for i in range(start, end):
+                tensor[i - start, :, : mels[i].shape[-1]] = mels[i][:]
+                mel_frame[i - start] = mels[i].shape[-1]
+
+            tensors.append(tensor)
+            mel_frames.append(mel_frame)
+
+    return tensors, mel_frames
+
+
+def load_model_config(args):
+    """Load model configurations (in args.json under checkpoint directory)
+
+    Args:
+        args (ArgumentParser): arguments to run bins/preprocess.py
+
+    Returns:
+        dict: dictionary that stores model configurations
+    """
+    if args.checkpoint_dir is None:
+        assert args.checkpoint_file is not None
+        checkpoint_dir = os.path.split(args.checkpoint_file)[0]
+    else:
+        checkpoint_dir = args.checkpoint_dir
+    config_path = os.path.join(checkpoint_dir, "args.json")
+    print("config_path: ", config_path)
+
+    config = load_config(config_path)
+    return config
+
+
+def remove_and_create(dir):
+    if os.path.exists(dir):
+        os.system("rm -r {}".format(dir))
+    os.makedirs(dir, exist_ok=True)
+
+
+def has_existed(path, warning=False):
+    if not warning:
+        return os.path.exists(path)
+
+    if os.path.exists(path):
+        answer = input(
+            "The path {} has existed. \nInput 'y' (or hit Enter) to skip it, and input 'n' to re-write it [y/n]\n".format(
+                path
+            )
+        )
+        if not answer == "n":
+            return True
+
+    return False
+
+
+def remove_older_ckpt(saved_model_name, checkpoint_dir, max_to_keep=5):
+    if os.path.exists(os.path.join(checkpoint_dir, "checkpoint")):
+        with open(os.path.join(checkpoint_dir, "checkpoint"), "r") as f:
+            ckpts = [x.strip() for x in f.readlines()]
+    else:
+        ckpts = []
+    ckpts.append(saved_model_name)
+    for item in ckpts[:-max_to_keep]:
+        if os.path.exists(os.path.join(checkpoint_dir, item)):
+            os.remove(os.path.join(checkpoint_dir, item))
+    with open(os.path.join(checkpoint_dir, "checkpoint"), "w") as f:
+        for item in ckpts[-max_to_keep:]:
+            f.write("{}\n".format(item))
+
+
+def set_all_random_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.random.manual_seed(seed)
+
+
+def save_checkpoint(
+    args,
+    generator,
+    g_optimizer,
+    step,
+    discriminator=None,
+    d_optimizer=None,
+    max_to_keep=5,
+):
+    saved_model_name = "model.ckpt-{}.pt".format(step)
+    checkpoint_path = os.path.join(args.checkpoint_dir, saved_model_name)
+
+    if discriminator and d_optimizer:
+        torch.save(
+            {
+                "generator": generator.state_dict(),
+                "discriminator": discriminator.state_dict(),
+                "g_optimizer": g_optimizer.state_dict(),
+                "d_optimizer": d_optimizer.state_dict(),
+                "global_step": step,
+            },
+            checkpoint_path,
+        )
+    else:
+        torch.save(
+            {
+                "generator": generator.state_dict(),
+                "g_optimizer": g_optimizer.state_dict(),
+                "global_step": step,
+            },
+            checkpoint_path,
+        )
+
+    print("Saved checkpoint: {}".format(checkpoint_path))
+
+    if os.path.exists(os.path.join(args.checkpoint_dir, "checkpoint")):
+        with open(os.path.join(args.checkpoint_dir, "checkpoint"), "r") as f:
+            ckpts = [x.strip() for x in f.readlines()]
+    else:
+        ckpts = []
+    ckpts.append(saved_model_name)
+    for item in ckpts[:-max_to_keep]:
+        if os.path.exists(os.path.join(args.checkpoint_dir, item)):
+            os.remove(os.path.join(args.checkpoint_dir, item))
+    with open(os.path.join(args.checkpoint_dir, "checkpoint"), "w") as f:
+        for item in ckpts[-max_to_keep:]:
+            f.write("{}\n".format(item))
+
+
+def attempt_to_restore(
+    generator, g_optimizer, checkpoint_dir, discriminator=None, d_optimizer=None
+):
+    checkpoint_list = os.path.join(checkpoint_dir, "checkpoint")
+    if os.path.exists(checkpoint_list):
+        checkpoint_filename = open(checkpoint_list).readlines()[-1].strip()
+        checkpoint_path = os.path.join(checkpoint_dir, "{}".format(checkpoint_filename))
+        print("Restore from {}".format(checkpoint_path))
+        checkpoint = torch.load(checkpoint_path, map_location="cpu")
+        if generator:
+            if not list(generator.state_dict().keys())[0].startswith("module."):
+                raw_dict = checkpoint["generator"]
+                clean_dict = OrderedDict()
+                for k, v in raw_dict.items():
+                    if k.startswith("module."):
+                        clean_dict[k[7:]] = v
+                    else:
+                        clean_dict[k] = v
+                generator.load_state_dict(clean_dict)
+            else:
+                generator.load_state_dict(checkpoint["generator"])
+        if g_optimizer:
+            g_optimizer.load_state_dict(checkpoint["g_optimizer"])
+        global_step = 100000
+        if discriminator and "discriminator" in checkpoint.keys():
+            discriminator.load_state_dict(checkpoint["discriminator"])
+            global_step = checkpoint["global_step"]
+            print("restore discriminator")
+        if d_optimizer and "d_optimizer" in checkpoint.keys():
+            d_optimizer.load_state_dict(checkpoint["d_optimizer"])
+            print("restore d_optimizer...")
+    else:
+        global_step = 0
+    return global_step
+
+
+class ExponentialMovingAverage(object):
+    def __init__(self, decay):
+        self.decay = decay
+        self.shadow = {}
+
+    def register(self, name, val):
+        self.shadow[name] = val.clone()
+
+    def update(self, name, x):
+        assert name in self.shadow
+        update_delta = self.shadow[name] - x
+        self.shadow[name] -= (1.0 - self.decay) * update_delta
+
+
+def apply_moving_average(model, ema):
+    for name, param in model.named_parameters():
+        if name in ema.shadow:
+            ema.update(name, param.data)
+
+
+def register_model_to_ema(model, ema):
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            ema.register(name, param.data)
+
+
+class YParams(HParams):
+    def __init__(self, yaml_file):
+        if not os.path.exists(yaml_file):
+            raise IOError("yaml file: {} is not existed".format(yaml_file))
+        super().__init__()
+        self.d = collections.OrderedDict()
+        with open(yaml_file) as fp:
+            for _, v in yaml().load(fp).items():
+                for k1, v1 in v.items():
+                    try:
+                        if self.get(k1):
+                            self.set_hparam(k1, v1)
+                        else:
+                            self.add_hparam(k1, v1)
+                        self.d[k1] = v1
+                    except Exception:
+                        import traceback
+
+                        print(traceback.format_exc())
+
+    # @property
+    def get_elements(self):
+        return self.d.items()
+
+
+def override_config(base_config, new_config):
+    """Update new configurations in the original dict with the new dict
+
+    Args:
+        base_config (dict): original dict to be overridden
+        new_config (dict): dict with new configurations
+
+    Returns:
+        dict: updated configuration dict
+    """
+    for k, v in new_config.items():
+        if type(v) == dict:
+            if k not in base_config.keys():
+                base_config[k] = {}
+            base_config[k] = override_config(base_config[k], v)
+        else:
+            base_config[k] = v
+    return base_config
+
+
+def get_lowercase_keys_config(cfg):
+    """Change all keys in cfg to lower case
+
+    Args:
+        cfg (dict): dictionary that stores configurations
+
+    Returns:
+        dict: dictionary that stores configurations
+    """
+    updated_cfg = dict()
+    for k, v in cfg.items():
+        if type(v) == dict:
+            v = get_lowercase_keys_config(v)
+        updated_cfg[k.lower()] = v
+    return updated_cfg
+
+
+def _load_config(config_fn, lowercase=False):
+    """Load configurations into a dictionary
+
+    Args:
+        config_fn (str): path to configuration file
+        lowercase (bool, optional): whether changing keys to lower case. Defaults to False.
+
+    Returns:
+        dict: dictionary that stores configurations
+    """
+    with open(config_fn, "r") as f:
+        data = f.read()
+    config_ = json5.loads(data)
+    if "base_config" in config_:
+        # load configurations from new path
+        try:
+            p_config_path = os.path.join(os.getenv("WORK_DIR"), config_["base_config"])
+        except:
+            p_config_path = config_["base_config"]
+        p_config_ = _load_config(p_config_path)
+        config_ = override_config(p_config_, config_)
+    if lowercase:
+        # change keys in config_ to lower case
+        config_ = get_lowercase_keys_config(config_)
+    return config_
+
+
+def load_config(config_fn, lowercase=False):
+    """Load configurations into a dictionary
+
+    Args:
+        config_fn (str): path to configuration file
+        lowercase (bool, optional): _description_. Defaults to False.
+
+    Returns:
+        JsonHParams: an object that stores configurations
+    """
+    config_ = _load_config(config_fn, lowercase=lowercase)
+    # create an JsonHParams object with configuration dict
+    cfg = JsonHParams(**config_)
+    return cfg
+
+
+def save_config(save_path, cfg):
+    """Save configurations into a json file
+
+    Args:
+        save_path (str): path to save configurations
+        cfg (dict): dictionary that stores configurations
+    """
+    with open(save_path, "w") as f:
+        json5.dump(
+            cfg, f, ensure_ascii=False, indent=4, quote_keys=True, sort_keys=True
+        )
+
+
+class JsonHParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if type(v) == dict:
+                v = JsonHParams(**v)
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+
+class ValueWindow:
+    def __init__(self, window_size=100):
+        self._window_size = window_size
+        self._values = []
+
+    def append(self, x):
+        self._values = self._values[-(self._window_size - 1) :] + [x]
+
+    @property
+    def sum(self):
+        return sum(self._values)
+
+    @property
+    def count(self):
+        return len(self._values)
+
+    @property
+    def average(self):
+        return self.sum / max(1, self.count)
+
+    def reset(self):
+        self._values = []
+
+
+class Logger(object):
+    def __init__(
+        self,
+        filename,
+        level="info",
+        when="D",
+        backCount=10,
+        fmt="%(asctime)s : %(message)s",
+    ):
+        self.level_relations = {
+            "debug": logging.DEBUG,
+            "info": logging.INFO,
+            "warning": logging.WARNING,
+            "error": logging.ERROR,
+            "crit": logging.CRITICAL,
+        }
+        if level == "debug":
+            fmt = "%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s"
+        self.logger = logging.getLogger(filename)
+        format_str = logging.Formatter(fmt)
+        self.logger.setLevel(self.level_relations.get(level))
+        sh = logging.StreamHandler()
+        sh.setFormatter(format_str)
+        th = handlers.TimedRotatingFileHandler(
+            filename=filename, when=when, backupCount=backCount, encoding="utf-8"
+        )
+        th.setFormatter(format_str)
+        self.logger.addHandler(sh)
+        self.logger.addHandler(th)
+        self.logger.info(
+            "==========================New Starting Here=============================="
+        )
+
+
+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)
+
+
+def slice_segments(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = ids_str[i]
+        idx_end = idx_str + segment_size
+        ret[i] = x[i, :, idx_str:idx_end]
+    return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def subsequent_mask(length):
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def convert_pad_shape(pad_shape):
+    l = pad_shape[::-1]
+    pad_shape = [item for sublist in l for item in sublist]
+    return pad_shape
+
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def generate_path(duration, mask):
+    """
+    duration: [b, 1, t_x]
+    mask: [b, 1, t_y, t_x]
+    """
+    device = duration.device
+
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm
+
+
+def get_current_time():
+    pass
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """
+    Args:
+      lengths:
+        A 1-D tensor containing sentence lengths.
+      max_len:
+        The length of masks.
+    Returns:
+      Return a 2-D bool tensor, where masked positions
+      are filled with `True` and non-masked positions are
+      filled with `False`.
+
+    >>> lengths = torch.tensor([1, 3, 2, 5])
+    >>> make_pad_mask(lengths)
+    tensor([[False,  True,  True,  True,  True],
+            [False, False, False,  True,  True],
+            [False, False,  True,  True,  True],
+            [False, False, False, False, False]])
+    """
+    assert lengths.ndim == 1, lengths.ndim
+    max_len = max(max_len, lengths.max())
+    n = lengths.size(0)
+    seq_range = torch.arange(0, max_len, device=lengths.device)
+    expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
+
+    return expaned_lengths >= lengths.unsqueeze(-1)