Upload folder using huggingface_hub

README.md

@@ -8,7 +8,7 @@ sdk_version: 6.30.0
 app_file: app.py
 short_description: FocalCodec weights at int4–int32 vs fp4–fp32
 python_version: "3.12"
-startup_duration_timeout: 30m
+startup_duration_timeout: 1h
 ---
 
 # ⚡ FocalCodec 12.5 Hz — weight quantization lab
@@ -28,7 +28,7 @@ The token stream is untouched, so the codec still runs at **0.1625 kbps**
 | --- | --- |
 | **Compare** | One reconstructed clip + input/reconstruction spectrograms, side by side for all 8 levels |
 | **Playground** | Encode → decode a single clip at one chosen level, with its metrics row |
-| **Benchmark** | SI-SDR, STOI, LSD, stored model size, bitrate and real-time factor for all 8 levels, plus a summary chart |
+| **Benchmark** | SI-SDR, STOI, LSD, stored model size, bitrate and real-time factor for all 8 levels, plus an `lscodec-25hz` baseline row and a summary chart |
 
 ## Method
 
@@ -37,6 +37,10 @@ absmax scale, round-to-nearest in the target format, dequantize back to fp32, th
 the untouched FocalCodec forward pass. `fp4` is IEEE-style E2M1, `fp8` is E4M3, `fp16`
 and `fp32` are exact casts. Nothing is fine-tuned or calibrated on data.
 
+The Benchmark tab also runs [LSCodec-25 Hz](https://github.com/X-LANCE/LSCodec-Inference)
+(self-prompted with WavLM-Large) as a second codec for reference: it is a different model
+at a different bitrate (0.25 kbps vs 0.1625 kbps), so read it as a floor, not a control.
+
 ⚠️ **On SI-SDR:** it is deeply negative at every level, including `fp32`. That is
 expected — at 0.16 kbps FocalCodec does not preserve waveform phase, and the
 [FocalCodec paper](https://arxiv.org/abs/2502.04465) reports SI-SNR ≈ −37 dB for the
@@ -48,3 +52,8 @@ untouched model. Use **STOI** (intelligibility, higher is better) and **LSD**
 Model, method and example audio by
 [Luca Della Libera](https://huggingface.co/lucadellalib) et al. — Apache-2.0.
 [FocalCodec: Low-Bitrate Speech Coding via Focal Modulation Networks](https://arxiv.org/abs/2502.04465), NeurIPS 2025.
+
+The packed `int8` / `fp8` checkpoints produced by this lab live in
+[`ALWOAL/focalcodec_12_5hz_quantized`](https://huggingface.co/ALWOAL/focalcodec_12_5hz_quantized).
+The baseline codec is [LSCodec-Inference](https://github.com/X-LANCE/LSCodec-Inference) by
+Yiwei Guo et al. — MIT.

app.py

@@ -3,17 +3,22 @@
 import hashlib
 import time
 
-import spaces  # must precede torch: patches torch.cuda for ZeroGPU
-
 import gradio as gr
+import numpy as np
 import soundfile as sf
+import spaces  # patches torch.cuda for ZeroGPU; must stay above `import torch`
 import torch
+from scipy.signal import resample_poly
 
 import core
 from core import QUANT_LEVELS
 
 MODEL_REPO = "lucadellalib/focalcodec_12_5hz"
 HUB_REPO = "lucadellalib/focalcodec"
+LS_REPO = "cantabile-kwok/lscodec_25hz"
+LS_LEVEL = "lscodec-25hz"
+LS_SIZE_MB = 196.3  # encoder 70.4 + vocoder 125.6 + codebook 0.3 MB
+LS_BITRATE_KBPS = 25 * 10 / 1000  # 25 tokens/s x 10-bit codebook
 MAX_SECONDS = 20
 HEADERS = ["Level", "SI-SDR (dB)", "STOI", "LSD (dB)", "Model size (MB)", "Bitrate (kbps)", "RTF"]
 DATATYPES = ["str", "number", "number", "number", "number", "number", "number"]
@@ -32,33 +37,125 @@ SAMPLE_RATE = codec.sample_rate_input
 assert codec.sample_rate_output == SAMPLE_RATE, "demo assumes a 16/16 kHz codec"
 
 
+def _boot_lscodec() -> dict | None:
+    """LSCodec-25 Hz baseline: WavLM-Large (first 6 layers) prompt encoder + encoder/vocoder."""
+    from scipy import signal
+
+    if not hasattr(signal, "kaiser"):  # their pqmf.py imports scipy<1.15's top-level kaiser
+        from scipy.signal.windows import kaiser
+
+        signal.kaiser = kaiser
+
+    from pathlib import Path
+
+    import yaml
+    from huggingface_hub import hf_hub_download
+    from transformers import WavLMModel
+
+    from lscodec.utils import load_model, load_vocoder
+
+    files = {
+        name: hf_hub_download(LS_REPO, name, repo_type="model")
+        for name in (
+            "encoder_config.yml",
+            "vocoder_config.yml",
+            "codebook.npy",
+            "lscodec_encoder.pt",
+            "lscodec_vocoder.pt",
+        )
+    }
+    encoder_config = yaml.safe_load(Path(files["encoder_config.yml"]).read_text())
+    vocoder_config = yaml.safe_load(Path(files["vocoder_config.yml"]).read_text())
+    encoder_config["pretrain_codebook"] = files["codebook.npy"]
+    vocoder_config["vq_codebook"] = files["codebook.npy"]
+
+    enc = load_model(encoder_config, files["lscodec_encoder.pt"]).eval().requires_grad_(False)
+    dec = load_vocoder(vocoder_config, files["lscodec_vocoder.pt"]).eval().requires_grad_(False)
+    codebook = torch.as_tensor(np.load(files["codebook.npy"], allow_pickle=True))
+    books = torch.nn.ModuleList(
+        [torch.nn.Embedding.from_pretrained(b, freeze=True) for b in (codebook if codebook.ndim == 3 else [codebook])]
+    )
+    wavlm = WavLMModel.from_pretrained("microsoft/wavlm-large").eval().requires_grad_(False)
+    wavlm.encoder.layers = wavlm.encoder.layers[:6]  # prompt features come from layers 0..5
+    for model in (enc, dec, books, wavlm):
+        model.cuda()  # ZeroGPU wants module-level cuda placement (docs: model loading)
+    return {
+        "enc": enc,
+        "dec": dec,
+        "books": books,
+        "wavlm": wavlm,
+        "repeat": bool(vocoder_config.get("repeat_input_tokens")),
+        "out_sr": int(vocoder_config["sampling_rate"]),
+    }
+
+
+try:
+    LS = _boot_lscodec()
+except Exception as exc:  # noqa: BLE001  a third-party dep breaking must not take the demo down
+    LS = None
+    print(f"[lscodec] baseline disabled: {exc!r}")
+
+BENCH_LEVELS = QUANT_LEVELS + ((LS_LEVEL,) if LS is not None else ())
+
+
+def _lscodec_decode(ref: np.ndarray, sr: int) -> np.ndarray:
+    """Self-prompted LSCodec-25 Hz reconstruction, resampled back to `sr`."""
+    audio = torch.from_numpy(ref).float().view(1, -1).cuda()
+    with torch.inference_mode():
+        prompt_input = torch.nn.functional.layer_norm(audio, audio.shape)  # fairseq WavLM cfg.normalize
+        prompt = LS["wavlm"](prompt_input, output_hidden_states=True).hidden_states[6]
+        _, _, index = LS["enc"].encode(audio.unsqueeze(0))
+        if LS["repeat"]:
+            index = index.repeat_interleave(2, dim=0)
+        vqvec = torch.cat([book(index[:, i]) for i, book in enumerate(LS["books"])], dim=-1).unsqueeze(0)
+        rec = LS["dec"].inference(vqvec, prompt)[-1].view(-1).float().cpu().numpy()
+    out_sr = LS["out_sr"]
+    if out_sr != sr:
+        g = np.gcd(out_sr, sr)
+        rec = resample_poly(rec, sr // g, out_sr // g).astype(np.float32)
+    return rec
+
+
 def _estimate_duration(audio_path: str, levels: tuple) -> float:
     try:
         seconds = min(sf.info(audio_path).duration, MAX_SECONDS)
-    except Exception:
+    except Exception:  # noqa: BLE001  unreadable input falls back to the full cap
         seconds = float(MAX_SECONDS)
-    return min(240.0, 5.0 + 2.0 * len(levels) + 2.5 * seconds)
+    # ZeroGPU rejects a request above the visitor's tier cap, so never ask for more than 60 s.
+    return min(60.0, 4.0 + 2.0 * len(levels) + 2.0 * seconds + (8.0 if LS_LEVEL in levels else 0.0))
 
 
 @spaces.GPU(duration=_estimate_duration)
 def _encode_decode(audio_path: str, levels: tuple) -> dict:
     """Reconstruct `audio_path` once per weight-quantization level, scoring each result."""
     ref, sr = core.load_audio(audio_path, SAMPLE_RATE, MAX_SECONDS)
-    sig = torch.from_numpy(ref).unsqueeze(0).cuda()
-    with torch.inference_mode():  # untimed warmup so the first row's RTF isn't inflated
-        codec.sig_to_sig(sig)
+    seconds = len(ref) / sr
     rows, reconstructed = [], {}
-    for level in levels:
-        codec.load_state_dict(BASE_STATE if level == "fp32" else core.quantize_state(BASE_STATE, level))
+    focal_levels = tuple(level for level in levels if level in QUANT_LEVELS)
+    if focal_levels:
+        sig = torch.from_numpy(ref).unsqueeze(0).cuda()
+        with torch.inference_mode():  # untimed warmup so the first row's RTF isn't inflated
+            codec.sig_to_sig(sig)
+        for level in focal_levels:
+            codec.load_state_dict(BASE_STATE if level == "fp32" else core.quantize_state(BASE_STATE, level))
+            torch.cuda.synchronize()
+            start = time.perf_counter()
+            with torch.inference_mode():
+                rec = codec.sig_to_sig(sig).squeeze(0).float().cpu().numpy()
+            torch.cuda.synchronize()
+            rtf = (time.perf_counter() - start) / seconds
+            rows.append(core.evaluate(ref, rec, sr, rtf, SIZES[level], level))
+            reconstructed[level] = (sr, rec)
+        codec.load_state_dict(BASE_STATE)
+    if LS_LEVEL in levels:
+        with torch.inference_mode():  # untimed warmup, as above
+            _lscodec_decode(ref, sr)
         torch.cuda.synchronize()
         start = time.perf_counter()
-        with torch.inference_mode():
-            rec = codec.sig_to_sig(sig).squeeze(0).float().cpu().numpy()
+        rec = _lscodec_decode(ref, sr)
         torch.cuda.synchronize()
-        rtf = (time.perf_counter() - start) / (len(ref) / sr)
-        rows.append(core.evaluate(ref, rec, sr, rtf, SIZES[level], level))
-        reconstructed[level] = (sr, rec)
-    codec.load_state_dict(BASE_STATE)
+        rtf = (time.perf_counter() - start) / seconds
+        rows.append(core.evaluate(ref, rec, sr, rtf, LS_SIZE_MB, LS_LEVEL, LS_BITRATE_KBPS))
     return {"input": (sr, ref), "rows": rows, "audio": reconstructed}
 
 
@@ -67,6 +164,7 @@ def _table(rows: list[dict]) -> list[list]:
 
 
 _CACHE: dict[str, dict] = {}
+_LS_CACHE: dict[str, dict] = {}
 
 
 def run(audio_path: str, levels: list | tuple) -> dict:
@@ -109,10 +207,23 @@ def playground(audio_path: str, level: str):
     )
 
 
+def _lscodec_row(audio_path: str) -> dict:
+    """LSCodec-25 Hz score for one input, cached separately from the quantization sweep."""
+    with open(audio_path, "rb") as fh:
+        key = hashlib.md5(fh.read()).hexdigest()
+    if key not in _LS_CACHE:
+        if len(_LS_CACHE) >= 6:
+            _LS_CACHE.clear()
+        _LS_CACHE[key] = _encode_decode(audio_path, (LS_LEVEL,))["rows"][0]
+    return _LS_CACHE[key]
+
+
 def benchmark(audio_path: str):
     """Score every quantization level on one input: metrics table plus summary chart."""
-    result = run(audio_path, QUANT_LEVELS)
-    return _table(result["rows"]), core.metrics_figure(result["rows"])
+    rows = run(audio_path, QUANT_LEVELS)["rows"]
+    if LS_LEVEL in BENCH_LEVELS:
+        rows = rows + [_lscodec_row(audio_path)]
+    return _table(rows), core.metrics_figure(rows)
 
 
 HEADER = f"""# ⚡ FocalCodec 12.5 Hz — weight quantization lab
@@ -127,7 +238,9 @@ The bitrate never changes — **{core.BITRATE_KBPS:g} kbps** (12.5 tokens/s × 1
 """
 
 FOOTER = """Audio examples from the [FocalCodec repository](https://github.com/lucadellalib/focalcodec)
-(LibriSpeech, VoiceBank, MagicData). Model and method by Luca Della Libera et al., Apache-2.0."""
+(LibriSpeech, VoiceBank, MagicData). Model and method by Luca Della Libera et al., Apache-2.0.
+The **lscodec-25hz** baseline is [LSCodec](https://github.com/X-LANCE/LSCodec-Inference) by
+Yiwei Guo et al. (MIT), self-prompted with WavLM-Large from [`microsoft/wavlm-large`](https://huggingface.co/microsoft/wavlm-large)."""
 
 with gr.Blocks(title="FocalCodec quantization lab") as demo:
     gr.Markdown(HEADER)
@@ -183,13 +296,20 @@ with gr.Blocks(title="FocalCodec quantization lab") as demo:
                     sources=["upload", "microphone"], type="filepath", label="Input audio"
                 )
                 gr.Examples(examples=EXAMPLES, inputs=bm_input, label="Examples")
-                bm_button = gr.Button("Run all 8 levels", variant="primary")
+                bm_button = gr.Button(f"Run all {len(BENCH_LEVELS)} levels", variant="primary")
             with gr.Column(scale=3):
                 gr.Markdown(
                     "**STOI** ↑ intelligibility · **LSD** ↓ spectral error · "
                     "**RTF** = encode+decode time ÷ audio duration · size assumes "
                     "the stored weight format."
                 )
+                if LS_LEVEL in BENCH_LEVELS:
+                    gr.Markdown(
+                        f"**{LS_LEVEL}** is a second codec, not a weight format: it runs at "
+                        f"**{LS_BITRATE_KBPS:g} kbps** (25 tokens/s × 10 bits) versus FocalCodec's "
+                        f"**{core.BITRATE_KBPS:g} kbps**, and its size counts only the codec "
+                        f"({LS_SIZE_MB:g} MB) — the WavLM-Large prompt encoder it shares is 1.2 GB."
+                    )
         bm_table = gr.Dataframe(headers=HEADERS, datatype=DATATYPES, interactive=False)
         bm_plot = gr.Plot(label="Quality and size per level")
         bm_button.click(benchmark, inputs=bm_input, outputs=[bm_table, bm_plot])

core.py

@@ -24,6 +24,7 @@ BITRATE_KBPS = TOKEN_RATE_HZ * BITS_PER_TOKEN / 1000.0
 
 _BITS = {"int4": 4, "int8": 8, "int16": 16, "int32": 32, "fp4": 4, "fp8": 8, "fp16": 16, "fp32": 32}
 _FP_FORMATS = {"fp4": (2, 1), "fp8": (4, 3)}  # (exponent bits, mantissa bits) -> E2M1 / E4M3
+_FP_MAX_NORM = {"fp4": 6.0, "fp8": 448.0}  # E2M1 max; E4M3-FN finite max (top encoding is NaN)
 
 
 def _ref_scale(w: torch.Tensor) -> torch.Tensor:
@@ -33,24 +34,40 @@ def _ref_scale(w: torch.Tensor) -> torch.Tensor:
     return w.abs().amax()
 
 
-def _quant_int(w: torch.Tensor, bits: int) -> torch.Tensor:
-    if bits >= 32:
-        return w
+def int_scale(w: torch.Tensor, bits: int) -> torch.Tensor:
+    return (_ref_scale(w) / (1 << (bits - 1))).clamp_min(1e-12)
+
+
+def int_levels(w: torch.Tensor, bits: int) -> torch.Tensor:
+    """Signed integer levels before rescaling (what an int-format export stores)."""
     qmax = 1 << (bits - 1)
-    scale = (_ref_scale(w) / qmax).clamp_min(1e-12)
-    return (w / scale).round().clamp(-qmax, qmax - 1) * scale
+    return (w / int_scale(w, bits)).round().clamp(-qmax, qmax - 1)
+
+
+def fp_scale(w: torch.Tensor, level: str) -> torch.Tensor:
+    return (_ref_scale(w) / _FP_MAX_NORM[level]).clamp_min(1e-12)
 
 
-def _quant_fp(w: torch.Tensor, e_bits: int, m_bits: int) -> torch.Tensor:
+def fp_levels(w: torch.Tensor, level: str) -> torch.Tensor:
+    """Values on the format's own grid (what a float-format export stores, unscaled)."""
+    e_bits, m_bits = _FP_FORMATS[level]
     bias = (1 << (e_bits - 1)) - 1
     max_exp = (1 << e_bits) - 1 - bias
-    max_norm = (2.0 - 2.0 ** -m_bits) * 2.0**max_exp
-    scale = (_ref_scale(w) / max_norm).clamp_min(1e-12)
-    x = w / scale
+    x = w / fp_scale(w, level)
     ax = x.abs()
     exp = torch.floor(torch.log2(ax.clamp_min(1e-30))).clamp(1 - bias, max_exp)
     step = torch.pow(2.0, exp - m_bits)
-    return torch.sign(x) * ((ax / step).round() * step) * scale
+    return torch.sign(x) * ((ax / step).round() * step)
+
+
+def _quant_int(w: torch.Tensor, bits: int) -> torch.Tensor:
+    if bits >= 32:
+        return w
+    return int_levels(w, bits) * int_scale(w, bits)
+
+
+def _quant_fp(w: torch.Tensor, level: str) -> torch.Tensor:
+    return fp_levels(w, level) * fp_scale(w, level)
 
 
 def quantize_state(state: dict, level: str) -> dict:
@@ -66,7 +83,7 @@ def quantize_state(state: dict, level: str) -> dict:
         elif level.startswith("int"):
             out[key] = _quant_int(value, int(level[3:]))
         else:
-            out[key] = _quant_fp(value, *_FP_FORMATS[level])
+            out[key] = _quant_fp(value, level)
     return out
 
 
@@ -104,16 +121,21 @@ def log_spectral_distance(ref: np.ndarray, rec: np.ndarray, sr: int) -> float:
     return float(np.mean(np.sqrt(np.mean(diff**2, axis=0))))
 
 
-def evaluate(ref: np.ndarray, rec: np.ndarray, sr: int, rtf: float, size_mb: float, level: str) -> dict:
+def evaluate(
+    ref: np.ndarray, rec: np.ndarray, sr: int, rtf: float, size_mb: float, level: str,
+    bitrate_kbps: float | None = None,
+) -> dict:
     n = min(len(ref), len(rec))
     ref, rec = ref[:n], rec[:n]
+    if bitrate_kbps is None:
+        bitrate_kbps = BITRATE_KBPS
     return {
         "Level": level,
         "SI-SDR (dB)": round(si_sdr(ref, rec), 2),
         "STOI": round(float(stoi(ref, rec, sr, extended=False)), 3),
         "LSD (dB)": round(log_spectral_distance(ref, rec, sr), 2),
         "Model size (MB)": round(size_mb, 1),
-        "Bitrate (kbps)": round(BITRATE_KBPS, 4),
+        "Bitrate (kbps)": round(bitrate_kbps, 4),
         "RTF": round(rtf, 3),
     }
 
@@ -150,9 +172,18 @@ if __name__ == "__main__":
         q = quantize_state({"w": w}, level)["w"]
         assert torch.isfinite(q).all(), f"{level} not finite"
         assert q.abs().max() <= w.abs().max() * 1.01 + 1e-9, f"{level} broke the magnitude bound"
+    # the packed levels a real export stores must reproduce the fake-quant demo exactly
+    assert torch.equal(int_levels(w, 8) * int_scale(w, 8), quantize_state({"w": w}, "int8")["w"])
+    assert torch.equal(fp_levels(w, "fp8") * fp_scale(w, "fp8"), quantize_state({"w": w}, "fp8")["w"])
+    # fp8 grid values are exactly float8_e4m3fn-representable, so no precision is lost on store
+    grid8 = fp_levels(w, "fp8")
+    assert torch.equal(grid8.to(torch.float8_e4m3fn).float(), grid8), "fp8 grid escapes E4M3-FN"
+    grid4 = fp_levels(w, "fp4")
+    assert grid4.abs().max() <= _FP_MAX_NORM["fp4"] * 1.0000001, "fp4 grid escapes E2M1"
     x = np.random.randn(16000).astype(np.float32)
     assert si_sdr(x, x) > 60, "SI-SDR of identical signals should be large"
     assert si_sdr(x, 3 * x) > 60, "SI-SDR must be scale invariant"
     assert abs(evaluate(x, x, 16000, 1.0, 1.0, "fp32")["STOI"] - 1.0) < 1e-3
     assert evaluate(x, x, 16000, 1.0, 1.0, "fp32")["Bitrate (kbps)"] == round(BITRATE_KBPS, 4)
+    assert evaluate(x, x, 16000, 1.0, 1.0, "lscodec-25hz", 0.25)["Bitrate (kbps)"] == 0.25
     print("core self-check ok")

lscodec/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

lscodec/bin/decode_feat_prompt.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+# MIT License
+# Copyright (c) 2025 Yiwei Guo
+# This file is based on ParallelWaveGAN (https://github.com/kan-bayashi/ParallelWaveGAN). 
+# Above is the original copyright notice.
+
+
+"""Decode with Trained LSCodec Vocoder."""
+
+import argparse
+import logging
+import os
+import time
+
+import numpy as np
+import soundfile as sf
+import torch
+import yaml
+import kaldiio
+from tqdm import tqdm
+
+from lscodec.utils import load_vocoder
+
+
+def main():
+    """Run decoding process."""
+    parser = argparse.ArgumentParser(
+        description="Decode from LSCodec tokens to waveforms"
+    )
+    parser.add_argument(
+        "--feats-scp",
+        "--scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file. "
+        "you need to specify either feats-scp or dumpdir.",
+    )
+    parser.add_argument(
+        "--prompt-scp", 
+        default=None, 
+        type=str
+    )
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        required=True,
+        help="directory to save generated speech.",
+    )
+    # parser.add_argument(
+    #     "--checkpoint",
+    #     type=str,
+    #     default='pretrained/lscodec_vocoder.pt',
+    #     help="checkpoint file to be loaded.",
+    # )
+    # parser.add_argument(
+    #     "--config",
+    #     default='pretrained/vocoder_config.yml',
+    #     type=str,
+    #     help="yaml format configuration file."
+    # )
+    parser.add_argument(
+        '--pretrained-dir',
+        type=str,
+        default='pretrained',
+        help='directory of pretrained models. Should contain WavLM-Large.pt, lscodec_vocoder.pt, vocoder_config.yml',
+    )
+    parser.add_argument(
+        "--verbose",
+        type=int,
+        default=1,
+        help="logging level. higher is more logging. (default=1)",
+    )
+    args = parser.parse_args()
+
+    # set logger
+    if args.verbose > 1:
+        logging.basicConfig(
+            level=logging.DEBUG,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    elif args.verbose > 0:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    else:
+        logging.basicConfig(
+            level=logging.WARN,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+        logging.warning("Skip DEBUG/INFO messages")
+
+    # check directory existence
+    if not os.path.exists(args.outdir):
+        os.makedirs(args.outdir)
+
+    # load config
+    config_path = os.path.join(args.pretrained_dir, "vocoder_config.yml")
+    checkpoint_path = os.path.join(args.pretrained_dir, "lscodec_vocoder.pt")
+    with open(config_path) as f:
+        config = yaml.load(f, Loader=yaml.Loader)
+    config.update(vars(args))
+    config['vq_codebook'] = os.path.join(args.pretrained_dir, "codebook.npy")  # overwrite
+
+    # check arguments
+    if args.feats_scp is None:
+        raise ValueError("Please specify --feats-scp.")
+
+    # setup model
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        logging.info("Using GPU.")
+    else:
+        device = torch.device("cpu")
+        logging.info("Using CPU.")
+    model = load_vocoder(checkpoint=checkpoint_path, config=config)
+    logging.info(f"Loaded model parameters from {checkpoint_path}.")
+    model.backend.remove_weight_norm()
+    model = model.eval().to(device)
+
+    # load vq codebook
+    feat_codebook = torch.tensor(np.load(config["vq_codebook"], allow_pickle=True)).to(device)  # (V, D)
+    if feat_codebook.ndim == 2:
+        feat_codebook = feat_codebook.unsqueeze(0)
+    feat_codebook_numgroups = feat_codebook.shape[0]
+    feat_codebook = torch.nn.ModuleList([torch.nn.Embedding.from_pretrained(feat_codebook[i], freeze=True) for i in range(feat_codebook_numgroups)])
+
+    # start generation
+    total_rtf = 0.0
+    feats_scp = kaldiio.load_scp(args.feats_scp)
+    prompt_scp = kaldiio.load_scp(args.prompt_scp)
+    total_cnt = len(list(feats_scp.keys()))
+    logging.info(f"The number of features to be decoded = {total_cnt}.")
+
+    with torch.no_grad(), tqdm(feats_scp, desc="[decode]", total=total_cnt) as pbar:
+        for idx, utt_id in enumerate(pbar, 1):
+            c = feats_scp[utt_id]
+            prompt = prompt_scp[utt_id]
+
+            c = torch.tensor(c).to(device)  # (L, D)
+            prompt = torch.tensor(prompt).unsqueeze(0).to(device)  # (1, L', D')
+
+            vqidx = c.long()
+            if config.get("repeat_input_tokens", False):
+                vqidx = vqidx.repeat_interleave(2, dim=0)
+            vqvec = torch.cat([feat_codebook[i](vqidx[:, i]) for i in range(feat_codebook_numgroups)], dim=-1).unsqueeze(0)  # (1, L, D)
+
+            # generate
+            start = time.time()
+            y = model.inference(vqvec, prompt)[-1].view(-1)
+            rtf = (time.time() - start) / (len(y) / config["sampling_rate"])
+            pbar.set_postfix({"RTF": rtf})
+            total_rtf += rtf
+
+            tgt_dir = os.path.dirname(os.path.join(config["outdir"], f"{utt_id}.wav"))
+            os.makedirs(tgt_dir, exist_ok=True)
+            basename = os.path.basename(f"{utt_id}.wav")
+            # save as PCM 16 bit wav file
+            sf.write(
+                os.path.join(tgt_dir, basename),
+                y.cpu().numpy(),
+                config["sampling_rate"],
+                "PCM_16",
+            )
+
+    # report average RTF
+    logging.info(
+        f"Finished generation of {idx} utterances (RTF = {total_rtf / idx:.03f})."
+    )
+
+
+if __name__ == "__main__":
+    main()

lscodec/bin/decode_wav_prompt.py

@@ -0,0 +1,190 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+# MIT License
+# Copyright (c) 2025 Yiwei Guo
+# This file is based on ParallelWaveGAN (https://github.com/kan-bayashi/ParallelWaveGAN). 
+# Above is the original copyright notice.
+
+"""Decode with Trained LSCodec Vocoder."""
+
+import argparse
+import logging
+import os
+import time
+
+import numpy as np
+import soundfile as sf
+import torch
+import yaml
+import kaldiio
+from tqdm import tqdm
+import librosa
+
+from lscodec.utils import load_vocoder, load_audio
+from lscodec.ssl_models.wavlm_extractor import Extractor as WavLMExtractor
+
+
+def main():
+    """Run decoding process."""
+    parser = argparse.ArgumentParser(
+        description="Decode from LSCodec tokens to waveforms"
+    )
+    parser.add_argument(
+        "--feats-scp",
+        "--scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file. "
+        "you need to specify either feats-scp or dumpdir.",
+    )
+    parser.add_argument(
+        "--prompt-wav-scp", 
+        default=None, 
+        type=str,
+        help="kaldi-style prompt_wav.scp file."
+    )
+    parser.add_argument(
+        "--num-frames", 
+        default=None, 
+        type=str
+    )
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        required=True,
+        help="directory to save generated speech.",
+    )
+    parser.add_argument(
+        '--pretrained-dir',
+        type=str,
+        default='pretrained',
+        help='directory to save pretrained models. Should contain WavLM-Large.pt, lscodec_vocoder.pt, vocoder_config.yml',
+    )
+    parser.add_argument(
+        "--verbose",
+        type=int,
+        default=1,
+        help="logging level. higher is more logging. (default=1)",
+    )
+    args = parser.parse_args()
+
+    # set logger
+    if args.verbose > 1:
+        logging.basicConfig(
+            level=logging.DEBUG,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    elif args.verbose > 0:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    else:
+        logging.basicConfig(
+            level=logging.WARN,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+        logging.warning("Skip DEBUG/INFO messages")
+
+    # check directory existence
+    if not os.path.exists(args.outdir):
+        os.makedirs(args.outdir)
+
+    # load config
+    config_path = os.path.join(args.pretrained_dir, "vocoder_config.yml")
+    checkpoint_path = os.path.join(args.pretrained_dir, "lscodec_vocoder.pt")
+    wavlm_path = os.path.join(args.pretrained_dir, "WavLM-Large.pt")
+    with open(config_path) as f:
+        config = yaml.load(f, Loader=yaml.Loader)
+    config.update(vars(args))
+    config['vq_codebook'] = os.path.join(args.pretrained_dir, "codebook.npy")  # overwrite
+
+    # check arguments
+    if args.feats_scp is None:
+        raise ValueError("Please specify --feats-scp.")
+
+    # setup model
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        logging.info("Using GPU.")
+    else:
+        device = torch.device("cpu")
+        logging.info("Using CPU.")
+    model = load_vocoder(checkpoint=checkpoint_path, config=config)
+    logging.info(f"Loaded model parameters from {checkpoint_path}.")
+    model.backend.remove_weight_norm()
+    model = model.eval().to(device)
+
+    # load vq codebook
+    feat_codebook = torch.tensor(np.load(config["vq_codebook"], allow_pickle=True)).to(device)  # (V, D)
+    if feat_codebook.ndim == 2:
+        feat_codebook = feat_codebook.unsqueeze(0)
+    feat_codebook_numgroups = feat_codebook.shape[0]
+    feat_codebook = torch.nn.ModuleList([torch.nn.Embedding.from_pretrained(feat_codebook[i], freeze=True) for i in range(feat_codebook_numgroups)])
+
+    # start generation
+    total_rtf = 0.0
+    feats_scp = kaldiio.load_scp(args.feats_scp)
+    prompt_wav_scp = dict()
+    with open(args.prompt_wav_scp, 'r') as fr:
+        for line in fr.readlines():
+            terms = line.strip().split()
+            if len(terms) != 2:
+                raise ValueError(f"Invalid line in prompt_wav_scp: {line.strip()}")
+            utt_id, audio_path = terms
+            prompt_wav_scp[utt_id] = audio_path
+
+    total_cnt = len(list(feats_scp.keys()))
+    logging.info(f"The number of features to be decoded = {total_cnt}.")
+
+    wavlm = WavLMExtractor(checkpoint=wavlm_path, device=device)
+    
+    with torch.no_grad(), tqdm(feats_scp, desc="[decode]", total=total_cnt) as pbar:
+        for idx, utt_id in enumerate(pbar, 1):
+            c = feats_scp[utt_id]
+            # prompt = prompt_scp[utt_id]
+            prompt_path = prompt_wav_scp[utt_id]
+            prompt_audio, prompt_sr = load_audio(prompt_path)
+
+            if prompt_sr != 16000:
+                prompt_audio = librosa.resample(prompt_audio, orig_sr=prompt_sr, target_sr=16000)
+
+            c = torch.tensor(c).to(device)  # (L, D)
+
+            prompt = wavlm.extract(prompt_audio).unsqueeze(0)  # (1, L', D')
+
+            vqidx = c.long()
+            if config.get("repeat_input_tokens", False):
+                vqidx = vqidx.repeat_interleave(2, dim=0)
+            vqvec = torch.cat([feat_codebook[i](vqidx[:, i]) for i in range(feat_codebook_numgroups)], dim=-1).unsqueeze(0)  # (1, L, D)
+
+            # generate
+            start = time.time()
+            y = model.inference(vqvec, prompt)[-1].view(-1)
+            rtf = (time.time() - start) / (len(y) / config["sampling_rate"])
+            pbar.set_postfix({"RTF": rtf})
+            total_rtf += rtf
+
+            tgt_dir = os.path.dirname(os.path.join(config["outdir"], f"{utt_id}.wav"))
+            os.makedirs(tgt_dir, exist_ok=True)
+            basename = os.path.basename(f"{utt_id}.wav")
+            # save as PCM 16 bit wav file
+            sf.write(
+                os.path.join(tgt_dir, basename),
+                y.cpu().numpy(),
+                config["sampling_rate"],
+                "PCM_16",
+            )
+
+    # report average RTF
+    logging.info(
+        f"Finished generation of {idx} utterances (RTF = {total_rtf / idx:.03f})."
+    )
+
+
+if __name__ == "__main__":
+    main()

lscodec/bin/encode.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+# MIT License
+# Copyright (c) 2025 Yiwei Guo
+# This file is based on ParallelWaveGAN (https://github.com/kan-bayashi/ParallelWaveGAN). 
+# Above is the original copyright notice.
+
+
+"""Encode with Trained LSCodec Encoder."""
+
+import argparse
+import logging
+import os
+import time
+
+import numpy as np
+import re
+import librosa
+
+import soundfile as sf
+import torch
+from pathlib import Path
+import yaml
+
+from tqdm import tqdm
+import kaldiio
+
+from lscodec.utils import load_model, load_audio
+from kaldiio import WriteHelper
+
+
+def main():
+    """Run decoding process."""
+    parser = argparse.ArgumentParser(
+        description="Encode features with trained LSCodec"
+    )
+    parser.add_argument(
+        "--wav-scp",
+        "--scp",
+        default=None,
+        type=str,
+        help="kaldi-style wav.scp file. "
+        "you need to specify wav-scp.",
+    )
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        required=True,
+        help="directory to save generated speech.",
+    )
+    parser.add_argument(
+        "--pretrained-dir",
+        type=str,
+        default="pretrained/",
+        help="directory to save pretrained model. (default=None). It will need to contain lscodec_encoder.pt and encoder_config.yml",
+    )
+    parser.add_argument(
+        "--verbose",
+        type=int,
+        default=1,
+        help="logging level. higher is more logging. (default=1)",
+    )
+    args = parser.parse_args()
+
+    # set logger
+    if args.verbose > 1:
+        logging.basicConfig(
+            level=logging.DEBUG,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    elif args.verbose > 0:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    else:
+        logging.basicConfig(
+            level=logging.WARN,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+        logging.warning("Skip DEBUG/INFO messages")
+
+    # check directory existence
+    if not os.path.exists(args.outdir):
+        os.makedirs(args.outdir)
+
+    # load config
+    config_path = os.path.join(args.pretrained_dir, "encoder_config.yml")
+    checkpoint_path = os.path.join(args.pretrained_dir, "lscodec_encoder.pt")
+    with open(config_path) as f:
+        config = yaml.load(f, Loader=yaml.Loader)
+    config.update(vars(args))
+    config['pretrain_codebook'] = os.path.join(args.pretrained_dir, "codebook.npy")  # overwrite
+
+    # check arguments
+    if args.wav_scp is None:
+        raise ValueError("Please specify --wav-scp.")
+
+    utt2path = dict()
+    with open(args.wav_scp, 'r') as fr:
+        for line in fr.readlines():
+            terms = line.strip().split()
+            utt2path[terms[0]] = terms[1]
+    logging.info(f"The number of features to be decoded = {len(utt2path)}.")
+
+    # setup model
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        logging.info("Using GPU.")
+    else:
+        device = torch.device("cpu")
+        logging.info("Using CPU.")
+    model = load_model(config, checkpoint_path)
+    logging.info(f"Loaded model parameters from {checkpoint_path}.")
+    model = model.eval().to(device)
+
+    # start generation
+    outdir = str(Path(args.outdir).absolute())
+    with torch.no_grad(), tqdm(utt2path.items(), desc="[encode]") as pbar, WriteHelper(f"ark,scp:{outdir}/feats.ark,{outdir}/feats.scp") as writer:
+        for i, (utt_id, audio_path) in enumerate(pbar, 1):
+            audio, sr = load_audio(audio_path)
+
+            if sr != config['feature_extractor_params']['fs']:
+                audio_for_extractor = librosa.resample(audio, orig_sr=sr, target_sr=config['feature_extractor_params']['fs'])
+            else:
+                audio_for_extractor = audio
+            audio_for_extractor = torch.tensor(audio_for_extractor).float().to(device)  # (L)
+            audio_for_extractor = audio_for_extractor.unsqueeze(0).unsqueeze(0)  # (1, 1, L)
+            means, _, embed_index = model.encode(audio_for_extractor)
+
+            # now the embed_index is [L, 1] token tensor
+            if embed_index is not None:
+                writer(utt_id, embed_index.float().cpu().numpy())
+            else:
+                writer(utt_id, means.squeeze(0).cpu().numpy())
+
+
+if __name__ == "__main__":
+    main()

lscodec/bin/recon_with_prompt.py

@@ -0,0 +1,235 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+# MIT License
+# Copyright (c) 2025 Yiwei Guo
+# This file is based on ParallelWaveGAN (https://github.com/kan-bayashi/ParallelWaveGAN). 
+# Above is the original copyright notice.
+
+
+"""Decode with trained Parallel WaveGAN Generator."""
+
+import argparse
+import logging
+from kaldiio import WriteHelper
+from pathlib import Path
+import os
+import time
+import re
+
+import numpy as np
+import librosa
+
+import soundfile as sf
+import torch
+import yaml
+
+from tqdm import tqdm
+import kaldiio
+from lscodec.ssl_models.wavlm_extractor import Extractor as WavLMExtractor
+from lscodec.utils import load_model, load_vocoder, load_audio
+
+def main():
+    """Run decoding process."""
+    parser = argparse.ArgumentParser(
+        description="Decode dumped features with trained VQVAE "
+    )
+    parser.add_argument(
+        "--wav-scp",
+        "--scp",
+        default=None,
+        type=str,
+        help="kaldi-style wav.scp file. "
+        "you need to specify either wav-scp or dumpdir.",
+    )
+    parser.add_argument(
+        "--prompt-wav-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file. "
+    )
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        required=True,
+        help="directory to save generated speech.",
+    )
+    # parser.add_argument(
+    #     "--encoder-checkpoint",
+    #     type=str,
+    #     default="pretrained/lscodec_encoder.pt",
+    #     help="checkpoint file to be loaded.",
+    # )
+    # parser.add_argument(
+    #     "--encoder-config",
+    #     default='pretrained/encoder_config.yml',
+    #     type=str,
+    #     help="yaml format configuration file. if not explicitly provided, "
+    #     "it will be searched in the checkpoint directory. (default=None)",
+    # )
+    # parser.add_argument(
+    #     "--vocoder-checkpoint",
+    #     type=str,
+    #     default="pretrained/lscodec_vocoder.pt",
+    #     help="checkpoint file to be loaded.",
+    # )
+    # parser.add_argument(
+    #     "--vocoder-config",
+    #     default='pretrained/vocoder_config.yml',
+    #     type=str,
+    #     help="yaml format configuration file. if not explicitly provided, "
+    #     "it will be searched in the checkpoint directory. (default=None)",
+    # )
+    # parser.add_argument(
+    #     "--wavlm-path",
+    #     type=str,
+    #     default="pretrained/WavLM-Large.pt",
+    #     help="path to WavLM checkpoint.",
+    # )
+    parser.add_argument(
+        "--pretrained-dir",
+        type=str,
+        default="pretrained/",
+        help="directory to save pretrained model. (default=None). It will need to contain WavLM-Large.pt, lscodec_encoder.pt, lscodec_vocoder.pt, encoder_config.yml, vocoder_config.yml and codebook.npy",
+    )
+    parser.add_argument(
+        "--verbose",
+        type=int,
+        default=1,
+        help="logging level. higher is more logging. (default=1)",
+    )
+    args = parser.parse_args()
+
+    # set logger
+    if args.verbose > 1:
+        logging.basicConfig(
+            level=logging.DEBUG,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    elif args.verbose > 0:
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    else:
+        logging.basicConfig(
+            level=logging.WARN,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+        logging.warning("Skip DEBUG/INFO messages")
+
+    # check directory existence
+    if not os.path.exists(args.outdir):
+        os.makedirs(args.outdir)
+
+    # load encoder config
+    encoder_config_path = os.path.join(args.pretrained_dir, "encoder_config.yml")
+    encoder_checkpoint_path = os.path.join(args.pretrained_dir, "lscodec_encoder.pt")
+    with open(encoder_config_path) as f:
+        encoder_config = yaml.load(f, Loader=yaml.Loader)
+    encoder_config.update(vars(args))
+    encoder_config['pretrain_codebook'] = os.path.join(args.pretrained_dir, "codebook.npy")  # overwrite
+    
+    # load vocoder config
+    vocoder_config_path = os.path.join(args.pretrained_dir, "vocoder_config.yml")
+    vocoder_checkpoint_path = os.path.join(args.pretrained_dir, "lscodec_vocoder.pt")
+    with open(vocoder_config_path) as f:
+        vocoder_config = yaml.load(f, Loader=yaml.Loader)
+    vocoder_config.update(vars(args))
+    vocoder_config['vq_codebook'] = os.path.join(args.pretrained_dir, "codebook.npy")  # overwrite
+
+    # check arguments
+    if args.wav_scp is None or args.prompt_wav_scp is None:
+        raise ValueError("Please specify --wav-scp and --prompt-wav-scp.")
+
+    utt2path = dict()
+    with open(args.wav_scp, 'r') as fr:
+        for line in fr.readlines():
+            terms = line.strip().split()
+            if len(terms) != 2:
+                raise ValueError(f"Invalid line in wav_scp: {line.strip()}")
+            utt2path[terms[0]] = terms[1]
+    prompt_scp = dict()
+    with open(args.prompt_wav_scp, 'r') as fr:
+        for line in fr.readlines():
+            terms = line.strip().split()
+            if len(terms) != 2:
+                raise ValueError(f"Invalid line in prompt_wav_scp: {line.strip()}")
+            utt_id, audio_path = terms
+            prompt_scp[utt_id] = audio_path
+
+    logging.info(f"The number of features to be decoded = {len(utt2path)}.")
+
+    # setup model
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        logging.info("Using GPU.")
+    else:
+        device = torch.device("cpu")
+        logging.info("Using CPU.")
+    model = load_model(encoder_config, encoder_checkpoint_path)
+    logging.info(f"Loaded encoder parameters from {encoder_checkpoint_path}.")
+    model = model.eval().to(device)
+
+    vocoder = load_vocoder(vocoder_config, vocoder_checkpoint_path)
+    logging.info(f"Loaded vocoder parameters from {vocoder_checkpoint_path}.")
+    vocoder = vocoder.eval().to(device)
+    # load vq codebook
+    feat_codebook = torch.tensor(np.load(vocoder_config["vq_codebook"], allow_pickle=True)).to(device)  # (V, D)
+    if feat_codebook.ndim == 2:
+        feat_codebook = feat_codebook.unsqueeze(0)
+    feat_codebook_numgroups = feat_codebook.shape[0]
+    feat_codebook = torch.nn.ModuleList([torch.nn.Embedding.from_pretrained(feat_codebook[i], freeze=True) for i in range(feat_codebook_numgroups)])
+
+    wavlm_path = os.path.join(args.pretrained_dir, "WavLM-Large.pt")
+    wavlm = WavLMExtractor(checkpoint=wavlm_path, device=device)
+
+    # start generation
+    with torch.no_grad(), tqdm(utt2path.items(), desc="[prompted recon]") as pbar:
+        for idx, (utt_id, audio_path) in enumerate(pbar, 1):
+            prompt_path = prompt_scp[utt_id]
+            audio, sr = load_audio(audio_path)
+            prompt_audio, prompt_sr = load_audio(prompt_path)
+
+            if sr != encoder_config['feature_extractor_params']['fs']:
+                audio = librosa.resample(audio, orig_sr=sr, target_sr=encoder_config['feature_extractor_params']['fs'])
+
+            if prompt_sr != 16000:
+                prompt_audio = librosa.resample(prompt_audio, orig_sr=prompt_sr, target_sr=16000)
+
+            audio = torch.tensor(audio).float().to(device)  # (L)
+            audio = audio.unsqueeze(0).unsqueeze(0)  # (1, 1, L)
+
+            prompt = wavlm.extract(prompt_audio)
+            prompt = prompt.float().unsqueeze(0).to(device) # [1, L, D]
+
+            # generate
+            _, _, embed_index = model.encode(audio)
+            if vocoder_config.get("repeat_input_tokens", False):
+                # in LSCodec-25Hz, we still trained the vocoder on 50Hz input, so repeat_interleave is needed.
+                embed_index = embed_index.repeat_interleave(2, dim=0)
+            vqvec = torch.cat([feat_codebook[i](embed_index[:, i]) for i in range(feat_codebook_numgroups)], dim=-1).unsqueeze(0)  # (1, L, D)
+            y = vocoder.inference(vqvec, prompt)[-1].view(-1)
+
+            tgt_dir = os.path.dirname(os.path.join(args.outdir, f"{utt_id}.wav"))
+            os.makedirs(tgt_dir, exist_ok=True)
+            basename = os.path.basename(f"{utt_id}.wav")
+            # save as PCM 16 bit wav file
+            sf.write(
+                os.path.join(tgt_dir, basename),
+                y.cpu().numpy(),
+                vocoder_config["sampling_rate"],
+                "PCM_16",
+            )
+
+    # report average RTF
+    logging.info(
+        f"Finished generation of {idx} utterances."
+    )
+
+
+if __name__ == "__main__":
+    main()

lscodec/bin/train.py

@@ -0,0 +1,994 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+# Portions of this code are a modification of the ParallelWaveGAN project.
+# The original source code can be found at:
+# https://github.com/kan-bayashi/ParallelWaveGAN/blob/master/parallel_wavegan/bin/train.py
+#
+# Copyright (c) 2025 Yiwei Guo
+# Licensed under the MIT license.
+
+import argparse
+import logging
+import os
+import sys
+import random
+
+from collections import defaultdict
+
+import matplotlib
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+import yaml
+import torch.multiprocessing as mp
+from tensorboardX import SummaryWriter
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+import lscodec
+import lscodec.models
+from lscodec.models.quantization.core_vq import GroupVectorQuantization
+import lscodec.optimizers
+from torch.utils.data.distributed import DistributedSampler
+
+from lscodec.datasets.scp_dataset import AudioMelSCPDataset
+from lscodec.utils.espnet_utils import pad_list, make_non_pad_mask
+from lscodec.datasets.perturbs import Perturbs
+
+# set to avoid matplotlib error in CLI environment
+matplotlib.use("Agg")
+
+
+def KL_loss(means, logvars, mask):
+    means = means.masked_select(mask)
+    logvars = logvars.masked_select(mask)
+    logvars = torch.clamp(logvars, -3, 3)  # NOTE: added to ensure stability
+    loss = -0.5*logvars + 0.5*(logvars.exp() + means.pow(2)) - 0.5
+    # if torch.isnan(loss).any():
+        # logging.info("detected nan in kl_loss")
+        # raise ValueError()
+    # loss = torch.nan_to_num(loss, nan=10)
+    return loss.mean()
+
+
+class Trainer(object):
+    """Customized trainer module for Parallel WaveGAN training."""
+
+    def __init__(
+            self,
+            steps,
+            epochs,
+            data_loader,
+            sampler,
+            model,
+            optimizer,
+            scheduler,
+            config,
+            device=torch.device("cpu"),
+    ):
+        """Initialize trainer.
+
+        Args:
+            steps (int): Initial global steps.
+            epochs (int): Initial global epochs.
+            data_loader (dict): Dict of data loaders. It must contain "train" and "dev" loaders.
+            model (dict): Dict of models. It must contain "generator" model.
+            optimizer (dict): Dict of optimizers. It must contain "generator" and optimizer.
+            scheduler (dict): Dict of schedulers. It must contain "generator" and schedulers.
+            config (dict): Config dict loaded from yaml format configuration file.
+            device (torch.deive): Pytorch device instance.
+
+        """
+        self.steps = steps
+        self.epochs = epochs
+        self.data_loader = data_loader
+        self.sampler = sampler
+        self.model = model
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.config = config
+        self.device = device
+        self.writer = SummaryWriter(config["outdir"])
+        self.finish_train = False
+        self.total_train_loss = defaultdict(float)
+        self.total_eval_loss = defaultdict(float)
+
+    def run(self):
+        """Run training."""
+        self.tqdm = tqdm(
+            initial=self.steps, total=self.config["train_max_steps"], desc="[train]"
+        )
+        while True:
+            # train one epoch
+            self._train_epoch()
+
+            # check whether training is finished
+            if self.finish_train:
+                break
+
+        self.tqdm.close()
+        logging.info("Finished training.")
+
+    def save_checkpoint(self, checkpoint_path):
+        """Save checkpoint.
+
+        Args:
+            checkpoint_path (str): Checkpoint path to be saved.
+
+        """
+        state_dict = {
+            "optimizer": {
+                "generator": self.optimizer["generator"].state_dict(),
+            },
+            "scheduler": {
+                "generator": self.scheduler["generator"].state_dict(),
+            },
+            "steps": self.steps,
+            "epochs": self.epochs,
+        }
+        if self.config["distributed"]:
+            state_dict["model"] = {
+                "generator": self.model["generator"].module.state_dict(),
+            }
+        else:
+            state_dict["model"] = {
+                "generator": self.model["generator"].state_dict(),
+            }
+
+        if not os.path.exists(os.path.dirname(checkpoint_path)):
+            os.makedirs(os.path.dirname(checkpoint_path))
+        torch.save(state_dict, checkpoint_path)
+
+    def load_checkpoint(self, checkpoint_path, load_only_params=False):
+        """Load checkpoint.
+
+        Args:
+            checkpoint_path (str): Checkpoint path to be loaded.
+            load_only_params (bool): Whether to load only model parameters.
+
+        """
+        state_dict = torch.load(checkpoint_path, map_location="cpu")
+        if self.config["distributed"]:
+            self.model["generator"].module.load_state_dict(
+                state_dict["model"]["generator"]
+            )
+        else:
+            self.model["generator"].load_state_dict(state_dict["model"]["generator"])
+        if not load_only_params:
+            self.steps = state_dict["steps"]
+            self.epochs = state_dict["epochs"]
+            self.optimizer["generator"].load_state_dict(
+                state_dict["optimizer"]["generator"]
+            )
+            self.scheduler["generator"].load_state_dict(
+                state_dict["scheduler"]["generator"]
+            )
+
+    def load_vae_params(self, checkpoint_path):
+        state_dict = torch.load(checkpoint_path, map_location="cpu")
+        generator = state_dict['model']['generator']
+        frontend_state_dict = {}
+        backend_state_dict = {}
+        for key in generator.keys():
+            if key.startswith("frontend"):
+                frontend_state_dict[".".join(key.split(".")[1:])] = generator[key]
+            elif key.startswith("backend"):
+                backend_state_dict[".".join(key.split(".")[1:])] = generator[key]
+            else:
+                raise NotImplementedError(f"Unknown key: {key}")
+        if self.config["distributed"]:
+            self.model['generator'].module.frontend.load_state_dict(frontend_state_dict)
+            self.model['generator'].module.backend.load_state_dict(backend_state_dict)
+        else:
+            self.model['generator'].frontend.load_state_dict(frontend_state_dict)
+            self.model['generator'].backend.load_state_dict(backend_state_dict)
+
+    def open_codebook_expire(self):
+        if self.config['distributed']:
+            for group_i in range(len(self.model['generator'].module.quantizer.groups)):
+                self.model['generator'].module.quantizer.groups[group_i]._codebook.perform_expire = True
+        else:
+            for group_i in range(len(self.model['generator'].quantizer.groups)):
+                self.model['generator'].quantizer.groups[group_i]._codebook.perform_expire = True
+
+    def open_codebook_ema(self):
+        if self.config['distributed']:
+            for group_i in range(len(self.model['generator'].module.quantizer.groups)):
+                self.model['generator'].module.quantizer.groups[group_i]._codebook.perform_ema_update = True
+        else:
+            for group_i in range(len(self.model['generator'].quantizer.groups)):
+                self.model['generator'].quantizer.groups[group_i]._codebook.perform_ema_update = True
+
+    def _train_step(self, batch):
+        """Train model one step."""
+        if self.config['have_quantizer'] and self.config['quantizer_class'] == "KMeans":
+            if self.steps > self.config['quantizer_params']['perform_expire_after_steps']:
+                self.open_codebook_expire()
+            if self.steps > self.config['quantizer_params']['perform_ema_after_steps']:
+                self.open_codebook_ema()
+        if self.config['have_quantizer'] and self.config['quantizer_class'] == "Gumbel":
+            self.model['generator'].module.quantizer.set_num_updates(self.steps)
+
+        # parse batch
+        aux, target, idx, prompt, y_for_extractor, xlens, prompt_lens = batch
+        target = target.to(self.device)
+        prompt = prompt.to(self.device)
+        idx = idx.to(self.device)
+        y_for_extractor = y_for_extractor.unsqueeze(-2).to(self.device)  # (B, 1, T')
+
+        # build mask
+        mask = make_non_pad_mask(xlens).to(self.device)  # (B, L) True for valid values
+        prompt_mask = make_non_pad_mask(prompt_lens).to(self.device)  # (B, L_prompt)
+
+        #######################
+        #      Generator      #
+        #######################
+        if self.config['distributed']:
+            tolerance = self.model['generator'].module.length_mismatch_tolerance
+        else:
+            tolerance = self.model['generator'].length_mismatch_tolerance
+        target_, idx_logits_, _, commitment_loss, mask, means, logvars = self.model["generator"](y_for_extractor, prompt, mask, prompt_mask)
+        # (B, L, 80)
+
+        if target.shape[1] != mask.shape[-1]:
+            if abs(target.shape[1] - mask.shape[-1]) <= tolerance:
+                min_len = min(target.shape[1], mask.shape[-1])
+                target = target[:, :min_len, :]
+                mask = mask[..., :min_len]
+                target_ = target_[:, :min_len, :]
+                idx = idx[:, :min_len, :]
+                idx_logits_ = idx_logits_[:, :min_len, :]
+
+        # initialize
+        gen_loss = 0.0
+        # frontend target prediction loss
+        frontend_target_pred_loss = F.l1_loss(torch.masked_select(target, mask.unsqueeze(-1)),
+                                            torch.masked_select(target_, mask.unsqueeze(-1)))
+        self.total_train_loss["train/frontend_target_pred_loss"] += frontend_target_pred_loss.item()
+        gen_loss += self.config["lambda_frontend_target_prediction"] * frontend_target_pred_loss
+        
+        idx_criterion = torch.nn.CrossEntropyLoss()
+        V = idx_logits_.shape[-1]
+        frontend_idx_pred_loss = idx_criterion(torch.masked_select(idx_logits_, mask.unsqueeze(-1)).view(-1, V),
+                                               torch.masked_select(idx.squeeze(-1), mask))
+        self.total_train_loss["train/frontend_idx_pred_loss"] += frontend_idx_pred_loss.item()
+        gen_loss += self.config["lambda_frontend_idx_prediction"] * frontend_idx_pred_loss
+
+        if commitment_loss is None:
+            if self.config['quantizer_class'] != "Gumbel":  # In this case, no additional loss.
+                # VQVAE KL loss
+                kl_loss = KL_loss(means, logvars, mask.unsqueeze(-1))
+                gen_loss += self.config['lambda_kl_loss'] * kl_loss
+                self.total_train_loss['train/kl_loss'] += kl_loss.item()
+        else:
+            # VQ loss (if kmeans)
+            if self.config['quantizer_class'] == "KMeans":
+                gen_loss += self.config['lambda_kmeans_loss'] * commitment_loss
+                self.total_train_loss['train/kmeans_loss'] += commitment_loss.item()
+
+        self.total_train_loss["train/generator_loss"] += gen_loss.item()
+
+        # update generator
+        self.optimizer["generator"].zero_grad()
+        # gen_loss = torch.nan_to_num(gen_loss, nan=100)  # NOTE: will this ensure stability?
+        if torch.isnan(gen_loss):
+            logging.info(f"Find NAN in loss. Refusing to update.")
+            gen_loss = 0.
+            for p in self.model['generator'].parameters():
+                gen_loss = gen_loss + 0.0 * p.sum()
+            # gen_loss = torch.sum([0.0*p.sum() for p in self.model['generator'].parameters()])  # fake loss. Gradient=0
+
+        gen_loss.backward()
+        if self.config["generator_grad_norm"] > 0:
+            torch.nn.utils.clip_grad_norm_(
+                self.model["generator"].parameters(),
+                self.config["generator_grad_norm"],
+            )
+        self.optimizer["generator"].step()
+        self.scheduler["generator"].step()
+        # update counts
+        self.steps += 1
+        self.tqdm.update(1)
+        self._check_train_finish()
+
+    def _train_epoch(self):
+        """Train model one epoch."""
+        for train_steps_per_epoch, batch in enumerate(self.data_loader["train"], 1):
+            # train one step
+            self._train_step(batch)
+
+            # check interval
+            if self.config["rank"] == 0:
+                self._check_log_interval()
+                self._check_eval_interval()
+                self._check_save_interval()
+
+            # check whether training is finished
+            if self.finish_train:
+                return
+
+        # update
+        self.epochs += 1
+        self.train_steps_per_epoch = train_steps_per_epoch
+        logging.info(
+            f"(Steps: {self.steps}) Finished {self.epochs} epoch training "
+            f"({self.train_steps_per_epoch} steps per epoch)."
+        )
+
+        # needed for shuffle in distributed training
+        if self.config["distributed"]:
+            self.sampler["train"].set_epoch(self.epochs)
+
+    @torch.no_grad()
+    def _eval_step(self, batch):
+        """Evaluate model one step."""
+        # parse batch
+        aux, target, idx, prompt, y_for_extractor, xlens, prompt_lens = batch
+        target = target.to(self.device)
+        prompt = prompt.to(self.device)
+        idx = idx.to(self.device)
+        y_for_extractor = y_for_extractor.unsqueeze(-2).to(self.device)  # (B, 1, T')
+
+        # build mask
+        mask = make_non_pad_mask(xlens).to(self.device)  # (B, L)
+        prompt_mask = make_non_pad_mask(prompt_lens).to(self.device)  # (B, L_prompt)
+
+        #######################
+        #      Generator      #
+        #######################
+        if self.config['distributed']:
+            tolerance = self.model['generator'].module.length_mismatch_tolerance
+        else:
+            tolerance = self.model['generator'].length_mismatch_tolerance
+        target_, idx_logits_, _, commitment_loss, mask, means, logvars = self.model["generator"](y_for_extractor, prompt, mask, prompt_mask)  # (B, L, 80), (B, C, T)
+        if target.shape[1] != mask.shape[-1]:
+            if abs(target.shape[1] - mask.shape[-1]) <= tolerance:
+                min_len = min(target.shape[1], mask.shape[-1])
+                target = target[:, :min_len, :]
+                mask = mask[..., :min_len]
+                target_ = target_[:, :min_len, :]
+                idx = idx[:, :min_len, :]
+                idx_logits_ = idx_logits_[:, :min_len, :]
+
+        # frontend target prediction loss
+        frontend_target_pred_loss = F.l1_loss(torch.masked_select(target, mask.unsqueeze(-1)),
+                                            torch.masked_select(target_, mask.unsqueeze(-1)))
+        self.total_eval_loss["eval/frontend_target_pred_loss"] += frontend_target_pred_loss.item()
+        
+        idx_criterion = torch.nn.CrossEntropyLoss()
+        V = idx_logits_.shape[-1]
+        frontend_idx_pred_loss = idx_criterion(torch.masked_select(idx_logits_, mask.unsqueeze(-1)).view(-1, V),
+                                               torch.masked_select(idx.squeeze(-1), mask))
+        self.total_eval_loss["eval/frontend_idx_pred_loss"] += frontend_idx_pred_loss.item()
+
+        if commitment_loss is None:
+            if self.config['quantizer_class'] != "Gumbel":  # In this case, no additional loss.
+                # VAE KL loss
+                kl_loss = KL_loss(means, logvars, mask.unsqueeze(-1))
+                self.total_eval_loss['eval/kl_loss'] += kl_loss.item()
+        else:
+            # VQ loss (if kmeans)
+            if self.config['quantizer_class'] == "KMeans":
+                self.total_eval_loss['eval/kmeans_loss'] += commitment_loss.item()
+
+    def _eval_epoch(self):
+        """Evaluate model one epoch."""
+        logging.info(f"(Steps: {self.steps}) Start evaluation.")
+        # change mode
+        for key in self.model.keys():
+            self.model[key].eval()
+
+        # calculate loss for each batch
+        for eval_steps_per_epoch, batch in enumerate(
+                tqdm(self.data_loader["dev"], desc="[eval]"), 1
+        ):
+            # eval one step
+            self._eval_step(batch)
+
+        logging.info(
+            f"(Steps: {self.steps}) Finished evaluation "
+            f"({eval_steps_per_epoch} steps per epoch)."
+        )
+
+        # average loss
+        for key in self.total_eval_loss.keys():
+            self.total_eval_loss[key] /= eval_steps_per_epoch
+            logging.info(
+                f"(Steps: {self.steps}) {key} = {self.total_eval_loss[key]:.4f}."
+            )
+
+        # record
+        self._write_to_tensorboard(self.total_eval_loss)
+
+        # reset
+        self.total_eval_loss = defaultdict(float)
+
+        # restore mode
+        for key in self.model.keys():
+            self.model[key].train()
+
+    def _write_to_tensorboard(self, loss):
+        """Write to tensorboard."""
+        for key, value in loss.items():
+            self.writer.add_scalar(key, value, self.steps)
+
+    def _check_save_interval(self):
+        if self.steps % self.config["save_interval_steps"] == 0:
+            self.save_checkpoint(
+                os.path.join(self.config["outdir"], f"checkpoint-{self.steps}steps.pkl")
+            )
+            logging.info(f"Successfully saved checkpoint @ {self.steps} steps.")
+
+    def _check_eval_interval(self):
+        if self.steps % self.config["eval_interval_steps"] == 0:
+            self._eval_epoch()
+
+    def _check_log_interval(self):
+        if self.steps % self.config["log_interval_steps"] == 0:
+            for key in self.total_train_loss.keys():
+                self.total_train_loss[key] /= self.config["log_interval_steps"]
+                logging.info(
+                    f"(Steps: {self.steps}) {key} = {self.total_train_loss[key]:.4f}."
+                )
+            if self.config['have_quantizer'] and self.config['quantizer_class'] == "Gumbel":
+                logging.info(f"(Steps: {self.steps}) gumbel temperature = {self.model['generator'].module.quantizer.curr_temp}")
+            self._write_to_tensorboard(self.total_train_loss)
+
+            # reset
+            self.total_train_loss = defaultdict(float)
+
+    def _check_train_finish(self):
+        if self.steps >= self.config["train_max_steps"]:
+            self.finish_train = True
+
+
+class Collator(object):
+    """Customized collator for Pytorch DataLoader in training."""
+
+    def __init__(
+            self,
+            hop_size_for_extractor=256,
+            win_length_for_extractor=1024,
+            sampling_rate=16000,
+            n_mel=80,
+            force_from_start=False,
+    ):
+        """Initialize customized collator for PyTorch DataLoader.
+
+        Args:
+            hop_size_for_extractor (int): Hop size of features, in sampling points.
+            win_length_for_extractor (int): window length of features.
+            sampling_rate (int): sampling rate of waveform data
+            n_mel (int): number of mel spectrogram dimensions
+        """
+        self.hop_size_for_extractor = hop_size_for_extractor
+        self.win_length_for_extractor = win_length_for_extractor
+        self.sampling_rate = sampling_rate
+        self.n_mel = n_mel
+        self.force_from_start = force_from_start
+
+    def construct_prompt(self, mel_lens):
+        prompt_lens = [random.randint(int(l / 3), int(l / 2)) for l in mel_lens]
+        prompt_starts = []
+        is_from_start = []
+        for ml, pl in zip(mel_lens, prompt_lens):
+            random_number = random.random()
+            if self.force_from_start:
+                random_number = 1
+            if random_number > 0.5:
+                # from start
+                prompt_start = random.randint(0, 1 * self.sampling_rate//self.hop_size_for_extractor)
+                is_from_start.append(True)
+            else:
+                # from ending
+                prompt_start = random.randint((ml - 1 * self.sampling_rate//self.hop_size_for_extractor), ml) - pl
+                is_from_start.append(False)
+            prompt_starts.append(prompt_start)
+        return prompt_lens, prompt_starts, is_from_start
+
+    def __call__(self, batch):
+        """Convert into batch tensors.
+
+        Args:
+            batch (list): list of tuple of the pair of audio and features.
+
+        This collator will automatically determine the prompt segment (acoustic context) for each utterance.
+        The prompt is cut off from the current utterance, ranging from 2 to 3 seconds.
+        The starting point of the prompt lies between the first 0-1 second range.
+        Then, it parses the concatenated features into (3 dim auxiliary features, 2 dim VQ features, and 80 dim mel spectrograms)
+
+        Returns:
+            Tensor: Gaussian noise batch (B, 1, T).
+            Tensor: Auxiliary feature batch (B, C, T'), where
+                T = T' * hop_size.
+            Tensor: Target signal batch (B, 1, T).
+
+        """
+        batch = batch[0]
+
+        # check length
+        batch = [
+            self._adjust_length(*b) for b in batch
+        ]
+        ys_for_extractor, mels, auxs, targets, idxs = ([b[0] for b in batch], [b[1] for b in batch], [b[2] for b in batch], [b[3] for b in batch], [b[4] for b in batch])
+
+        batch_size = len(ys_for_extractor)
+        prompt_lengths, prompt_starts, is_from_starts = self.construct_prompt([len(m) for m in mels])
+
+        c_lengths = []
+
+        prompts = torch.zeros(batch_size, max(prompt_lengths), self.n_mel)
+        for i in range(batch_size):
+            # data_offset = prompt_offsets[i] + prompt_lengths[i]
+            prompts[i, :prompt_lengths[i]] = torch.tensor(mels[i][prompt_starts[i]:prompt_starts[i]+prompt_lengths[i], :])
+            if is_from_starts[i]:
+                targets[i] = targets[i][prompt_starts[i]+prompt_lengths[i]:]
+                idxs[i] = idxs[i][prompt_starts[i]+prompt_lengths[i]:]
+                auxs[i] = auxs[i][prompt_starts[i]+prompt_lengths[i]:]
+                ys_for_extractor[i] = ys_for_extractor[i][(prompt_starts[i]+prompt_lengths[i])*self.hop_size_for_extractor:]
+            else:
+                targets[i] = targets[i][:prompt_starts[i]]
+                idxs[i] = idxs[i][:prompt_starts[i]]
+                auxs[i] = auxs[i][:prompt_starts[i]]
+                ys_for_extractor[i] = ys_for_extractor[i][:prompt_starts[i]*self.hop_size_for_extractor]
+            c_lengths.append(len(targets[i]))
+        # NOTE: PAD_VALUES is set to 999 instead of 0, so that the code vectors to be quantized are not affected.
+        targets = pad_list([torch.tensor(c) for c in targets], pad_value=0)  # (B, L, 80)
+        idxs = pad_list([torch.tensor(c) for c in idxs], pad_value=0)  # (B, L, 80)
+        idxs = idxs.long()
+        auxs = pad_list([torch.tensor(c) for c in auxs], pad_value=0)  # (B, L, 3)
+        ys_for_extractor = pad_list([torch.tensor(y, dtype=torch.float) for y in ys_for_extractor],
+                                    pad_value=0)[:, :targets.size(1) * self.hop_size_for_extractor]
+        assert targets.size(1) == auxs.size(1) == idxs.size(1)
+        assert ys_for_extractor.size(1) == targets.size(1) * self.hop_size_for_extractor == auxs.size(1) * self.hop_size_for_extractor
+
+        return auxs, targets, idxs, prompts, ys_for_extractor, c_lengths, prompt_lengths
+
+    def _adjust_length(self, x, c, *args):
+        """Adjust the audio and feature lengths.
+
+        Note:
+            Basically we assume that the length of x and c are adjusted
+            through preprocessing stage, but if we use other library processed
+            features, this process will be needed.
+
+        """
+        if len(x) > len(c) * self.hop_size_for_extractor:
+            x = x[(self.win_length_for_extractor - self.hop_size_for_extractor) // 2:]
+            x = x[:len(c) * self.hop_size_for_extractor]
+
+        # check the legnth is valid
+        assert len(x) == len(c) * self.hop_size_for_extractor, f"{len(x)}, {len(c)}, {self.hop_size_for_extractor}"
+
+        return x, c, *args
+
+
+def main(rank, n_gpus):
+    """Run training process."""
+    parser = argparse.ArgumentParser(
+        description="Train Parallel WaveGAN (See detail in ctx_vec2wav/bin/train.py)."
+    )
+    parser.add_argument(
+        "--train-wav-scp",
+        default=None,
+        type=str,
+        help="kaldi-style wav.scp file for training. "
+    )
+    parser.add_argument(
+        "--train-idx-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file for training. "
+    )
+    parser.add_argument(
+        "--train-prompt-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file for training. "
+    )
+    parser.add_argument(
+        "--train-target-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file for training. "
+    )
+    parser.add_argument(
+        "--train-aux-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file for training. "
+    )
+    parser.add_argument(
+        "--train-segments",
+        default=None,
+        type=str,
+        help="kaldi-style segments file for training.",
+    )
+    # parser.add_argument(
+    #     "--train-xvector-scp",
+    #     default=None,
+    #     type=str,
+    #     help="kaldi-style xvector.scp file for training.",
+    # )
+    parser.add_argument(
+        "--train-num-frames",
+        default=None,
+        type=str,
+        help="kaldi-style utt2num_frames file for training.",
+    )
+    parser.add_argument(
+        "--train-utt2spk",
+        default=None,
+        type=str,
+        help="kaldi-style utt2spk file for training.",
+    )
+    parser.add_argument(
+        "--train-spk2gender",
+        default=None,
+        type=str,
+        help="kaldi-style spk2gender file for training.",
+    )
+    parser.add_argument(
+        "--dev-wav-scp",
+        default=None,
+        type=str,
+        help="kaldi-style wav.scp file for validation. "
+    )
+    parser.add_argument(
+        "--dev-idx-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file for vaidation. "
+    )
+    parser.add_argument(
+        "--dev-prompt-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file for vaidation. "
+    )
+    parser.add_argument(
+        "--dev-target-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file for vaidation. "
+    )
+    parser.add_argument(
+        "--dev-aux-scp",
+        default=None,
+        type=str,
+        help="kaldi-style feats.scp file for vaidation. "
+    )
+    parser.add_argument(
+        "--dev-segments",
+        default=None,
+        type=str,
+        help="kaldi-style segments file for validation.",
+    )
+    # parser.add_argument(
+    #     "--dev-xvector-scp",
+    #     default=None,
+    #     type=str,
+    #     help="kaldi-style xvector.scp file for validation.",
+    # )
+    parser.add_argument(
+        "--dev-num-frames",
+        default=None,
+        type=str,
+        help="kaldi-style utt2num_frames file for validation.",
+    )
+    parser.add_argument(
+        "--dev-utt2spk",
+        default=None,
+        type=str,
+        help="kaldi-style utt2spk file for training.",
+    )
+    parser.add_argument(
+        "--dev-spk2gender",
+        default=None,
+        type=str,
+        help="kaldi-style spk2gender file for training.",
+    )
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        required=True,
+        help="directory to save checkpoints.",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        required=True,
+        help="yaml format configuration file.",
+    )
+    parser.add_argument(
+        "--pretrain",
+        default="",
+        type=str,
+        nargs="?",
+        help='checkpoint file path to load pretrained params. (default="")',
+    )
+    # parser.add_argument("--pretrain-feature-extractor",
+    #                     default="",
+    #                     type=str,
+    #                     help="checkpoint file path to load pretrained params for CNN+trans feature extractor. Default=''")
+    parser.add_argument("--pretrain-vae",
+                        default="",
+                        type=str,
+                        help="pretrained VAE (no VQ)")
+    parser.add_argument("--pretrain-codebook",
+                        default="",
+                        type=str,
+                        help="checkpoint file path to load pretrained params for codebook. Default=''")
+    # parser.add_argument("--pretrain-ctx-vec2wav",
+    #                     default="",
+    #                     type=str,
+    #                     help="checkpoint file path to load pretrained params for ctx-vec2wav. Default=''")
+    parser.add_argument(
+        "--resume",
+        default="",
+        type=str,
+        nargs="?",
+        help='checkpoint file path to resume training. (default="")',
+    )
+    parser.add_argument(
+        "--verbose",
+        type=int,
+        default=1,
+        help="logging level. higher is more logging. (default=1)",
+    )
+    # parser.add_argument("--vq-codebook", default=None, type=str)
+    parser.add_argument("--sampling-rate", type=int)
+    parser.add_argument("--num-mels", type=int)
+    parser.add_argument("--hop-size", type=int)
+    parser.add_argument("--win-length", type=int)
+    args = parser.parse_args()
+
+    # init distributed training
+    device = torch.device("cuda")
+    # effective when using fixed size inputs
+    # see https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936
+    torch.backends.cudnn.benchmark = True
+    # setup for distributed training
+    # see example: https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed
+    if n_gpus == 1:
+        assert rank == 0
+
+    # set logger
+    if args.verbose > 1:
+        logging.basicConfig(
+            level=logging.DEBUG,
+            stream=sys.stdout,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    elif args.verbose > 0:
+        logging.basicConfig(
+            level=logging.INFO,
+            stream=sys.stdout,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+    else:
+        logging.basicConfig(
+            level=logging.WARN,
+            stream=sys.stdout,
+            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
+        )
+        logging.warning("Skip DEBUG/INFO messages")
+
+    # check directory existence
+    if not os.path.exists(args.outdir):
+        os.makedirs(args.outdir)
+
+    # init process group
+    logging.info("Synchronizing between all workers.")
+    torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=n_gpus, rank=rank)
+    torch.cuda.set_device(rank)
+    logging.info("Finished init process group.")
+
+    # load and save config
+    with open(args.config) as f:
+        config = yaml.load(f, Loader=yaml.Loader)
+    config.update(vars(args))
+    config['rank'] = rank
+    config['distributed'] = True
+    config['world_size'] = n_gpus
+    # config["version"] = lscodec.__version__  # add version info
+    if rank == 0:
+        with open(os.path.join(args.outdir, "config.yml"), "w") as f:
+            yaml.dump(config, f, Dumper=yaml.Dumper)
+        for key, value in config.items():
+            logging.info(f"{key} = {value}")
+    
+    # get perturber
+    perturber = Perturbs(
+        formant_f0_perturb_params=config.get("formant_f0_perturb_params", dict()),
+        speed_perturb_params=config.get("speed_perturb_params", dict()),
+        speed_perturb_prob=config.get("speed_perturb_prob", 1.0)
+    )
+
+    # get dataset
+    train_dataset = AudioMelSCPDataset(
+        wav_scp=args.train_wav_scp,
+        prompt_scp=args.train_prompt_scp,
+        target_scp=args.train_target_scp,
+        aux_scp=args.train_aux_scp,
+        idx_scp=args.train_idx_scp,
+        utt2num_frames=args.train_num_frames,
+        segments=args.train_segments,
+        batch_frames=config.get("batch_frames", None),
+        batch_size=config.get("batch_size", None),
+        min_num_frames=config.get("min_num_frames", None),
+        max_num_frames=config.get("max_num_frames", None),
+        allow_cache=config.get("allow_cache", False),  # keep compatibility
+        length_tolerance=config.get("length_tolerance", 2),
+        prompt_min_length=config.get("prompt_min_length", 500),
+        extractor_fs=config['feature_extractor_params'].get("fs", 16000),
+        # speed_alpha_range=eval(config['speed_alpha_range'])
+        utt2spk=args.train_utt2spk,
+        spk2gender=args.train_spk2gender,
+        perturber=perturber
+    )
+    if rank == 0:
+        logging.info(f"The number of training files = {len(train_dataset)}.")
+    dev_dataset = AudioMelSCPDataset(
+        wav_scp=args.dev_wav_scp,
+        prompt_scp=args.dev_prompt_scp,
+        target_scp=args.dev_target_scp,
+        aux_scp=args.dev_aux_scp,
+        idx_scp=args.dev_idx_scp,
+        utt2num_frames=args.dev_num_frames,
+        segments=args.dev_segments,
+        min_num_frames=config.get("min_num_frames", None),
+        max_num_frames=config.get("max_num_frames", None),
+        allow_cache=config.get("allow_cache", False),  # keep compatibility
+        length_tolerance=config.get("length_tolerance", 2),
+        prompt_min_length=config.get("prompt_min_length", 500),
+        extractor_fs=config['feature_extractor_params'].get("fs", 16000),
+        # speed_alpha_range=eval(config['speed_alpha_range'])
+        utt2spk=args.dev_utt2spk,
+        spk2gender=args.dev_spk2gender,
+        perturber=perturber
+    )
+    if rank == 0:
+        logging.info(f"The number of development files = {len(dev_dataset)}.")
+    dataset = {
+        "train": train_dataset,
+        "dev": dev_dataset,
+    }
+
+    # get data loader
+    collator = Collator(
+        hop_size_for_extractor=config['feature_extractor_params']['hop_size'],
+        win_length_for_extractor=config['feature_extractor_params']['win_length'],
+        sampling_rate=config["feature_extractor_params"]["fs"],
+        n_mel=1024,
+        force_from_start=config.get("force_from_start", False)
+    )
+
+    sampler = {
+        "train": DistributedSampler(
+            dataset=dataset["train"],
+            num_replicas=n_gpus,
+            rank=rank,
+            shuffle=True,
+        ),
+        "dev": DistributedSampler(
+            dataset=dataset["dev"],
+            num_replicas=n_gpus,
+            rank=rank,
+            shuffle=False,
+        )}
+    data_loader = {
+        "train": DataLoader(
+            dataset=dataset["train"],
+            shuffle=False,
+            collate_fn=collator,
+            num_workers=config["num_workers"],
+            sampler=sampler["train"],
+            pin_memory=config["pin_memory"],
+        ),
+        "dev": DataLoader(
+            dataset=dataset["dev"],
+            shuffle=False,
+            collate_fn=collator,
+            num_workers=config["num_workers"],
+            sampler=sampler["dev"],
+            pin_memory=config["pin_memory"],
+        ),
+    }
+
+    from lscodec.utils import load_model
+    config['pretrain_codebook'] = args.pretrain_codebook
+    generator = load_model(config, None)
+
+    model = {
+        "generator": generator.to(device),
+    }
+
+    # define optimizers and schedulers
+    generator_optimizer_class = getattr(
+        lscodec.optimizers,
+        # keep compatibility
+        config.get("generator_optimizer_type", "RAdam"),
+    )
+    optimizer = {
+        "generator": generator_optimizer_class(
+            model["generator"].parameters(),
+            **config["generator_optimizer_params"],
+        ),
+    }
+    generator_scheduler_class = getattr(
+        torch.optim.lr_scheduler,
+        # keep compatibility
+        config.get("generator_scheduler_type", "StepLR"),
+    )
+    scheduler = {
+        "generator": generator_scheduler_class(
+            optimizer=optimizer["generator"],
+            **config["generator_scheduler_params"],
+        ),
+    }
+    from torch.nn.parallel import DistributedDataParallel
+    model["generator"] = DistributedDataParallel(model["generator"], device_ids=[rank], find_unused_parameters=True)
+
+    if rank == 0:
+        # show settings
+        logging.info(model["generator"])
+        logging.info(optimizer["generator"])
+
+    # define trainer
+    trainer = Trainer(
+        steps=0,
+        epochs=0,
+        data_loader=data_loader,
+        sampler=sampler,
+        model=model,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        config=config,
+        device=device,
+    )
+
+    if len(args.pretrain_vae) != 0:
+        trainer.load_vae_params(args.pretrain_vae)
+        if rank == 0:
+            logging.info(f"Successfully load parameters from {args.pretrain_vae}.")
+
+    # load pretrained parameters from checkpoint
+    if len(args.pretrain) != 0:
+        trainer.load_checkpoint(args.pretrain, load_only_params=True)
+        if rank == 0:
+            logging.info(f"Successfully load parameters from {args.pretrain}.")
+
+    if len(args.resume) != 0:
+        trainer.load_checkpoint(args.resume)
+        if rank == 0:
+            logging.info(f"Successfully resumed from {args.resume}.")
+
+    # run training loop
+    try:
+        trainer.run()
+    finally:
+        if rank == 0:
+            trainer.save_checkpoint(
+                os.path.join(config["outdir"], f"checkpoint-{trainer.steps}steps.pkl")
+            )
+            logging.info(f"Successfully saved checkpoint @ {trainer.steps}steps.")
+
+
+if __name__ == "__main__":
+    assert torch.cuda.is_available(), "CPU training is not allowed."
+    n_gpus = torch.cuda.device_count()
+    print(f"============> using {n_gpus} GPUS")
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "8001"
+
+    mp.spawn(
+        main,
+        nprocs=n_gpus,
+        args=(n_gpus,)
+    )

lscodec/datasets/perturbs.py

@@ -0,0 +1,308 @@
+# -*- coding: utf-8 -*-
+
+# Portions of this code are based on ContentVec.
+# The original source code can be found at:
+# https://github.com/auspicious3000/contentvec/blob/d746688a32940f4bee410ed7c87ec9cf8ff04f74/contentvec/data/audio/audio_utils_1.py
+#
+# Copyright (c) 2025 Yiwei Guo
+# Licensed under the MIT license.
+
+import os
+import numpy as np
+from scipy.signal import sosfilt
+from scipy import signal as sg
+import sox
+
+
+def make_lowshelf(g, fc, Q, fs=44100):
+    """Generate filter coefficients for 2nd order Lowshelf filter.
+    This function follows the code from the JUCE DSP library 
+    which can be found in `juce_IIRFilter.cpp`. 
+    
+    The design equations are based upon those found in the Cookbook 
+    formulae for audio equalizer biquad filter coefficients
+    by Robert Bristow-Johnson. 
+    https://www.w3.org/2011/audio/audio-eq-cookbook.html
+    Args:
+        g  (float): Gain factor in dB.
+        fc (float): Cutoff frequency in Hz.
+        Q  (float): Q factor.
+        fs (float): Sampling frequency in Hz.
+    Returns:
+        tuple: (b, a) filter coefficients 
+    """
+    # convert gain from dB to linear
+    g = np.power(10,(g/20))
+
+    # initial values
+    A = np.max([0.0, np.sqrt(g)])
+    aminus1 = A - 1
+    aplus1 = A + 1
+    omega = (2 * np.pi * np.max([fc, 2.0])) / fs
+    coso = np.cos(omega)
+    beta = np.sin(omega) * np.sqrt(A) / Q 
+    aminus1TimesCoso = aminus1 * coso
+
+    # coefs calculation
+    b0 = A * (aplus1 - aminus1TimesCoso + beta)
+    b1 = A * 2 * (aminus1 - aplus1 * coso)
+    b2 = A * (aplus1 - aminus1TimesCoso - beta)
+    a0 = aplus1 + aminus1TimesCoso + beta
+    a1 = -2 * (aminus1 + aplus1 * coso)
+    a2 = aplus1 + aminus1TimesCoso - beta
+
+    # output coefs 
+    #b = np.array([b0/a0, b1/a0, b2/a0])
+    #a = np.array([a0/a0, a1/a0, a2/a0])
+
+    return np.array([[b0/a0, b1/a0, b2/a0, 1.0, a1/a0, a2/a0]])
+
+
+
+def make_highself(g, fc, Q, fs=44100):
+    """Generate filter coefficients for 2nd order Highshelf filter.
+    This function follows the code from the JUCE DSP library 
+    which can be found in `juce_IIRFilter.cpp`. 
+    
+    The design equations are based upon those found in the Cookbook 
+    formulae for audio equalizer biquad filter coefficients
+    by Robert Bristow-Johnson. 
+    https://www.w3.org/2011/audio/audio-eq-cookbook.html
+    Args:
+        g  (float): Gain factor in dB.
+        fc (float): Cutoff frequency in Hz.
+        Q  (float): Q factor.
+        fs (float): Sampling frequency in Hz.
+    Returns:
+        tuple: (b, a) filter coefficients 
+    """
+    # convert gain from dB to linear
+    g = np.power(10,(g/20))
+
+    # initial values
+    A = np.max([0.0, np.sqrt(g)])
+    aminus1 = A - 1
+    aplus1 = A + 1
+    omega = (2 * np.pi * np.max([fc, 2.0])) / fs
+    coso = np.cos(omega)
+    beta = np.sin(omega) * np.sqrt(A) / Q 
+    aminus1TimesCoso = aminus1 * coso
+
+    # coefs calculation
+    b0 = A * (aplus1 + aminus1TimesCoso + beta)
+    b1 = A * -2 * (aminus1 + aplus1 * coso)
+    b2 = A * (aplus1 + aminus1TimesCoso - beta)
+    a0 = aplus1 - aminus1TimesCoso + beta
+    a1 = 2 * (aminus1 - aplus1 * coso)
+    a2 = aplus1 - aminus1TimesCoso - beta
+
+    # output coefs
+    #b = np.array([b0/a0, b1/a0, b2/a0])
+    #a = np.array([a0/a0, a1/a0, a2/a0])
+      
+    return np.array([[b0/a0, b1/a0, b2/a0, 1.0, a1/a0, a2/a0]])
+
+
+
+def make_peaking(g, fc, Q, fs=44100):
+    """Generate filter coefficients for 2nd order Peaking EQ.
+    This function follows the code from the JUCE DSP library 
+    which can be found in `juce_IIRFilter.cpp`. 
+    
+    The design equations are based upon those found in the Cookbook 
+    formulae for audio equalizer biquad filter coefficients
+    by Robert Bristow-Johnson. 
+    https://www.w3.org/2011/audio/audio-eq-cookbook.html
+    Args:
+        g  (float): Gain factor in dB.
+        fc (float): Cutoff frequency in Hz.
+        Q  (float): Q factor.
+        fs (float): Sampling frequency in Hz.
+    Returns:
+        tuple: (b, a) filter coefficients 
+    """
+    # convert gain from dB to linear
+    g = np.power(10,(g/20))
+
+    # initial values
+    A = np.max([0.0, np.sqrt(g)])
+    omega = (2 * np.pi * np.max([fc, 2.0])) / fs
+    alpha = np.sin(omega) / (Q * 2)
+    c2 = -2 * np.cos(omega)
+    alphaTimesA = alpha * A
+    alphaOverA = alpha / A
+
+    # coefs calculation
+    b0 = 1 + alphaTimesA
+    b1 = c2
+    b2 = 1 - alphaTimesA
+    a0 = 1 + alphaOverA
+    a1 = c2
+    a2 = 1 - alphaOverA
+
+    # output coefs
+    #b = np.array([b0/a0, b1/a0, b2/a0])
+    #a = np.array([a0/a0, a1/a0, a2/a0])
+    
+    return np.array([[b0/a0, b1/a0, b2/a0, 1.0, a1/a0, a2/a0]])
+
+
+
+def params2sos(G, Fc, Q, fs):
+    """Convert 5 band EQ paramaters to 2nd order sections.
+    Takes a vector with shape (13,) of denormalized EQ parameters
+    and calculates filter coefficients for each of the 5 filters.
+    These coefficients (2nd order sections) are then stored into a
+    single (5,6) matrix. This matrix can be fed to `scipy.signal.sosfreqz()`
+    in order to determine the frequency response of the cascasd of
+    all five biquad filters.
+    Args:
+        x  (float): Gain factor in dB.       
+        fs (float): Sampling frequency in Hz.
+    Returns:
+        ndarray: filter coefficients for 5 band EQ stored in (5,6) matrix.
+        [[b1_0, b1_1, b1_2, a1_0, a1_1, a1_2],  # lowshelf coefficients
+         [b2_0, b2_1, b2_2, a2_0, a2_1, a2_2],  # first band coefficients
+         [b3_0, b3_1, b3_2, a3_0, a3_1, a3_2],  # second band coefficients
+         [b4_0, b4_1, b4_2, a4_0, a4_1, a4_2],  # third band coefficients
+         [b5_0, b5_1, b5_2, a5_0, a5_1, a5_2]]  # highshelf coefficients
+    """
+    # generate filter coefficients from eq params
+    c0 = make_lowshelf(G[0], Fc[0], Q[0], fs=fs)
+    c1 = make_peaking (G[1], Fc[1], Q[1], fs=fs)
+    c2 = make_peaking (G[2], Fc[2], Q[2], fs=fs)
+    c3 = make_peaking (G[3], Fc[3], Q[3], fs=fs)
+    c4 = make_peaking (G[4], Fc[4], Q[4], fs=fs)
+    c5 = make_peaking (G[5], Fc[5], Q[5], fs=fs)
+    c6 = make_peaking (G[6], Fc[6], Q[6], fs=fs)
+    c7 = make_peaking (G[7], Fc[7], Q[7], fs=fs)
+    c8 = make_peaking (G[8], Fc[8], Q[8], fs=fs)
+    c9 = make_highself(G[9], Fc[9], Q[9], fs=fs)
+
+    # stuff coefficients into second order sections structure
+    sos = np.concatenate([c0,c1,c2,c3,c4,c5,c6,c7,c8,c9], axis=0)
+
+    return sos
+
+
+import parselmouth
+def change_gender(x, fs, lo, hi, ratio_fs, ratio_ps, ratio_pr):
+    s = parselmouth.Sound(x, sampling_frequency=fs)
+    f0 = s.to_pitch_ac(pitch_floor=lo, pitch_ceiling=hi, time_step=0.8/lo)
+    f0_np = f0.selected_array['frequency']
+    f0_med = np.median(f0_np[f0_np!=0]).item()
+    ss = parselmouth.praat.call([s, f0], "Change gender", ratio_fs, f0_med*ratio_ps, ratio_pr, 1.0)
+    return ss.values.squeeze(0)
+
+def change_gender_f0(x, fs, lo, hi, ratio_fs, new_f0_med, ratio_pr):
+    s = parselmouth.Sound(x, sampling_frequency=fs)
+    ss = parselmouth.praat.call(s, "Change gender", lo, hi, ratio_fs, new_f0_med, ratio_pr, 1.0)
+    return ss.values.squeeze(0)
+
+def parse_range(range_str):
+    # (a,b)|(c,d)|p
+    
+    terms = range_str.split("|")
+    assert len(terms)==3, range_str
+    prob = float(terms[-1])
+    ranges = [eval(terms[0]), eval(terms[1])]
+    return ranges, prob
+
+class FormantF0Perturb:
+    def __init__(self, male_formant_range, female_formant_range, female_f0_range):
+        # male_formant_range: (0.9,0.95)|(1.05,1.3)|0.25
+        # first range, second range, prob of first range
+        
+        self.Qmin = 2
+        self.Qmax = 5
+        self.male_formant_range, self.male_formant_prob = parse_range(male_formant_range)
+        self.female_formant_range, self.female_formant_prob = parse_range(female_formant_range)
+        self.female_f0_range, self.female_f0_prob = parse_range(female_f0_range)
+        
+    def get_perturb_param(self, gender):
+        # formant shift and pitch shift.
+        if gender == "M":
+            if np.random.random()<=self.male_formant_prob:
+                ratio_fs = np.random.uniform(self.male_formant_range[0][0], self.male_formant_range[0][1])
+            else:
+                ratio_fs = np.random.uniform(self.male_formant_range[1][0], self.male_formant_range[1][1])
+            return 75, 250, ratio_fs, 1.0
+        else:
+            if np.random.random()<=self.female_formant_prob:
+                ratio_fs = np.random.uniform(self.female_formant_range[0][0], self.female_formant_range[0][1])
+            else:
+                ratio_fs = np.random.uniform(self.female_formant_range[1][0], self.female_formant_range[1][1])
+            if np.random.random()<=self.female_f0_prob:
+                ratio_ps = np.random.uniform(self.female_f0_range[0][0], self.female_f0_range[0][1])
+            else:
+                ratio_ps = np.random.uniform(self.female_f0_range[1][0], self.female_f0_range[1][1])
+            return 100, 400, ratio_fs, ratio_ps
+
+    def random_eq(self, wav, sr):
+        Fc = np.exp(np.linspace(np.log(60), np.log(7600), 10))
+        # rng = np.random.default_rng()
+        z = np.random.uniform(0, 1, size=(10,))
+        Q = self.Qmin * (self.Qmax / self.Qmin)**z
+        G = np.random.uniform(-12, 12, size=(10,))
+        sos = params2sos(G, Fc, Q, sr)
+        wav = sosfilt(sos, wav)
+        return wav
+
+    def __call__(self, wav, sr, gender):
+        # gender: "F" or "M"
+        
+        lo, hi, ratio_fs, ratio_ps = self.get_perturb_param(gender)
+        # fs: formant shift. This controls timbre (higher -> more female). ps: pitch shift. This controls absolute pitch level
+        # We fix the pitch range parameter to 1. So that we don't change the pitch variations
+        ratio_pr = 1
+        
+        ss = change_gender(wav, sr, lo, hi, ratio_fs, ratio_ps, ratio_pr)
+        
+        return ss
+
+class SpeedPerturb:
+    def __init__(self, male_range, female_range):
+        self.male_range, self.male_prob = parse_range(male_range)
+        self.female_range, self.female_prob = parse_range(female_range)
+        
+    def __call__(self, wav, sr, gender):
+        
+        if gender == "F":
+            ranges, prob = self.male_range, self.male_prob
+        else:
+            ranges, prob = self.female_range, self.female_prob
+        
+        if np.random.random()<=prob:
+            speed_alpha = np.random.uniform(ranges[0][0], ranges[0][1])
+        else:
+            speed_alpha = np.random.uniform(ranges[1][0], ranges[1][1])
+        
+        self.transformer = sox.Transformer()
+        self.transformer.set_globals(verbosity=1)  # disable warnings.
+        self.transformer.speed(factor=speed_alpha)
+        self.transformer.tempo(factor=1/speed_alpha)  # recover the original length. Only affects timbre and pitch.
+        wav_perturb = self.transformer.build_array(input_array=wav, sample_rate_in=sr)
+        return wav_perturb
+
+class IdentityPerturb:
+    def __init__(self, ):
+        pass
+    def __call__(self,wav, sr, gender):
+        return wav
+        
+class Perturbs:
+    def __init__(self, formant_f0_perturb_params, speed_perturb_params, speed_perturb_prob):
+        self.speed_perturb_prob = speed_perturb_prob
+        if speed_perturb_prob!=0.0:
+            self.speed_perturber = SpeedPerturb(**speed_perturb_params)
+        else:
+            self.speed_perturber = IdentityPerturb()
+        if speed_perturb_prob!=1.0:
+            self.formant_perturber = FormantF0Perturb(**formant_f0_perturb_params)
+        else:
+            self.formant_perturber = IdentityPerturb()
+    def __call__(self, wav, sr, gender):
+        if np.random.random()<=self.speed_perturb_prob:
+            return self.speed_perturber(wav,sr,gender)
+        else:
+            return self.formant_perturber(wav,sr,gender)

lscodec/datasets/scp_dataset.py

@@ -0,0 +1,381 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+# Portions of this code are a modification of the ParallelWaveGAN project.
+# The original source code can be found at:
+# https://github.com/kan-bayashi/ParallelWaveGAN/blob/master/parallel_wavegan/datasets/scp_dataset.py
+#
+# Copyright (c) 2025 Yiwei Guo
+# Licensed under the MIT license.
+
+
+"""Dataset modules based on kaldi-style scp files."""
+
+import torch
+import logging
+import random
+import copy
+import sox
+from multiprocessing import Manager
+
+import kaldiio
+import librosa
+import numpy as np
+
+from torch.utils.data import Dataset
+from tqdm import tqdm
+from lscodec.utils import HDF5ScpLoader
+from lscodec.utils import NpyScpLoader
+from lscodec.datasets.perturbs import Perturbs
+
+logging.getLogger("sox").setLevel(logging.ERROR)
+
+
+def _get_feats_scp_loader(feats_scp):
+    # read the first line of feats.scp file
+    with open(feats_scp) as f:
+        key, value = f.readlines()[0].replace("\n", "").split()
+
+    # check scp type
+    if ":" in value:
+        value_1, value_2 = value.split(":")
+        if value_1.endswith(".ark"):
+            # kaldi-ark case: utt_id_1 /path/to/utt_id_1.ark:index
+            return kaldiio.load_scp(feats_scp)
+        elif value_1.endswith(".h5"):
+            # hdf5 case with path in hdf5: utt_id_1 /path/to/utt_id_1.h5:feats
+            return HDF5ScpLoader(feats_scp)
+        else:
+            raise ValueError("Not supported feats.scp type.")
+    else:
+        if value.endswith(".h5"):
+            # hdf5 case without path in hdf5: utt_id_1 /path/to/utt_id_1.h5
+            return HDF5ScpLoader(feats_scp)
+        elif value.endswith(".npy"):
+            # npy case: utt_id_1 /path/to/utt_id_1.npy
+            return NpyScpLoader(feats_scp)
+        else:
+            raise ValueError("Not supported feats.scp type.")
+
+
+class AudioMelSCPDataset(Dataset):
+    """PyTorch compatible audio and feat dataset based on kaldi-stype scp files."""
+
+    def __init__(
+        self,
+        wav_scp,
+        prompt_scp,
+        target_scp,
+        aux_scp,
+        idx_scp,
+        utt2num_frames=None,
+        segments=None,
+        batch_frames=None,
+        batch_size=None,
+        min_num_frames=None,
+        max_num_frames=None,
+        return_utt_id=False,
+        return_sampling_rate=False,
+        allow_cache=False,
+        length_tolerance=2,
+        prompt_min_length=500,  # in terms of frames
+        extractor_fs=16000,
+        # speed_alpha_range=(0.8, 1.2),
+        perturber=None,
+        utt2spk=None,
+        spk2gender=None
+    ):
+        """Initialize dataset.
+
+        Args:
+            wav_scp (str): Kaldi-style wav.scp file.
+            prompt_scp (str): Kaldi-style fests.scp file.
+            aux_scp (str): Kaldi-style fests.scp file.
+            segments (str): Kaldi-style segments file.
+            min_num_frames (int): Threshold to remove short feature files.
+            max_num_frames (int): Threshold to remove long feature files.
+            return_utt_id (bool): Whether to return utterance id.
+            return_sampling_rate (bool): Whether to return sampling rate.
+            allow_cache (bool): Whether to allow cache of the loaded files.
+
+        """
+        # load scp as lazy dict
+        self.audio_loader = kaldiio.load_scp(wav_scp, segments=segments)
+        self.prompt_loader = _get_feats_scp_loader(prompt_scp)
+        # with open(prompt_scp, 'r') as fr:
+        #     self.prompt_loader = dict()
+        #     for line in fr.readlines():
+        #         terms = line.strip().split()
+        #         utt, path = terms
+        #         self.prompt_loader[utt] = path
+
+        self.target_loader = _get_feats_scp_loader(target_scp)
+        self.aux_loader = _get_feats_scp_loader(aux_scp)
+        self.idx_loader = _get_feats_scp_loader(idx_scp)
+        
+        self.utt_ids = list(self.prompt_loader.keys())
+        self.return_utt_id = return_utt_id
+        self.return_sampling_rate = return_sampling_rate
+        self.allow_cache = allow_cache
+        self.extractor_fs = extractor_fs
+
+        utt2num_frames_loader = None
+        if utt2num_frames is not None:
+            with open(utt2num_frames, 'r') as f:
+                utt2num_frames_loader = dict([(x.split()[0], int(x.split()[1])) for x in f.readlines()])
+        else:
+            utt2num_frames_loader = dict([(k, mel.shape[0]) for k, mel in self.prompt_loader.items()])
+        self.utt2num_frames_loader = utt2num_frames_loader
+        # filter by threshold
+        if (min_num_frames or max_num_frames) is not None:
+            mel_lengths = [utt2num_frames_loader[key] for key in self.utt_ids]
+            idxs = [
+                idx
+                for idx in range(len(self.utt_ids))
+                if (min_num_frames and mel_lengths[idx] >= min_num_frames) and (max_num_frames and mel_lengths[idx] <= max_num_frames)
+            ]
+            if len(self.utt_ids) != len(idxs):
+                logging.warning(
+                    f"Some files are filtered by mel length threshold "
+                    f"({len(self.utt_ids)} -> {len(idxs)})."
+                )
+            self.utt_ids = [self.utt_ids[idx] for idx in idxs]
+
+        # batchify
+        if batch_frames is not None:
+            self.batches = self.batchify(utt2num_frames_loader, batch_frames=batch_frames)
+        elif batch_size is not None:
+            self.batches = self.batchify(utt2num_frames_loader, batch_size=batch_size)
+        else:
+            self.batches = [[utt_id] for utt_id in self.utt_ids]
+
+        if allow_cache:
+            # NOTE(kan-bayashi): Manager is need to share memory in dataloader with num_workers > 0
+            self.manager = Manager()
+            self.caches = self.manager.dict()
+        self.length_tolerance = length_tolerance
+        self.prompt_min_length = prompt_min_length
+
+        # assert len(speed_alpha_range)==2, speed_alpha_range
+        # self.alpha_range = speed_alpha_range
+        self.perturber = perturber
+        
+        if utt2spk is not None:
+            self.utt2spk = dict()
+            with open(utt2spk, 'r') as f:
+                for line in f.readlines():
+                    terms = line.strip().split()
+                    self.utt2spk[terms[0]] = terms[1]
+            assert spk2gender is not None
+            self.spk2gender = dict()
+            with open(spk2gender, 'r') as f:
+                for line in f.readlines():
+                    terms = line.strip().split()
+                    assert terms[1] in {"F", "M"}
+                    self.spk2gender[terms[0]] = terms[1]
+        else:
+            self.utt2spk = None
+
+    def batchify(self, utt2num_frames_loader, batch_frames=None, batch_size=None, min_batch_size=1, drop_last=True):
+
+        assert batch_size is None or batch_size > min_batch_size
+
+        batches = []
+        batch = []
+        accum_num_frames = 0
+        utt_ids_set = set(self.utt_ids)
+        for utt_id, mel_length in tqdm(sorted(list(utt2num_frames_loader.items()), key=lambda x: x[1], reverse=True)):
+            if utt_id not in utt_ids_set:
+                continue
+            if (batch_frames is not None and accum_num_frames + mel_length > batch_frames and len(batch) > min_batch_size) or (batch_size is not None and len(batch) == batch_size):
+                batches.append(batch)
+                batch = []
+                accum_num_frames = 0
+            batch.append(utt_id)
+            accum_num_frames += mel_length
+        if len(batch) > min_batch_size and not drop_last:
+            batches.append(batch)
+        return batches
+
+    def __getitem__(self, idx):
+        """Get specified idx items.
+
+        Args:
+            idx (int): Index of the item.
+
+        Returns:
+            str: Utterance id (only in return_utt_id = True).
+            ndarray or tuple: Audio signal (T,) or (w/ sampling rate if return_sampling_rate = True).
+            ndarray: Feature (T', C).
+
+        """
+        batch = self.batches[idx]
+        batch_items = []
+
+        for utt_id in batch:
+            if self.allow_cache and self.caches.get(utt_id) is not None:
+                items = self.caches[utt_id]
+            else:
+                fs, audio = self.audio_loader[utt_id]
+                prompt = self.prompt_loader[utt_id]
+                # prompt = torch.load(prompt_path).numpy()
+                target = self.target_loader[utt_id]
+                ssl_idx = self.idx_loader[utt_id]
+                aux = self.aux_loader[utt_id]
+
+                min_len = min(len(prompt), len(aux), len(target), len(ssl_idx))
+                assert ((abs(len(prompt) - min_len) <= self.length_tolerance) and
+                        (abs(len(target) - min_len) <= self.length_tolerance) and
+                        (abs(len(ssl_idx) - min_len) <= self.length_tolerance) and
+                        (abs(len(aux) - min_len) <= self.length_tolerance)), \
+                    f"Audio feature lengths difference exceeds length tolerance for {utt_id}"
+                prompt, aux, target, ssl_idx = prompt[:min_len], aux[:min_len], target[:min_len], ssl_idx[:min_len]
+
+                if self.utt2spk is not None:
+                    spk=self.utt2spk[utt_id]
+                    gender=self.spk2gender[spk]
+                else:
+                    gender = "M"  # Default to male, so that not use pitch shift
+                    
+                # normalize audio signal to be [-1, 1]
+                audio = audio.astype(np.float32)
+                audio /= 1 << (16 - 1)  # assume that wav is PCM 16 bit
+
+                if fs != self.extractor_fs:
+                    audio_for_extractor = librosa.resample(audio, orig_sr=fs, target_sr=self.extractor_fs)
+                else:
+                    audio_for_extractor = audio
+
+                # # perturb 
+                # if self.alpha_range[0] == self.alpha_range[1] == 1.0:
+                #     pass
+                # else:
+                #     speed_alpha = random.random() * (self.alpha_range[1] - self.alpha_range[0]) + self.alpha_range[0]
+                #     self.transformer = sox.Transformer()
+                #     self.transformer.set_globals(verbosity=1)  # disable warnings.
+                #     self.transformer.speed(factor=speed_alpha)
+                #     self.transformer.tempo(factor=1/speed_alpha)  # recover the original length. Only affects timbre and pitch.
+                #     audio_for_extractor = self.transformer.build_array(input_array=audio_for_extractor, sample_rate_in=self.extractor_fs)
+                audio_for_extractor = self.perturber(audio_for_extractor, self.extractor_fs, gender)
+
+                if self.return_utt_id:
+                    items = utt_id, audio_for_extractor, prompt, aux, target, ssl_idx
+                else:
+                    items = audio_for_extractor, prompt, aux, target, ssl_idx
+
+                if self.allow_cache:
+                    self.caches[utt_id] = items
+
+            batch_items.append(items)
+
+        return batch_items
+
+    def __len__(self):
+        """Return dataset length.
+        Returns:
+            int: The length of dataset.
+        """
+        return len(self.batches)
+
+
+class MelSCPDataset(Dataset):
+    """PyTorch compatible feat dataset based on kaldi-stype scp files."""
+
+    def __init__(
+        self,
+        vqidx_scp,
+        prompt_scp,
+        utt2num_frames=None,
+        min_num_frames=None,
+        max_num_frames=None,
+        return_utt_id=False,
+        allow_cache=False,
+    ):
+        """Initialize dataset.
+
+        Args:
+            vqidx_scp (str): Kaldi-style fests.scp file.
+            prompt_scp (str): Kaldi-style scp file. In this file, every utt is associated with its prompt's mel-spectrogram.
+            min_num_frames (int): Threshold to remove short feature files.
+            max_num_frames (int): Threshold to remove long feature files.
+            return_utt_id (bool): Whether to return utterance id.
+            allow_cache (bool): Whether to allow cache of the loaded files.
+        """
+        # load scp as lazy dict
+        vqidx_loader = _get_feats_scp_loader(vqidx_scp)
+        prompt_loader = _get_feats_scp_loader(prompt_scp)
+        vqidx_keys = list(set(prompt_loader.keys()) & set(vqidx_loader.keys()))
+
+        utt2num_frames_loader = None
+        if utt2num_frames is not None:
+            with open(utt2num_frames, 'r') as f:
+                utt2num_frames_loader = dict([(x.split()[0], int(x.split()[1])) for x in f.readlines()])
+        else:
+            utt2num_frames_loader = dict([(k, mel.shape[0]) for k, mel in vqidx_loader.items()])
+
+        # filter by threshold
+        if (min_num_frames or max_num_frames) is not None:
+            mel_lengths = [utt2num_frames_loader[key] for key in vqidx_keys]
+            idxs = [
+                idx
+                for idx in range(len(vqidx_keys))
+                if (min_num_frames and mel_lengths[idx] > min_num_frames) and (max_num_frames and mel_lengths[idx] < max_num_frames)
+            ]
+            if len(vqidx_keys) != len(idxs):
+                logging.warning(
+                    f"Some files are filtered by mel length threshold "
+                    f"({len(vqidx_keys)} -> {len(idxs)})."
+                )
+            vqidx_keys = [vqidx_keys[idx] for idx in idxs]
+
+        self.vqidx_loader = vqidx_loader
+        self.prompt_loader = prompt_loader
+        self.utt_ids = vqidx_keys
+        self.return_utt_id = return_utt_id
+        self.allow_cache = allow_cache
+
+        if allow_cache:
+            # NOTE(kan-bayashi): Manager is need to share memory in dataloader with num_workers > 0
+            self.manager = Manager()
+            self.caches = self.manager.list()
+            self.caches += [() for _ in range(len(self.utt_ids))]
+
+    def __getitem__(self, idx):
+        """Get specified idx items.
+
+        Args:
+            idx (int): Index of the item.
+
+        Returns:
+            str: Utterance id (only in return_utt_id = True).
+            ndarray: Feature (T', C).
+
+        """
+        if self.allow_cache and len(self.caches[idx]) != 0:
+            return self.caches[idx]
+
+        utt_id = self.utt_ids[idx]
+        vqidx = self.vqidx_loader[utt_id].astype(int)
+
+        prompt = self.prompt_loader[utt_id].copy()
+
+        if self.return_utt_id:
+            items = utt_id, vqidx, prompt
+        else:
+            items = vqidx, prompt
+
+        if self.allow_cache:
+            self.caches[idx] = items
+
+        return items
+
+    def __len__(self):
+        """Return dataset length.
+
+        Returns:
+            int: The length of dataset.
+
+        """
+        return len(self.utt_ids)

lscodec/distributed/distrib.py

@@ -0,0 +1,124 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Torch distributed utilities."""
+
+import typing as tp
+
+import torch
+
+
+def rank():
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_rank()
+    else:
+        return 0
+
+
+def world_size():
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_world_size()
+    else:
+        return 1
+
+
+def is_distributed():
+    return world_size() > 1
+
+
+def all_reduce(tensor: torch.Tensor, op=torch.distributed.ReduceOp.SUM):
+    if is_distributed():
+        return torch.distributed.all_reduce(tensor, op)
+
+
+def _is_complex_or_float(tensor):
+    return torch.is_floating_point(tensor) or torch.is_complex(tensor)
+
+
+def _check_number_of_params(params: tp.List[torch.Tensor]):
+    # utility function to check that the number of params in all workers is the same,
+    # and thus avoid a deadlock with distributed all reduce.
+    if not is_distributed() or not params:
+        return
+    tensor = torch.tensor([len(params)], device=params[0].device, dtype=torch.long)
+    all_reduce(tensor)
+    if tensor.item() != len(params) * world_size():
+        # If not all the workers have the same number, for at least one of them,
+        # this inequality will be verified.
+        raise RuntimeError(f"Mismatch in number of params: ours is {len(params)}, "
+                           "at least one worker has a different one.")
+
+
+def broadcast_tensors(tensors: tp.Iterable[torch.Tensor], src: int = 0):
+    """Broadcast the tensors from the given parameters to all workers.
+    This can be used to ensure that all workers have the same model to start with.
+    """
+    if not is_distributed():
+        return
+    tensors = [tensor for tensor in tensors if _is_complex_or_float(tensor)]
+    _check_number_of_params(tensors)
+    handles = []
+    for tensor in tensors:
+        handle = torch.distributed.broadcast(tensor.data, src=src, async_op=True)
+        handles.append(handle)
+    for handle in handles:
+        handle.wait()
+
+
+def sync_buffer(buffers, average=True):
+    """
+    Sync grad for buffers. If average is False, broadcast instead of averaging.
+    """
+    if not is_distributed():
+        return
+    handles = []
+    for buffer in buffers:
+        if torch.is_floating_point(buffer.data):
+            if average:
+                handle = torch.distributed.all_reduce(
+                    buffer.data, op=torch.distributed.ReduceOp.SUM, async_op=True)
+            else:
+                handle = torch.distributed.broadcast(
+                    buffer.data, src=0, async_op=True)
+            handles.append((buffer, handle))
+    for buffer, handle in handles:
+        handle.wait()
+        if average:
+            buffer.data /= world_size
+
+
+def sync_grad(params):
+    """
+    Simpler alternative to DistributedDataParallel, that doesn't rely
+    on any black magic. For simple models it can also be as fast.
+    Just call this on your model parameters after the call to backward!
+    """
+    if not is_distributed():
+        return
+    handles = []
+    for p in params:
+        if p.grad is not None:
+            handle = torch.distributed.all_reduce(
+                p.grad.data, op=torch.distributed.ReduceOp.SUM, async_op=True)
+            handles.append((p, handle))
+    for p, handle in handles:
+        handle.wait()
+        p.grad.data /= world_size()
+
+
+def average_metrics(metrics: tp.Dict[str, float], count=1.):
+    """Average a dictionary of metrics across all workers, using the optional
+    `count` as unormalized weight.
+    """
+    if not is_distributed():
+        return metrics
+    keys, values = zip(*metrics.items())
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    tensor = torch.tensor(list(values) + [1], device=device, dtype=torch.float32)
+    tensor *= count
+    all_reduce(tensor)
+    averaged = (tensor[:-1] / tensor[-1]).cpu().tolist()
+    return dict(zip(keys, averaged))

lscodec/distributed/launch.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""Distributed process launcher.
+
+This code is modified from https://github.com/pytorch/pytorch/blob/v1.3.0/torch/distributed/launch.py.
+
+"""
+import os
+import subprocess
+import sys
+
+from argparse import ArgumentParser
+from argparse import REMAINDER
+
+
+def parse_args():
+    """Parse arguments."""
+    parser = ArgumentParser(
+        description="PyTorch distributed training launch "
+        "helper utilty that will spawn up "
+        "multiple distributed processes"
+    )
+
+    # Optional arguments for the launch helper
+    parser.add_argument(
+        "--nnodes",
+        type=int,
+        default=1,
+        help="The number of nodes to use for distributed " "training",
+    )
+    parser.add_argument(
+        "--node_rank",
+        type=int,
+        default=0,
+        help="The rank of the node for multi-node distributed " "training",
+    )
+    parser.add_argument(
+        "--nproc_per_node",
+        type=int,
+        default=1,
+        help="The number of processes to launch on each node, "
+        "for GPU training, this is recommended to be set "
+        "to the number of GPUs in your system so that "
+        "each process can be bound to a single GPU.",
+    )
+    parser.add_argument(
+        "--master_addr",
+        default="127.0.0.1",
+        type=str,
+        help="Master node (rank 0)'s address, should be either "
+        "the IP address or the hostname of node 0, for "
+        "single node multi-proc training, the "
+        "--master_addr can simply be 127.0.0.1",
+    )
+    parser.add_argument(
+        "--master_port",
+        default=29500,
+        type=int,
+        help="Master node (rank 0)'s free port that needs to "
+        "be used for communciation during distributed "
+        "training",
+    )
+    parser.add_argument(
+        "--use_env",
+        default=False,
+        action="store_true",
+        help="Use environment variable to pass "
+        "'local rank'. For legacy reasons, the default value is False. "
+        "If set to True, the script will not pass "
+        "--local_rank as argument, and will instead set LOCAL_RANK.",
+    )
+    parser.add_argument(
+        "-m",
+        "--module",
+        default=False,
+        action="store_true",
+        help="Changes each process to interpret the launch script "
+        "as a python module, executing with the same behavior as"
+        "'python -m'.",
+    )
+    parser.add_argument(
+        "-c",
+        "--command",
+        default=False,
+        action="store_true",
+        help="Changes each process to interpret the launch script " "as a command.",
+    )
+
+    # positional
+    parser.add_argument(
+        "training_script",
+        type=str,
+        help="The full path to the single GPU training "
+        "program/script/command to be launched in parallel, "
+        "followed by all the arguments for the "
+        "training script",
+    )
+
+    # rest from the training program
+    parser.add_argument("training_script_args", nargs=REMAINDER)
+    return parser.parse_args()
+
+
+def main():
+    """Launch distributed processes."""
+    args = parse_args()
+
+    # world size in terms of number of processes
+    dist_world_size = args.nproc_per_node * args.nnodes
+
+    # set PyTorch distributed related environmental variables
+    current_env = os.environ.copy()
+    current_env["MASTER_ADDR"] = args.master_addr
+    current_env["MASTER_PORT"] = str(args.master_port)
+    current_env["WORLD_SIZE"] = str(dist_world_size)
+
+    processes = []
+
+    if "OMP_NUM_THREADS" not in os.environ and args.nproc_per_node > 1:
+        current_env["OMP_NUM_THREADS"] = str(1)
+        print(
+            "*****************************************\n"
+            "Setting OMP_NUM_THREADS environment variable for each process "
+            "to be {} in default, to avoid your system being overloaded, "
+            "please further tune the variable for optimal performance in "
+            "your application as needed. \n"
+            "*****************************************".format(
+                current_env["OMP_NUM_THREADS"]
+            )
+        )
+
+    for local_rank in range(0, args.nproc_per_node):
+        # each process's rank
+        dist_rank = args.nproc_per_node * args.node_rank + local_rank
+        current_env["RANK"] = str(dist_rank)
+        current_env["LOCAL_RANK"] = str(local_rank)
+
+        # spawn the processes
+        if args.command:
+            cmd = [args.training_script]
+        else:
+            cmd = [sys.executable, "-u"]
+            if args.module:
+                cmd.append("-m")
+            cmd.append(args.training_script)
+
+        if not args.use_env:
+            cmd.append("--local_rank={}".format(local_rank))
+
+        cmd.extend(args.training_script_args)
+
+        process = subprocess.Popen(cmd, env=current_env)
+        processes.append(process)
+
+    for process in processes:
+        process.wait()
+        if process.returncode != 0:
+            raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
+
+
+if __name__ == "__main__":
+    main()

lscodec/layers/__init__.py

@@ -0,0 +1,6 @@
+from .causal_conv import *  # NOQA
+from .pqmf import *  # NOQA
+from .residual_block import *  # NOQA
+from .residual_stack import *  # NOQA
+from .tade_res_block import *  # NOQA
+from .upsample import *  # NOQA

lscodec/layers/causal_conv.py

@@ -0,0 +1,66 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+"""Causal convolusion layer modules."""
+
+
+import torch
+
+
+class CausalConv1d(torch.nn.Module):
+    """CausalConv1d module with customized initialization."""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        dilation=1,
+        bias=True,
+        pad="ConstantPad1d",
+        pad_params={"value": 0.0},
+    ):
+        """Initialize CausalConv1d module."""
+        super(CausalConv1d, self).__init__()
+        self.pad = getattr(torch.nn, pad)((kernel_size - 1) * dilation, **pad_params)
+        self.conv = torch.nn.Conv1d(
+            in_channels, out_channels, kernel_size, dilation=dilation, bias=bias
+        )
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input tensor (B, in_channels, T).
+
+        Returns:
+            Tensor: Output tensor (B, out_channels, T).
+
+        """
+        return self.conv(self.pad(x))[:, :, : x.size(2)]
+
+
+class CausalConvTranspose1d(torch.nn.Module):
+    """CausalConvTranspose1d module with customized initialization."""
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride, bias=True):
+        """Initialize CausalConvTranspose1d module."""
+        super(CausalConvTranspose1d, self).__init__()
+        self.deconv = torch.nn.ConvTranspose1d(
+            in_channels, out_channels, kernel_size, stride, bias=bias
+        )
+        self.stride = stride
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input tensor (B, in_channels, T_in).
+
+        Returns:
+            Tensor: Output tensor (B, out_channels, T_out).
+
+        """
+        return self.deconv(x)[:, :, : -self.stride]

lscodec/layers/pqmf.py

@@ -0,0 +1,150 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+"""Pseudo QMF modules."""
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from scipy.signal import kaiser
+
+
+def design_prototype_filter(taps=62, cutoff_ratio=0.142, beta=9.0):
+    """Design prototype filter for PQMF.
+
+    This method is based on `A Kaiser window approach for the design of prototype
+    filters of cosine modulated filterbanks`_.
+
+    Args:
+        taps (int): The number of filter taps.
+        cutoff_ratio (float): Cut-off frequency ratio.
+        beta (float): Beta coefficient for kaiser window.
+
+    Returns:
+        ndarray: Impluse response of prototype filter (taps + 1,).
+
+    .. _`A Kaiser window approach for the design of prototype filters of cosine modulated filterbanks`:
+        https://ieeexplore.ieee.org/abstract/document/681427
+
+    """
+    # check the arguments are valid
+    assert taps % 2 == 0, "The number of taps mush be even number."
+    assert 0.0 < cutoff_ratio < 1.0, "Cutoff ratio must be > 0.0 and < 1.0."
+
+    # make initial filter
+    omega_c = np.pi * cutoff_ratio
+    with np.errstate(invalid="ignore"):
+        h_i = np.sin(omega_c * (np.arange(taps + 1) - 0.5 * taps)) / (
+            np.pi * (np.arange(taps + 1) - 0.5 * taps)
+        )
+    h_i[taps // 2] = np.cos(0) * cutoff_ratio  # fix nan due to indeterminate form
+
+    # apply kaiser window
+    w = kaiser(taps + 1, beta)
+    h = h_i * w
+
+    return h
+
+
+class PQMF(torch.nn.Module):
+    """PQMF module.
+
+    This module is based on `Near-perfect-reconstruction pseudo-QMF banks`_.
+
+    .. _`Near-perfect-reconstruction pseudo-QMF banks`:
+        https://ieeexplore.ieee.org/document/258122
+
+    """
+
+    def __init__(self, subbands=4, taps=62, cutoff_ratio=0.142, beta=9.0):
+        """Initilize PQMF module.
+
+        The cutoff_ratio and beta parameters are optimized for #subbands = 4.
+        See dicussion in https://github.com/kan-bayashi/ParallelWaveGAN/issues/195.
+
+        Args:
+            subbands (int): The number of subbands.
+            taps (int): The number of filter taps.
+            cutoff_ratio (float): Cut-off frequency ratio.
+            beta (float): Beta coefficient for kaiser window.
+
+        """
+        super(PQMF, self).__init__()
+
+        # build analysis & synthesis filter coefficients
+        h_proto = design_prototype_filter(taps, cutoff_ratio, beta)
+        h_analysis = np.zeros((subbands, len(h_proto)))
+        h_synthesis = np.zeros((subbands, len(h_proto)))
+        for k in range(subbands):
+            h_analysis[k] = (
+                2
+                * h_proto
+                * np.cos(
+                    (2 * k + 1)
+                    * (np.pi / (2 * subbands))
+                    * (np.arange(taps + 1) - (taps / 2))
+                    + (-1) ** k * np.pi / 4
+                )
+            )
+            h_synthesis[k] = (
+                2
+                * h_proto
+                * np.cos(
+                    (2 * k + 1)
+                    * (np.pi / (2 * subbands))
+                    * (np.arange(taps + 1) - (taps / 2))
+                    - (-1) ** k * np.pi / 4
+                )
+            )
+
+        # convert to tensor
+        analysis_filter = torch.from_numpy(h_analysis).float().unsqueeze(1)
+        synthesis_filter = torch.from_numpy(h_synthesis).float().unsqueeze(0)
+
+        # register coefficients as beffer
+        self.register_buffer("analysis_filter", analysis_filter)
+        self.register_buffer("synthesis_filter", synthesis_filter)
+
+        # filter for downsampling & upsampling
+        updown_filter = torch.zeros((subbands, subbands, subbands)).float()
+        for k in range(subbands):
+            updown_filter[k, k, 0] = 1.0
+        self.register_buffer("updown_filter", updown_filter)
+        self.subbands = subbands
+
+        # keep padding info
+        self.pad_fn = torch.nn.ConstantPad1d(taps // 2, 0.0)
+
+    def analysis(self, x):
+        """Analysis with PQMF.
+
+        Args:
+            x (Tensor): Input tensor (B, 1, T).
+
+        Returns:
+            Tensor: Output tensor (B, subbands, T // subbands).
+
+        """
+        x = F.conv1d(self.pad_fn(x), self.analysis_filter)
+        return F.conv1d(x, self.updown_filter, stride=self.subbands)
+
+    def synthesis(self, x):
+        """Synthesis with PQMF.
+
+        Args:
+            x (Tensor): Input tensor (B, subbands, T // subbands).
+
+        Returns:
+            Tensor: Output tensor (B, 1, T).
+
+        """
+        # NOTE(kan-bayashi): Power will be dreased so here multipy by # subbands.
+        #   Not sure this is the correct way, it is better to check again.
+        # TODO(kan-bayashi): Understand the reconstruction procedure
+        x = F.conv_transpose1d(
+            x, self.updown_filter * self.subbands, stride=self.subbands
+        )
+        return F.conv1d(self.pad_fn(x), self.synthesis_filter)

lscodec/layers/residual_block.py

@@ -0,0 +1,222 @@
+# -*- coding: utf-8 -*-
+
+"""Residual block modules.
+
+References:
+    - https://github.com/r9y9/wavenet_vocoder
+    - https://github.com/jik876/hifi-gan
+
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+class Conv1d(torch.nn.Conv1d):
+    """Conv1d module with customized initialization."""
+
+    def __init__(self, *args, **kwargs):
+        """Initialize Conv1d module."""
+        super(Conv1d, self).__init__(*args, **kwargs)
+
+    def reset_parameters(self):
+        """Reset parameters."""
+        torch.nn.init.kaiming_normal_(self.weight, nonlinearity="relu")
+        if self.bias is not None:
+            torch.nn.init.constant_(self.bias, 0.0)
+
+
+class Conv1d1x1(Conv1d):
+    """1x1 Conv1d with customized initialization."""
+
+    def __init__(self, in_channels, out_channels, bias):
+        """Initialize 1x1 Conv1d module."""
+        super(Conv1d1x1, self).__init__(
+            in_channels, out_channels, kernel_size=1, padding=0, dilation=1, bias=bias
+        )
+
+
+class WaveNetResidualBlock(torch.nn.Module):
+    """Residual block module in WaveNet."""
+
+    def __init__(
+        self,
+        kernel_size=3,
+        residual_channels=64,
+        gate_channels=128,
+        skip_channels=64,
+        aux_channels=80,
+        dropout=0.0,
+        dilation=1,
+        bias=True,
+        use_causal_conv=False,
+    ):
+        """Initialize WaveNetResidualBlock module.
+
+        Args:
+            kernel_size (int): Kernel size of dilation convolution layer.
+            residual_channels (int): Number of channels for residual connection.
+            skip_channels (int): Number of channels for skip connection.
+            aux_channels (int): Local conditioning channels i.e. auxiliary input dimension.
+            dropout (float): Dropout probability.
+            dilation (int): Dilation factor.
+            bias (bool): Whether to add bias parameter in convolution layers.
+            use_causal_conv (bool): Whether to use use_causal_conv or non-use_causal_conv convolution.
+
+        """
+        super().__init__()
+        self.dropout = dropout
+        # no future time stamps available
+        if use_causal_conv:
+            padding = (kernel_size - 1) * dilation
+        else:
+            assert (kernel_size - 1) % 2 == 0, "Not support even number kernel size."
+            padding = (kernel_size - 1) // 2 * dilation
+        self.use_causal_conv = use_causal_conv
+
+        # dilation conv
+        self.conv = Conv1d(
+            residual_channels,
+            gate_channels,
+            kernel_size,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+
+        # local conditioning
+        if aux_channels > 0:
+            self.conv1x1_aux = Conv1d1x1(aux_channels, gate_channels, bias=False)
+        else:
+            self.conv1x1_aux = None
+
+        # conv output is split into two groups
+        gate_out_channels = gate_channels // 2
+        self.conv1x1_out = Conv1d1x1(gate_out_channels, residual_channels, bias=bias)
+        self.conv1x1_skip = Conv1d1x1(gate_out_channels, skip_channels, bias=bias)
+
+    def forward(self, x, c):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input tensor (B, residual_channels, T).
+            c (Tensor): Local conditioning auxiliary tensor (B, aux_channels, T).
+
+        Returns:
+            Tensor: Output tensor for residual connection (B, residual_channels, T).
+            Tensor: Output tensor for skip connection (B, skip_channels, T).
+
+        """
+        residual = x
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.conv(x)
+
+        # remove future time steps if use_causal_conv conv
+        x = x[:, :, : residual.size(-1)] if self.use_causal_conv else x
+
+        # split into two part for gated activation
+        splitdim = 1
+        xa, xb = x.split(x.size(splitdim) // 2, dim=splitdim)
+
+        # local conditioning
+        if c is not None:
+            assert self.conv1x1_aux is not None
+            c = self.conv1x1_aux(c)
+            ca, cb = c.split(c.size(splitdim) // 2, dim=splitdim)
+            xa, xb = xa + ca, xb + cb
+
+        x = torch.tanh(xa) * torch.sigmoid(xb)
+
+        # for skip connection
+        s = self.conv1x1_skip(x)
+
+        # for residual connection
+        x = (self.conv1x1_out(x) + residual) * math.sqrt(0.5)
+
+        return x, s
+
+
+class HiFiGANResidualBlock(torch.nn.Module):
+    """Residual block module in HiFiGAN."""
+
+    def __init__(
+        self,
+        kernel_size=3,
+        channels=512,
+        dilations=(1, 3, 5),
+        bias=True,
+        use_additional_convs=True,
+        nonlinear_activation="LeakyReLU",
+        nonlinear_activation_params={"negative_slope": 0.1},
+    ):
+        """Initialize HiFiGANResidualBlock module.
+
+        Args:
+            kernel_size (int): Kernel size of dilation convolution layer.
+            channels (int): Number of channels for convolution layer.
+            dilations (List[int]): List of dilation factors.
+            use_additional_convs (bool): Whether to use additional convolution layers.
+            bias (bool): Whether to add bias parameter in convolution layers.
+            nonlinear_activation (str): Activation function module name.
+            nonlinear_activation_params (dict): Hyperparameters for activation function.
+
+        """
+        super().__init__()
+        self.use_additional_convs = use_additional_convs
+        self.convs1 = torch.nn.ModuleList()
+        if use_additional_convs:
+            self.convs2 = torch.nn.ModuleList()
+        assert kernel_size % 2 == 1, "Kernel size must be odd number."
+        for dilation in dilations:
+            self.convs1 += [
+                torch.nn.Sequential(
+                    getattr(torch.nn, nonlinear_activation)(
+                        **nonlinear_activation_params
+                    ),
+                    torch.nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation,
+                        bias=bias,
+                        padding=(kernel_size - 1) // 2 * dilation,
+                    ),
+                )
+            ]
+            if use_additional_convs:
+                self.convs2 += [
+                    torch.nn.Sequential(
+                        getattr(torch.nn, nonlinear_activation)(
+                            **nonlinear_activation_params
+                        ),
+                        torch.nn.Conv1d(
+                            channels,
+                            channels,
+                            kernel_size,
+                            1,
+                            dilation=1,
+                            bias=bias,
+                            padding=(kernel_size - 1) // 2,
+                        ),
+                    )
+                ]
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input tensor (B, channels, T).
+
+        Returns:
+            Tensor: Output tensor (B, channels, T).
+
+        """
+        for idx in range(len(self.convs1)):
+            xt = self.convs1[idx](x)
+            if self.use_additional_convs:
+                xt = self.convs2[idx](xt)
+            x = xt + x
+        return x

lscodec/layers/residual_stack.py

@@ -0,0 +1,85 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+"""Residual stack module in MelGAN."""
+
+import torch
+
+from lscodec.layers import CausalConv1d
+
+
+class ResidualStack(torch.nn.Module):
+    """Residual stack module introduced in MelGAN."""
+
+    def __init__(
+        self,
+        kernel_size=3,
+        channels=32,
+        dilation=1,
+        bias=True,
+        nonlinear_activation="LeakyReLU",
+        nonlinear_activation_params={"negative_slope": 0.2},
+        pad="ReflectionPad1d",
+        pad_params={},
+        use_causal_conv=False,
+    ):
+        """Initialize ResidualStack module.
+
+        Args:
+            kernel_size (int): Kernel size of dilation convolution layer.
+            channels (int): Number of channels of convolution layers.
+            dilation (int): Dilation factor.
+            bias (bool): Whether to add bias parameter in convolution layers.
+            nonlinear_activation (str): Activation function module name.
+            nonlinear_activation_params (dict): Hyperparameters for activation function.
+            pad (str): Padding function module name before dilated convolution layer.
+            pad_params (dict): Hyperparameters for padding function.
+            use_causal_conv (bool): Whether to use causal convolution.
+
+        """
+        super(ResidualStack, self).__init__()
+
+        # defile residual stack part
+        if not use_causal_conv:
+            assert (kernel_size - 1) % 2 == 0, "Not support even number kernel size."
+            self.stack = torch.nn.Sequential(
+                getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params),
+                getattr(torch.nn, pad)((kernel_size - 1) // 2 * dilation, **pad_params),
+                torch.nn.Conv1d(
+                    channels, channels, kernel_size, dilation=dilation, bias=bias
+                ),
+                getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params),
+                torch.nn.Conv1d(channels, channels, 1, bias=bias),
+            )
+        else:
+            self.stack = torch.nn.Sequential(
+                getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params),
+                CausalConv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    dilation=dilation,
+                    bias=bias,
+                    pad=pad,
+                    pad_params=pad_params,
+                ),
+                getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params),
+                torch.nn.Conv1d(channels, channels, 1, bias=bias),
+            )
+
+        # defile extra layer for skip connection
+        self.skip_layer = torch.nn.Conv1d(channels, channels, 1, bias=bias)
+
+    def forward(self, c):
+        """Calculate forward propagation.
+
+        Args:
+            c (Tensor): Input tensor (B, channels, T).
+
+        Returns:
+            Tensor: Output tensor (B, chennels, T).
+
+        """
+        return self.stack(c) + self.skip_layer(c)

lscodec/layers/tade_res_block.py

@@ -0,0 +1,160 @@
+# Copyright 2021 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+"""StyleMelGAN's TADEResBlock Modules."""
+
+from functools import partial
+
+import torch
+
+
+class TADELayer(torch.nn.Module):
+    """TADE Layer module."""
+
+    def __init__(
+        self,
+        in_channels=64,
+        aux_channels=80,
+        kernel_size=9,
+        bias=True,
+        upsample_factor=2,
+        upsample_mode="nearest",
+    ):
+        """Initilize TADE layer."""
+        super().__init__()
+        self.norm = torch.nn.InstanceNorm1d(in_channels)
+        self.aux_conv = torch.nn.Sequential(
+            torch.nn.Conv1d(
+                aux_channels,
+                in_channels,
+                kernel_size,
+                1,
+                bias=bias,
+                padding=(kernel_size - 1) // 2,
+            ),
+            # NOTE(kan-bayashi): Use non-linear activation?
+        )
+        self.gated_conv = torch.nn.Sequential(
+            torch.nn.Conv1d(
+                in_channels,
+                in_channels * 2,
+                kernel_size,
+                1,
+                bias=bias,
+                padding=(kernel_size - 1) // 2,
+            ),
+            # NOTE(kan-bayashi): Use non-linear activation?
+        )
+        self.upsample = torch.nn.Upsample(
+            scale_factor=upsample_factor, mode=upsample_mode
+        )
+
+    def forward(self, x, c):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input tensor (B, in_channels, T).
+            c (Tensor): Auxiliary input tensor (B, aux_channels, T').
+
+        Returns:
+            Tensor: Output tensor (B, in_channels, T * in_upsample_factor).
+            Tensor: Upsampled aux tensor (B, in_channels, T * aux_upsample_factor).
+
+        """
+        x = self.norm(x)
+        c = self.upsample(c)
+        c = self.aux_conv(c)
+        cg = self.gated_conv(c)
+        cg1, cg2 = cg.split(cg.size(1) // 2, dim=1)
+        # NOTE(kan-bayashi): Use upsample for noise input here?
+        y = cg1 * self.upsample(x) + cg2
+        # NOTE(kan-bayashi): Return upsampled aux here?
+        return y, c
+
+
+class TADEResBlock(torch.nn.Module):
+    """TADEResBlock module."""
+
+    def __init__(
+        self,
+        in_channels=64,
+        aux_channels=80,
+        kernel_size=9,
+        dilation=2,
+        bias=True,
+        upsample_factor=2,
+        upsample_mode="nearest",
+        gated_function="softmax",
+    ):
+        """Initialize TADEResBlock module."""
+        super().__init__()
+        self.tade1 = TADELayer(
+            in_channels=in_channels,
+            aux_channels=aux_channels,
+            kernel_size=kernel_size,
+            bias=bias,
+            # NOTE(kan-bayashi): Use upsample in the first TADE layer?
+            upsample_factor=1,
+            upsample_mode=upsample_mode,
+        )
+        self.gated_conv1 = torch.nn.Conv1d(
+            in_channels,
+            in_channels * 2,
+            kernel_size,
+            1,
+            bias=bias,
+            padding=(kernel_size - 1) // 2,
+        )
+        self.tade2 = TADELayer(
+            in_channels=in_channels,
+            aux_channels=in_channels,
+            kernel_size=kernel_size,
+            bias=bias,
+            upsample_factor=upsample_factor,
+            upsample_mode=upsample_mode,
+        )
+        self.gated_conv2 = torch.nn.Conv1d(
+            in_channels,
+            in_channels * 2,
+            kernel_size,
+            1,
+            bias=bias,
+            dilation=dilation,
+            padding=(kernel_size - 1) // 2 * dilation,
+        )
+        self.upsample = torch.nn.Upsample(
+            scale_factor=upsample_factor, mode=upsample_mode
+        )
+        if gated_function == "softmax":
+            self.gated_function = partial(torch.softmax, dim=1)
+        elif gated_function == "sigmoid":
+            self.gated_function = torch.sigmoid
+        else:
+            raise ValueError(f"{gated_function} is not supported.")
+
+    def forward(self, x, c):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input tensor (B, in_channels, T).
+            c (Tensor): Auxiliary input tensor (B, aux_channels, T').
+
+        Returns:
+            Tensor: Output tensor (B, in_channels, T * in_upsample_factor).
+            Tensor: Upsampled auxirialy tensor (B, in_channels, T * in_upsample_factor).
+
+        """
+        residual = x
+
+        x, c = self.tade1(x, c)
+        x = self.gated_conv1(x)
+        xa, xb = x.split(x.size(1) // 2, dim=1)
+        x = self.gated_function(xa) * torch.tanh(xb)
+
+        x, c = self.tade2(x, c)
+        x = self.gated_conv2(x)
+        xa, xb = x.split(x.size(1) // 2, dim=1)
+        x = self.gated_function(xa) * torch.tanh(xb)
+
+        # NOTE(kan-bayashi): Return upsampled aux here?
+        return self.upsample(residual) + x, c

lscodec/layers/upsample.py

@@ -0,0 +1,194 @@
+# -*- coding: utf-8 -*-
+
+"""Upsampling module.
+
+This code is modified from https://github.com/r9y9/wavenet_vocoder.
+
+"""
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from lscodec.layers import Conv1d
+
+
+class Stretch2d(torch.nn.Module):
+    """Stretch2d module."""
+
+    def __init__(self, x_scale, y_scale, mode="nearest"):
+        """Initialize Stretch2d module.
+
+        Args:
+            x_scale (int): X scaling factor (Time axis in spectrogram).
+            y_scale (int): Y scaling factor (Frequency axis in spectrogram).
+            mode (str): Interpolation mode.
+
+        """
+        super(Stretch2d, self).__init__()
+        self.x_scale = x_scale
+        self.y_scale = y_scale
+        self.mode = mode
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input tensor (B, C, F, T).
+
+        Returns:
+            Tensor: Interpolated tensor (B, C, F * y_scale, T * x_scale),
+
+        """
+        return F.interpolate(
+            x, scale_factor=(self.y_scale, self.x_scale), mode=self.mode
+        )
+
+
+class Conv2d(torch.nn.Conv2d):
+    """Conv2d module with customized initialization."""
+
+    def __init__(self, *args, **kwargs):
+        """Initialize Conv2d module."""
+        super(Conv2d, self).__init__(*args, **kwargs)
+
+    def reset_parameters(self):
+        """Reset parameters."""
+        self.weight.data.fill_(1.0 / np.prod(self.kernel_size))
+        if self.bias is not None:
+            torch.nn.init.constant_(self.bias, 0.0)
+
+
+class UpsampleNetwork(torch.nn.Module):
+    """Upsampling network module."""
+
+    def __init__(
+        self,
+        upsample_scales,
+        nonlinear_activation=None,
+        nonlinear_activation_params={},
+        interpolate_mode="nearest",
+        freq_axis_kernel_size=1,
+        use_causal_conv=False,
+    ):
+        """Initialize upsampling network module.
+
+        Args:
+            upsample_scales (list): List of upsampling scales.
+            nonlinear_activation (str): Activation function name.
+            nonlinear_activation_params (dict): Arguments for specified activation function.
+            interpolate_mode (str): Interpolation mode.
+            freq_axis_kernel_size (int): Kernel size in the direction of frequency axis.
+
+        """
+        super(UpsampleNetwork, self).__init__()
+        self.use_causal_conv = use_causal_conv
+        self.up_layers = torch.nn.ModuleList()
+        for scale in upsample_scales:
+            # interpolation layer
+            stretch = Stretch2d(scale, 1, interpolate_mode)
+            self.up_layers += [stretch]
+
+            # conv layer
+            assert (
+                freq_axis_kernel_size - 1
+            ) % 2 == 0, "Not support even number freq axis kernel size."
+            freq_axis_padding = (freq_axis_kernel_size - 1) // 2
+            kernel_size = (freq_axis_kernel_size, scale * 2 + 1)
+            if use_causal_conv:
+                padding = (freq_axis_padding, scale * 2)
+            else:
+                padding = (freq_axis_padding, scale)
+            conv = Conv2d(1, 1, kernel_size=kernel_size, padding=padding, bias=False)
+            self.up_layers += [conv]
+
+            # nonlinear
+            if nonlinear_activation is not None:
+                nonlinear = getattr(torch.nn, nonlinear_activation)(
+                    **nonlinear_activation_params
+                )
+                self.up_layers += [nonlinear]
+
+    def forward(self, c):
+        """Calculate forward propagation.
+
+        Args:
+            c : Input tensor (B, C, T).
+
+        Returns:
+            Tensor: Upsampled tensor (B, C, T'), where T' = T * prod(upsample_scales).
+
+        """
+        c = c.unsqueeze(1)  # (B, 1, C, T)
+        for f in self.up_layers:
+            if self.use_causal_conv and isinstance(f, Conv2d):
+                c = f(c)[..., : c.size(-1)]
+            else:
+                c = f(c)
+        return c.squeeze(1)  # (B, C, T')
+
+
+class ConvInUpsampleNetwork(torch.nn.Module):
+    """Convolution + upsampling network module."""
+
+    def __init__(
+        self,
+        upsample_scales,
+        nonlinear_activation=None,
+        nonlinear_activation_params={},
+        interpolate_mode="nearest",
+        freq_axis_kernel_size=1,
+        aux_channels=80,
+        aux_context_window=0,
+        use_causal_conv=False,
+    ):
+        """Initialize convolution + upsampling network module.
+
+        Args:
+            upsample_scales (list): List of upsampling scales.
+            nonlinear_activation (str): Activation function name.
+            nonlinear_activation_params (dict): Arguments for specified activation function.
+            mode (str): Interpolation mode.
+            freq_axis_kernel_size (int): Kernel size in the direction of frequency axis.
+            aux_channels (int): Number of channels of pre-convolutional layer.
+            aux_context_window (int): Context window size of the pre-convolutional layer.
+            use_causal_conv (bool): Whether to use causal structure.
+
+        """
+        super(ConvInUpsampleNetwork, self).__init__()
+        self.aux_context_window = aux_context_window
+        self.use_causal_conv = use_causal_conv and aux_context_window > 0
+        # To capture wide-context information in conditional features
+        kernel_size = (
+            aux_context_window + 1 if use_causal_conv else 2 * aux_context_window + 1
+        )
+        # NOTE(kan-bayashi): Here do not use padding because the input is already padded
+        self.conv_in = Conv1d(
+            aux_channels, aux_channels, kernel_size=kernel_size, bias=False
+        )
+        self.upsample = UpsampleNetwork(
+            upsample_scales=upsample_scales,
+            nonlinear_activation=nonlinear_activation,
+            nonlinear_activation_params=nonlinear_activation_params,
+            interpolate_mode=interpolate_mode,
+            freq_axis_kernel_size=freq_axis_kernel_size,
+            use_causal_conv=use_causal_conv,
+        )
+
+    def forward(self, c):
+        """Calculate forward propagation.
+
+        Args:
+            c : Input tensor (B, C, T').
+
+        Returns:
+            Tensor: Upsampled tensor (B, C, T),
+                where T = (T' - aux_context_window * 2) * prod(upsample_scales).
+
+        Note:
+            The length of inputs considers the context window size.
+
+        """
+        c_ = self.conv_in(c)
+        c = c_[:, :, : -self.aux_context_window] if self.use_causal_conv else c_
+        return self.upsample(c)

lscodec/models/__init__.py

@@ -0,0 +1,3 @@
+from .hifigan import *  # NOQA
+from .melgan import *  # NOQA
+from .lscodec import *

lscodec/models/conformer/__init__.py

@@ -0,0 +1 @@
+"""Initialize sub package."""

lscodec/models/conformer/argument.py

@@ -0,0 +1,87 @@
+# Copyright 2020 Hirofumi Inaguma
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Conformer common arguments."""
+
+
+from distutils.util import strtobool
+import logging
+
+
+def add_arguments_conformer_common(group):
+    """Add Transformer common arguments."""
+    group.add_argument(
+        "--transformer-encoder-pos-enc-layer-type",
+        type=str,
+        default="abs_pos",
+        choices=["abs_pos", "scaled_abs_pos", "rel_pos"],
+        help="Transformer encoder positional encoding layer type",
+    )
+    group.add_argument(
+        "--transformer-encoder-activation-type",
+        type=str,
+        default="swish",
+        choices=["relu", "hardtanh", "selu", "swish"],
+        help="Transformer encoder activation function type",
+    )
+    group.add_argument(
+        "--macaron-style",
+        default=False,
+        type=strtobool,
+        help="Whether to use macaron style for positionwise layer",
+    )
+    # Attention
+    group.add_argument(
+        "--zero-triu",
+        default=False,
+        type=strtobool,
+        help="If true, zero the uppper triangular part of attention matrix.",
+    )
+    # Relative positional encoding
+    group.add_argument(
+        "--rel-pos-type",
+        type=str,
+        default="legacy",
+        choices=["legacy", "latest"],
+        help="Whether to use the latest relative positional encoding or the legacy one."
+        "The legacy relative positional encoding will be deprecated in the future."
+        "More Details can be found in https://github.com/espnet/espnet/pull/2816.",
+    )
+    # CNN module
+    group.add_argument(
+        "--use-cnn-module",
+        default=False,
+        type=strtobool,
+        help="Use convolution module or not",
+    )
+    group.add_argument(
+        "--cnn-module-kernel",
+        default=31,
+        type=int,
+        help="Kernel size of convolution module.",
+    )
+    return group
+
+
+def verify_rel_pos_type(args):
+    """Verify the relative positional encoding type for compatibility.
+
+    Args:
+        args (Namespace): original arguments
+    Returns:
+        args (Namespace): modified arguments
+    """
+    rel_pos_type = getattr(args, "rel_pos_type", None)
+    if rel_pos_type is None or rel_pos_type == "legacy":
+        if args.transformer_encoder_pos_enc_layer_type == "rel_pos":
+            args.transformer_encoder_pos_enc_layer_type = "legacy_rel_pos"
+            logging.warning(
+                "Using legacy_rel_pos and it will be deprecated in the future."
+            )
+        if args.transformer_encoder_selfattn_layer_type == "rel_selfattn":
+            args.transformer_encoder_selfattn_layer_type = "legacy_rel_selfattn"
+            logging.warning(
+                "Using legacy_rel_selfattn and it will be deprecated in the future."
+            )
+
+    return args

lscodec/models/conformer/attention.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Multi-Head Attention layer definition."""
+
+import math
+
+import numpy
+import torch
+from torch import nn
+
+
+class MultiHeadedAttention(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, n_head, n_feat, dropout_rate):
+        """Construct an MultiHeadedAttention object."""
+        super(MultiHeadedAttention, self).__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self, query, key, value):
+        """Transform query, key and value.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+
+        Returns:
+            torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k).
+            torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k).
+            torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k).
+
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(self, value, scores, mask):
+        """Compute attention context vector.
+
+        Args:
+            value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k).
+            scores (torch.Tensor): Attention score (#batch, n_head, time1, time2).
+            mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Transformed value (#batch, time1, d_model)
+                weighted by the attention score (#batch, time1, time2).
+
+        """
+        n_batch = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            min_value = float(
+                numpy.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min
+            )
+            scores = scores.masked_fill(mask, min_value)
+            self.attn = torch.softmax(scores, dim=-1).masked_fill(
+                mask, 0.0
+            )  # (batch, head, time1, time2)
+        else:
+            self.attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(self.attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = (
+            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
+        )  # (batch, time1, d_model)
+
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(self, query, key, value, mask):
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask)
+
+
+class LegacyRelPositionMultiHeadedAttention(MultiHeadedAttention):
+    """Multi-Head Attention layer with relative position encoding (old version).
+
+    Details can be found in https://github.com/espnet/espnet/pull/2816.
+
+    Paper: https://arxiv.org/abs/1901.02860
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+        zero_triu (bool): Whether to zero the upper triangular part of attention matrix.
+
+    """
+
+    def __init__(self, n_head, n_feat, dropout_rate, zero_triu=False):
+        """Construct an RelPositionMultiHeadedAttention object."""
+        super().__init__(n_head, n_feat, dropout_rate)
+        self.zero_triu = zero_triu
+        # linear transformation for positional encoding
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        # these two learnable bias are used in matrix c and matrix d
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        torch.nn.init.xavier_uniform_(self.pos_bias_u)
+        torch.nn.init.xavier_uniform_(self.pos_bias_v)
+
+    def rel_shift(self, x):
+        """Compute relative positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, head, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor.
+
+        """
+        zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype)
+        x_padded = torch.cat([zero_pad, x], dim=-1)
+
+        x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))
+        x = x_padded[:, :, 1:].view_as(x)
+
+        if self.zero_triu:
+            ones = torch.ones((x.size(2), x.size(3)))
+            x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :]
+
+        return x
+
+    def forward(self, query, key, value, pos_emb, mask):
+        """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            pos_emb (torch.Tensor): Positional embedding tensor (#batch, time1, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
+
+        n_batch_pos = pos_emb.size(0)
+        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
+        p = p.transpose(1, 2)  # (batch, head, time1, d_k)
+
+        # (batch, head, time1, d_k)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        # (batch, head, time1, d_k)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+        # compute attention score
+        # first compute matrix a and matrix c
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        # (batch, head, time1, time2)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+
+        # compute matrix b and matrix d
+        # (batch, head, time1, time1)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        matrix_bd = self.rel_shift(matrix_bd)
+
+        scores = (matrix_ac + matrix_bd) / math.sqrt(
+            self.d_k
+        )  # (batch, head, time1, time2)
+
+        return self.forward_attention(v, scores, mask)
+
+
+class RelPositionMultiHeadedAttention(MultiHeadedAttention):
+    """Multi-Head Attention layer with relative position encoding (new implementation).
+
+    Details can be found in https://github.com/espnet/espnet/pull/2816.
+
+    Paper: https://arxiv.org/abs/1901.02860
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+        zero_triu (bool): Whether to zero the upper triangular part of attention matrix.
+
+    """
+
+    def __init__(self, n_head, n_feat, dropout_rate, zero_triu=False):
+        """Construct an RelPositionMultiHeadedAttention object."""
+        super().__init__(n_head, n_feat, dropout_rate)
+        self.zero_triu = zero_triu
+        # linear transformation for positional encoding
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        # these two learnable bias are used in matrix c and matrix d
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        torch.nn.init.xavier_uniform_(self.pos_bias_u)
+        torch.nn.init.xavier_uniform_(self.pos_bias_v)
+
+    def rel_shift(self, x):
+        """Compute relative positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1).
+            time1 means the length of query vector.
+
+        Returns:
+            torch.Tensor: Output tensor.
+
+        """
+        zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype)
+        x_padded = torch.cat([zero_pad, x], dim=-1)
+
+        x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))
+        x = x_padded[:, :, 1:].view_as(x)[
+            :, :, :, : x.size(-1) // 2 + 1
+        ]  # only keep the positions from 0 to time2
+
+        if self.zero_triu:
+            ones = torch.ones((x.size(2), x.size(3)), device=x.device)
+            x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :]
+
+        return x
+
+    def forward(self, query, key, value, pos_emb, mask):
+        """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            pos_emb (torch.Tensor): Positional embedding tensor
+                (#batch, 2*time1-1, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
+
+        n_batch_pos = pos_emb.size(0)
+        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
+        p = p.transpose(1, 2)  # (batch, head, 2*time1-1, d_k)
+
+        # (batch, head, time1, d_k)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        # (batch, head, time1, d_k)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+        # compute attention score
+        # first compute matrix a and matrix c
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        # (batch, head, time1, time2)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+
+        # compute matrix b and matrix d
+        # (batch, head, time1, 2*time1-1)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        matrix_bd = self.rel_shift(matrix_bd)
+
+        scores = (matrix_ac + matrix_bd) / math.sqrt(
+            self.d_k
+        )  # (batch, head, time1, time2)
+
+        return self.forward_attention(v, scores, mask)

lscodec/models/conformer/contextual_block_encoder_layer.py

@@ -0,0 +1,309 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Sat Aug 21 16:57:31 2021.
+
+@author: Keqi Deng (UCAS)
+"""
+
+from espnet.nets.pytorch_backend.transformer.layer_norm import LayerNorm
+import torch
+from torch import nn
+
+
+class ContextualBlockEncoderLayer(nn.Module):
+    """Contexutal Block Encoder layer module.
+
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` instance
+            can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        feed_forward_macaron (torch.nn.Module): Additional feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        conv_module (torch.nn.Module): Convolution module instance.
+            `ConvlutionModule` instance can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        total_layer_num (int): Total number of layers
+        normalize_before (bool): Whether to use layer_norm before the first block.
+        concat_after (bool): Whether to concat attention layer's input and output.
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied. i.e. x -> x + att(x)
+
+    """
+
+    def __init__(
+        self,
+        size,
+        self_attn,
+        feed_forward,
+        feed_forward_macaron,
+        conv_module,
+        dropout_rate,
+        total_layer_num,
+        normalize_before=True,
+        concat_after=False,
+    ):
+        """Construct an EncoderLayer object."""
+        super(ContextualBlockEncoderLayer, self).__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.feed_forward_macaron = feed_forward_macaron
+        self.conv_module = conv_module
+        self.norm1 = LayerNorm(size)
+        self.norm2 = LayerNorm(size)
+        if feed_forward_macaron is not None:
+            self.norm_ff_macaron = LayerNorm(size)
+            self.ff_scale = 0.5
+        else:
+            self.ff_scale = 1.0
+        if self.conv_module is not None:
+            self.norm_conv = LayerNorm(size)  # for the CNN module
+            self.norm_final = LayerNorm(size)  # for the final output of the block
+        self.dropout = nn.Dropout(dropout_rate)
+        self.size = size
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        self.total_layer_num = total_layer_num
+        if self.concat_after:
+            self.concat_linear = nn.Linear(size + size, size)
+
+    def forward(
+        self,
+        x,
+        mask,
+        infer_mode=False,
+        past_ctx=None,
+        next_ctx=None,
+        is_short_segment=False,
+        layer_idx=0,
+        cache=None,
+    ):
+        """Calculate forward propagation."""
+        if self.training or not infer_mode:
+            return self.forward_train(x, mask, past_ctx, next_ctx, layer_idx, cache)
+        else:
+            return self.forward_infer(
+                x, mask, past_ctx, next_ctx, is_short_segment, layer_idx, cache
+            )
+
+    def forward_train(
+        self, x, mask, past_ctx=None, next_ctx=None, layer_idx=0, cache=None
+    ):
+        """Compute encoded features.
+
+        Args:
+            x_input (torch.Tensor): Input tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            past_ctx (torch.Tensor): Previous contexutal vector
+            next_ctx (torch.Tensor): Next contexutal vector
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+            cur_ctx (torch.Tensor): Current contexutal vector
+            next_ctx (torch.Tensor): Next contexutal vector
+            layer_idx (int): layer index number
+
+        """
+        nbatch = x.size(0)
+        nblock = x.size(1)
+
+        if past_ctx is not None:
+            if next_ctx is None:
+                # store all context vectors in one tensor
+                next_ctx = past_ctx.new_zeros(
+                    nbatch, nblock, self.total_layer_num, x.size(-1)
+                )
+            else:
+                x[:, :, 0] = past_ctx[:, :, layer_idx]
+
+        # reshape ( nbatch, nblock, block_size + 2, dim )
+        #     -> ( nbatch * nblock, block_size + 2, dim )
+        x = x.view(-1, x.size(-2), x.size(-1))
+        if mask is not None:
+            mask = mask.view(-1, mask.size(-2), mask.size(-1))
+
+        # whether to use macaron style
+        if self.feed_forward_macaron is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_ff_macaron(x)
+            x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
+            if not self.normalize_before:
+                x = self.norm_ff_macaron(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm1(x)
+
+        if cache is None:
+            x_q = x
+        else:
+            assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size)
+            x_q = x[:, -1:, :]
+            residual = residual[:, -1:, :]
+            mask = None if mask is None else mask[:, -1:, :]
+
+        if self.concat_after:
+            x_concat = torch.cat((x, self.self_attn(x_q, x, x, mask)), dim=-1)
+            x = residual + self.concat_linear(x_concat)
+        else:
+            x = residual + self.dropout(self.self_attn(x_q, x, x, mask))
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        # convolution module
+        if self.conv_module is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_conv(x)
+            x = residual + self.dropout(self.conv_module(x))
+            if not self.normalize_before:
+                x = self.norm_conv(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        x = residual + self.ff_scale * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        if self.conv_module is not None:
+            x = self.norm_final(x)
+
+        if cache is not None:
+            x = torch.cat([cache, x], dim=1)
+
+        layer_idx += 1
+        # reshape ( nbatch * nblock, block_size + 2, dim )
+        #       -> ( nbatch, nblock, block_size + 2, dim )
+        x = x.view(nbatch, -1, x.size(-2), x.size(-1)).squeeze(1)
+        if mask is not None:
+            mask = mask.view(nbatch, -1, mask.size(-2), mask.size(-1)).squeeze(1)
+
+        if next_ctx is not None and layer_idx < self.total_layer_num:
+            next_ctx[:, 0, layer_idx, :] = x[:, 0, -1, :]
+            next_ctx[:, 1:, layer_idx, :] = x[:, 0:-1, -1, :]
+
+        return x, mask, False, next_ctx, next_ctx, layer_idx
+
+    def forward_infer(
+        self,
+        x,
+        mask,
+        past_ctx=None,
+        next_ctx=None,
+        is_short_segment=False,
+        layer_idx=0,
+        cache=None,
+    ):
+        """Compute encoded features.
+
+        Args:
+            x_input (torch.Tensor): Input tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            past_ctx (torch.Tensor): Previous contexutal vector
+            next_ctx (torch.Tensor): Next contexutal vector
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+            cur_ctx (torch.Tensor): Current contexutal vector
+            next_ctx (torch.Tensor): Next contexutal vector
+            layer_idx (int): layer index number
+
+        """
+        nbatch = x.size(0)
+        nblock = x.size(1)
+        # if layer_idx == 0, next_ctx has to be None
+        if layer_idx == 0:
+            assert next_ctx is None
+            next_ctx = x.new_zeros(nbatch, self.total_layer_num, x.size(-1))
+
+        # reshape ( nbatch, nblock, block_size + 2, dim )
+        #     -> ( nbatch * nblock, block_size + 2, dim )
+        x = x.view(-1, x.size(-2), x.size(-1))
+        if mask is not None:
+            mask = mask.view(-1, mask.size(-2), mask.size(-1))
+
+        # whether to use macaron style
+        if self.feed_forward_macaron is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_ff_macaron(x)
+            x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
+            if not self.normalize_before:
+                x = self.norm_ff_macaron(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm1(x)
+
+        if cache is None:
+            x_q = x
+        else:
+            assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size)
+            x_q = x[:, -1:, :]
+            residual = residual[:, -1:, :]
+            mask = None if mask is None else mask[:, -1:, :]
+
+        if self.concat_after:
+            x_concat = torch.cat((x, self.self_attn(x_q, x, x, mask)), dim=-1)
+            x = residual + self.concat_linear(x_concat)
+        else:
+            x = residual + self.dropout(self.self_attn(x_q, x, x, mask))
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        # convolution module
+        if self.conv_module is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_conv(x)
+            x = residual + self.dropout(self.conv_module(x))
+            if not self.normalize_before:
+                x = self.norm_conv(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        x = residual + self.ff_scale * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        if self.conv_module is not None:
+            x = self.norm_final(x)
+
+        if cache is not None:
+            x = torch.cat([cache, x], dim=1)
+
+        # reshape ( nbatch * nblock, block_size + 2, dim )
+        #       -> ( nbatch, nblock, block_size + 2, dim )
+        x = x.view(nbatch, nblock, x.size(-2), x.size(-1))
+        if mask is not None:
+            mask = mask.view(nbatch, nblock, mask.size(-2), mask.size(-1))
+
+        # Propagete context information (the last frame of each block)
+        # to the first frame
+        # of the next block
+
+        if not is_short_segment:
+            if past_ctx is None:
+                # First block of an utterance
+                x[:, 0, 0, :] = x[:, 0, -1, :]
+            else:
+                x[:, 0, 0, :] = past_ctx[:, layer_idx, :]
+            if nblock > 1:
+                x[:, 1:, 0, :] = x[:, 0:-1, -1, :]
+            next_ctx[:, layer_idx, :] = x[:, -1, -1, :]
+        else:
+            next_ctx = None
+
+        return x, mask, True, past_ctx, next_ctx, is_short_segment, layer_idx + 1

lscodec/models/conformer/convolution.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Johns Hopkins University (Shinji Watanabe)
+#                Northwestern Polytechnical University (Pengcheng Guo)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""ConvolutionModule definition."""
+
+from torch import nn
+
+
+class ConvolutionModule(nn.Module):
+    """ConvolutionModule in Conformer model.
+
+    Args:
+        channels (int): The number of channels of conv layers.
+        kernel_size (int): Kernerl size of conv layers.
+
+    """
+
+    def __init__(self, channels, kernel_size, activation=nn.ReLU(), bias=True):
+        """Construct an ConvolutionModule object."""
+        super(ConvolutionModule, self).__init__()
+        # kernerl_size should be a odd number for 'SAME' padding
+        assert (kernel_size - 1) % 2 == 0
+
+        self.pointwise_conv1 = nn.Conv1d(
+            channels,
+            2 * channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias,
+        )
+        self.depthwise_conv = nn.Conv1d(
+            channels,
+            channels,
+            kernel_size,
+            stride=1,
+            padding=(kernel_size - 1) // 2,
+            groups=channels,
+            bias=bias,
+        )
+        self.norm = nn.BatchNorm1d(channels)
+        self.pointwise_conv2 = nn.Conv1d(
+            channels,
+            channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias,
+        )
+        self.activation = activation
+
+    def forward(self, x):
+        """Compute convolution module.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, channels).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, channels).
+
+        """
+        # exchange the temporal dimension and the feature dimension
+        x = x.transpose(1, 2)
+
+        # GLU mechanism
+        x = self.pointwise_conv1(x)  # (batch, 2*channel, dim)
+        x = nn.functional.glu(x, dim=1)  # (batch, channel, dim)
+
+        # 1D Depthwise Conv
+        x = self.depthwise_conv(x)
+        x = self.activation(self.norm(x))
+
+        x = self.pointwise_conv2(x)
+
+        return x.transpose(1, 2)

lscodec/models/conformer/decoder.py

@@ -0,0 +1,247 @@
+# Copyright 2020 Johns Hopkins University (Shinji Watanabe)
+#                Northwestern Polytechnical University (Pengcheng Guo)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Encoder definition."""
+
+import logging
+import torch
+
+from lscodec.models.conformer.convolution import ConvolutionModule
+from lscodec.models.conformer.decoder_layer import DecoderLayer
+from lscodec.models.conformer.nets_utils import get_activation
+from lscodec.models.conformer.vgg2l import VGG2L
+from lscodec.models.conformer.attention import (
+    MultiHeadedAttention,  # noqa: H301
+    RelPositionMultiHeadedAttention,  # noqa: H301
+    LegacyRelPositionMultiHeadedAttention,  # noqa: H301
+)
+from lscodec.models.conformer.embedding import (
+    PositionalEncoding,  # noqa: H301
+    ScaledPositionalEncoding,  # noqa: H301
+    RelPositionalEncoding,  # noqa: H301
+    LegacyRelPositionalEncoding,  # noqa: H301
+)
+from lscodec.models.conformer.layer_norm import LayerNorm
+from lscodec.models.conformer.multi_layer_conv import Conv1dLinear, MultiLayeredConv1d
+from lscodec.models.conformer.positionwise_feed_forward import (
+    PositionwiseFeedForward,  # noqa: H301
+)
+from lscodec.models.conformer.repeat import repeat
+from lscodec.models.conformer.subsampling import Conv2dSubsampling
+
+
+class Decoder(torch.nn.Module):
+    """Conformer encoder module.
+
+    Args:
+        idim (int): Input dimension.
+        attention_dim (int): Dimention of attention.
+        attention_heads (int): The number of heads of multi head attention.
+        linear_units (int): The number of units of position-wise feed forward.
+        num_blocks (int): The number of decoder blocks.
+        dropout_rate (float): Dropout rate.
+        positional_dropout_rate (float): Dropout rate after adding positional encoding.
+        attention_dropout_rate (float): Dropout rate in attention.
+        input_layer (Union[str, torch.nn.Module]): Input layer type.
+        normalize_before (bool): Whether to use layer_norm before the first block.
+        concat_after (bool): Whether to concat attention layer's input and output.
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied. i.e. x -> x + att(x)
+        positionwise_layer_type (str): "linear", "conv1d", or "conv1d-linear".
+        positionwise_conv_kernel_size (int): Kernel size of positionwise conv1d layer.
+        macaron_style (bool): Whether to use macaron style for positionwise layer.
+        pos_enc_layer_type (str): Encoder positional encoding layer type.
+        selfattention_layer_type (str): Encoder attention layer type.
+        activation_type (str): Encoder activation function type.
+        use_cnn_module (bool): Whether to use convolution module.
+        zero_triu (bool): Whether to zero the upper triangular part of attention matrix.
+        cnn_module_kernel (int): Kernerl size of convolution module.
+        padding_idx (int): Padding idx for input_layer=embed.
+
+    """
+
+    def __init__(
+        self,
+        idim,
+        attention_dim=256,
+        attention_heads=4,
+        linear_units=2048,
+        num_blocks=6,
+        dropout_rate=0.1,
+        positional_dropout_rate=0.1,
+        attention_dropout_rate=0.0,
+        input_layer="conv2d",
+        normalize_before=True,
+        concat_after=False,
+        positionwise_layer_type="linear",
+        positionwise_conv_kernel_size=1,
+        macaron_style=False,
+        pos_enc_layer_type="abs_pos",
+        selfattention_layer_type="selfattn",
+        activation_type="swish",
+        use_cnn_module=False,
+        zero_triu=False,
+        cnn_module_kernel=31,
+        padding_idx=-1,
+    ):
+        """Construct an Encoder object."""
+        super(Decoder, self).__init__()
+
+        activation = get_activation(activation_type)
+        if pos_enc_layer_type == "abs_pos":
+            pos_enc_class = PositionalEncoding
+        elif pos_enc_layer_type == "scaled_abs_pos":
+            pos_enc_class = ScaledPositionalEncoding
+        elif pos_enc_layer_type == "rel_pos":
+            assert selfattention_layer_type == "rel_selfattn"
+            pos_enc_class = RelPositionalEncoding
+        elif pos_enc_layer_type == "legacy_rel_pos":
+            pos_enc_class = LegacyRelPositionalEncoding
+            assert selfattention_layer_type == "legacy_rel_selfattn"
+        else:
+            raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type)
+
+        self.conv_subsampling_factor = 1
+        if input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(idim, attention_dim),
+                torch.nn.LayerNorm(attention_dim),
+                torch.nn.Dropout(dropout_rate),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d":
+            self.embed = Conv2dSubsampling(
+                idim,
+                attention_dim,
+                dropout_rate,
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+            self.conv_subsampling_factor = 4
+        elif input_layer == "vgg2l":
+            self.embed = VGG2L(idim, attention_dim)
+            self.conv_subsampling_factor = 4
+        elif input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(idim, attention_dim, padding_idx=padding_idx),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        elif isinstance(input_layer, torch.nn.Module):
+            self.embed = torch.nn.Sequential(
+                input_layer,
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        elif input_layer is None:
+            self.embed = torch.nn.Sequential(
+                pos_enc_class(attention_dim, positional_dropout_rate)
+            )
+        else:
+            raise ValueError("unknown input_layer: " + input_layer)
+        self.normalize_before = normalize_before
+
+        # self-attention module definition
+        if selfattention_layer_type == "selfattn":
+            logging.info("encoder self-attention layer type = self-attention")
+            encoder_selfattn_layer = MultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                attention_dim,
+                attention_dropout_rate,
+            )
+        elif selfattention_layer_type == "legacy_rel_selfattn":
+            assert pos_enc_layer_type == "legacy_rel_pos"
+            encoder_selfattn_layer = LegacyRelPositionMultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                attention_dim,
+                attention_dropout_rate,
+            )
+        elif selfattention_layer_type == "rel_selfattn":
+            logging.info("encoder self-attention layer type = relative self-attention")
+            assert pos_enc_layer_type == "rel_pos"
+            encoder_selfattn_layer = RelPositionMultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                attention_dim,
+                attention_dropout_rate,
+                zero_triu,
+            )
+        else:
+            raise ValueError("unknown encoder_attn_layer: " + selfattention_layer_type)
+
+        # feed-forward module definition
+        if positionwise_layer_type == "linear":
+            positionwise_layer = PositionwiseFeedForward
+            positionwise_layer_args = (
+                attention_dim,
+                linear_units,
+                dropout_rate,
+                activation,
+            )
+        elif positionwise_layer_type == "conv1d":
+            positionwise_layer = MultiLayeredConv1d
+            positionwise_layer_args = (
+                attention_dim,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        elif positionwise_layer_type == "conv1d-linear":
+            positionwise_layer = Conv1dLinear
+            positionwise_layer_args = (
+                attention_dim,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        else:
+            raise NotImplementedError("Support only linear or conv1d.")
+
+        # convolution module definition
+        convolution_layer = ConvolutionModule
+        convolution_layer_args = (attention_dim, cnn_module_kernel, activation)
+
+        self.decoders = repeat(
+            num_blocks,
+            lambda lnum: DecoderLayer(
+                attention_dim,
+                encoder_selfattn_layer(*encoder_selfattn_layer_args),
+                MultiHeadedAttention(
+                    attention_heads, attention_dim, attention_dropout_rate
+                ),
+                positionwise_layer(*positionwise_layer_args),
+                positionwise_layer(*positionwise_layer_args) if macaron_style else None,
+                convolution_layer(*convolution_layer_args) if use_cnn_module else None,
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if self.normalize_before:
+            self.after_norm = LayerNorm(attention_dim)
+
+    def forward(self, xs, masks, memory, memory_mask):
+        """Encode input sequence.
+
+        Args:
+            xs (torch.Tensor): Input tensor (#batch, time, idim).
+            masks (torch.Tensor): Mask tensor (#batch, time).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, attention_dim).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+        if isinstance(self.embed, (Conv2dSubsampling, VGG2L)):
+            xs, masks = self.embed(xs, masks)
+        else:
+            xs = self.embed(xs)
+
+        xs, masks, memory, memory_mask = self.decoders(xs, masks, memory, memory_mask)
+        if isinstance(xs, tuple):
+            xs = xs[0]
+
+        if self.normalize_before:
+            xs = self.after_norm(xs)
+        return xs, masks

lscodec/models/conformer/decoder_layer.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Johns Hopkins University (Shinji Watanabe)
+#                Northwestern Polytechnical University (Pengcheng Guo)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Encoder self-attention layer definition."""
+
+import torch
+
+from torch import nn
+
+from lscodec.models.conformer.layer_norm import LayerNorm
+
+
+class DecoderLayer(nn.Module):
+    """Encoder layer module.
+
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` instance
+            can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        feed_forward_macaron (torch.nn.Module): Additional feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        conv_module (torch.nn.Module): Convolution module instance.
+            `ConvlutionModule` instance can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool): Whether to use layer_norm before the first block.
+        concat_after (bool): Whether to concat attention layer's input and output.
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied. i.e. x -> x + att(x)
+
+    """
+
+    def __init__(
+        self,
+        size,
+        self_attn,
+        src_attn,
+        feed_forward,
+        feed_forward_macaron,
+        conv_module,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+    ):
+        """Construct an EncoderLayer object."""
+        super(DecoderLayer, self).__init__()
+        self.self_attn = self_attn
+        self.src_attn = src_attn
+        self.feed_forward = feed_forward
+        self.feed_forward_macaron = feed_forward_macaron
+        self.conv_module = conv_module
+        self.norm_ff = LayerNorm(size)  # for the FNN module
+        self.norm_mha = LayerNorm(size)  # for the MHA module
+        self.norm2 = LayerNorm(size)  # for the MHA module
+        if feed_forward_macaron is not None:
+            self.norm_ff_macaron = LayerNorm(size)
+            self.ff_scale = 0.5
+        else:
+            self.ff_scale = 1.0
+        if self.conv_module is not None:
+            self.norm_conv = LayerNorm(size)  # for the CNN module
+            self.norm_final = LayerNorm(size)  # for the final output of the block
+        self.dropout = nn.Dropout(dropout_rate)
+        self.size = size
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear = nn.Linear(size + size, size)
+            self.concat_linear2 = nn.Linear(size + size, size)
+
+    def forward(self, x_input, mask, memory, memory_mask, cache=None):
+        """Compute encoded features.
+
+        Args:
+            x_input (Union[Tuple, torch.Tensor]): Input tensor w/ or w/o pos emb.
+                - w/ pos emb: Tuple of tensors [(#batch, time, size), (1, time, size)].
+                - w/o pos emb: Tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+        if isinstance(x_input, tuple):
+            x, pos_emb = x_input[0], x_input[1]
+        else:
+            x, pos_emb = x_input, None
+
+        # whether to use macaron style
+        if self.feed_forward_macaron is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_ff_macaron(x)
+            x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
+            if not self.normalize_before:
+                x = self.norm_ff_macaron(x)
+
+        # multi-headed self-attention module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_mha(x)
+
+        if cache is None:
+            x_q = x
+        else:
+            assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size)
+            x_q = x[:, -1:, :]
+            residual = residual[:, -1:, :]
+            mask = None if mask is None else mask[:, -1:, :]
+
+        if pos_emb is not None:
+            x_att = self.self_attn(x_q, x, x, pos_emb, mask)
+        else:
+            x_att = self.self_attn(x_q, x, x, mask)
+
+        if self.concat_after:
+            x_concat = torch.cat((x, x_att), dim=-1)
+            x = residual + self.concat_linear(x_concat)
+        else:
+            x = residual + self.dropout(x_att)
+        if not self.normalize_before:
+            x = self.norm_mha(x)
+
+        # cross attention
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        if self.concat_after:
+            x_concat = torch.cat(
+                (x, self.src_attn(x, memory, memory, memory_mask)), dim=-1
+            )
+            x = residual + self.concat_linear2(x_concat)
+        else:
+            x = residual + self.dropout(self.src_attn(x, memory, memory, memory_mask))
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        # convolution module
+        if self.conv_module is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_conv(x)
+            x = residual + self.dropout(self.conv_module(x))
+            if not self.normalize_before:
+                x = self.norm_conv(x)
+
+        # feed forward module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_ff(x)
+        x = residual + self.ff_scale * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm_ff(x)
+
+        if self.conv_module is not None:
+            x = self.norm_final(x)
+
+        if cache is not None:
+            x = torch.cat([cache, x], dim=1)
+
+        if pos_emb is not None:
+            return (x, pos_emb), mask, memory, memory_mask
+
+        return x, mask, memory, memory_mask

lscodec/models/conformer/embedding.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Positional Encoding Module."""
+
+import math
+
+import torch
+
+
+def _pre_hook(
+    state_dict,
+    prefix,
+    local_metadata,
+    strict,
+    missing_keys,
+    unexpected_keys,
+    error_msgs,
+):
+    """Perform pre-hook in load_state_dict for backward compatibility.
+
+    Note:
+        We saved self.pe until v.0.5.2 but we have omitted it later.
+        Therefore, we remove the item "pe" from `state_dict` for backward compatibility.
+
+    """
+    k = prefix + "pe"
+    if k in state_dict:
+        state_dict.pop(k)
+
+
+class PositionalEncoding(torch.nn.Module):
+    """Positional encoding.
+
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+        reverse (bool): Whether to reverse the input position. Only for
+        the class LegacyRelPositionalEncoding. We remove it in the current
+        class RelPositionalEncoding.
+
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000, reverse=False):
+        """Construct an PositionalEncoding object."""
+        super(PositionalEncoding, self).__init__()
+        self.d_model = d_model
+        self.reverse = reverse
+        self.xscale = math.sqrt(self.d_model)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+        self._register_load_state_dict_pre_hook(_pre_hook)
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            if self.pe.size(1) >= x.size(1):
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        pe = torch.zeros(x.size(1), self.d_model)
+        if self.reverse:
+            position = torch.arange(
+                x.size(1) - 1, -1, -1.0, dtype=torch.float32
+            ).unsqueeze(1)
+        else:
+            position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+    def forward(self, x: torch.Tensor):
+        """Add positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+
+        """
+        self.extend_pe(x)
+        x = x * self.xscale + self.pe[:, : x.size(1)]
+        return self.dropout(x)
+
+
+class ScaledPositionalEncoding(PositionalEncoding):
+    """Scaled positional encoding module.
+
+    See Sec. 3.2  https://arxiv.org/abs/1809.08895
+
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000):
+        """Initialize class."""
+        super().__init__(d_model=d_model, dropout_rate=dropout_rate, max_len=max_len)
+        self.alpha = torch.nn.Parameter(torch.tensor(1.0))
+
+    def reset_parameters(self):
+        """Reset parameters."""
+        self.alpha.data = torch.tensor(1.0)
+
+    def forward(self, x):
+        """Add positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+
+        """
+        self.extend_pe(x)
+        x = x + self.alpha * self.pe[:, : x.size(1)]
+        return self.dropout(x)
+
+
+class LegacyRelPositionalEncoding(PositionalEncoding):
+    """Relative positional encoding module (old version).
+
+    Details can be found in https://github.com/espnet/espnet/pull/2816.
+
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000):
+        """Initialize class."""
+        super().__init__(
+            d_model=d_model,
+            dropout_rate=dropout_rate,
+            max_len=max_len,
+            reverse=True,
+        )
+
+    def forward(self, x):
+        """Compute positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+            torch.Tensor: Positional embedding tensor (1, time, `*`).
+
+        """
+        self.extend_pe(x)
+        x = x * self.xscale
+        pos_emb = self.pe[:, : x.size(1)]
+        return self.dropout(x), self.dropout(pos_emb)
+
+
+class RelPositionalEncoding(torch.nn.Module):
+    """Relative positional encoding module (new implementation).
+
+    Details can be found in https://github.com/espnet/espnet/pull/2816.
+
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000):
+        """Construct an PositionalEncoding object."""
+        super(RelPositionalEncoding, self).__init__()
+        self.d_model = d_model
+        self.xscale = math.sqrt(self.d_model)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            # self.pe contains both positive and negative parts
+            # the length of self.pe is 2 * input_len - 1
+            if self.pe.size(1) >= x.size(1) * 2 - 1:
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        # Suppose `i` means to the position of query vecotr and `j` means the
+        # position of key vector. We use position relative positions when keys
+        # are to the left (i>j) and negative relative positions otherwise (i<j).
+        pe_positive = torch.zeros(x.size(1), self.d_model)
+        pe_negative = torch.zeros(x.size(1), self.d_model)
+        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe_positive[:, 0::2] = torch.sin(position * div_term)
+        pe_positive[:, 1::2] = torch.cos(position * div_term)
+        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
+        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
+
+        # Reserve the order of positive indices and concat both positive and
+        # negative indices. This is used to support the shifting trick
+        # as in https://arxiv.org/abs/1901.02860
+        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
+        pe_negative = pe_negative[1:].unsqueeze(0)
+        pe = torch.cat([pe_positive, pe_negative], dim=1)
+        self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+    def forward(self, x: torch.Tensor):
+        """Add positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+
+        """
+        self.extend_pe(x)
+        x = x * self.xscale
+        pos_emb = self.pe[
+            :,
+            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),
+        ]
+        return self.dropout(x), self.dropout(pos_emb)

lscodec/models/conformer/encoder.py

@@ -0,0 +1,247 @@
+# Copyright 2020 Johns Hopkins University (Shinji Watanabe)
+#                Northwestern Polytechnical University (Pengcheng Guo)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Encoder definition."""
+
+import logging
+import torch
+from lscodec.models.conformer.convolution import ConvolutionModule
+from lscodec.models.conformer.encoder_layer import EncoderLayer
+from lscodec.utils.espnet_utils import get_activation
+from lscodec.models.conformer.vgg2l import VGG2L
+from lscodec.models.conformer.attention import (
+    MultiHeadedAttention,  # noqa: H301
+    RelPositionMultiHeadedAttention,  # noqa: H301
+    LegacyRelPositionMultiHeadedAttention,  # noqa: H301
+)
+from lscodec.models.conformer.embedding import (
+    PositionalEncoding,  # noqa: H301
+    ScaledPositionalEncoding,  # noqa: H301
+    RelPositionalEncoding,  # noqa: H301
+    LegacyRelPositionalEncoding,  # noqa: H301
+)
+from lscodec.models.conformer.layer_norm import LayerNorm
+from lscodec.models.conformer.multi_layer_conv import Conv1dLinear
+from lscodec.models.conformer.multi_layer_conv import MultiLayeredConv1d
+from lscodec.models.conformer.positionwise_feed_forward import (
+    PositionwiseFeedForward,  # noqa: H301
+)
+from lscodec.models.conformer.repeat import repeat
+from lscodec.models.conformer.subsampling import Conv2dSubsampling
+
+
+
+class Encoder(torch.nn.Module):
+    """Conformer encoder module.
+
+    Args:
+        idim (int): Input dimension.
+        attention_dim (int): Dimention of attention.
+        attention_heads (int): The number of heads of multi head attention.
+        linear_units (int): The number of units of position-wise feed forward.
+        num_blocks (int): The number of decoder blocks.
+        dropout_rate (float): Dropout rate.
+        positional_dropout_rate (float): Dropout rate after adding positional encoding.
+        attention_dropout_rate (float): Dropout rate in attention.
+        input_layer (Union[str, torch.nn.Module]): Input layer type.
+        normalize_before (bool): Whether to use layer_norm before the first block.
+        concat_after (bool): Whether to concat attention layer's input and output.
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied. i.e. x -> x + att(x)
+        positionwise_layer_type (str): "linear", "conv1d", or "conv1d-linear".
+        positionwise_conv_kernel_size (int): Kernel size of positionwise conv1d layer.
+        macaron_style (bool): Whether to use macaron style for positionwise layer.
+        pos_enc_layer_type (str): Encoder positional encoding layer type.
+        selfattention_layer_type (str): Encoder attention layer type.
+        activation_type (str): Encoder activation function type.
+        use_cnn_module (bool): Whether to use convolution module.
+        zero_triu (bool): Whether to zero the upper triangular part of attention matrix.
+        cnn_module_kernel (int): Kernerl size of convolution module.
+        padding_idx (int): Padding idx for input_layer=embed.
+
+    """
+
+    def __init__(
+        self,
+        idim,
+        attention_dim=256,
+        attention_heads=4,
+        linear_units=2048,
+        num_blocks=6,
+        dropout_rate=0.1,
+        positional_dropout_rate=0.1,
+        attention_dropout_rate=0.0,
+        input_layer="conv2d",
+        normalize_before=True,
+        concat_after=False,
+        positionwise_layer_type="linear",
+        positionwise_conv_kernel_size=1,
+        macaron_style=False,
+        pos_enc_layer_type="abs_pos",
+        selfattention_layer_type="selfattn",
+        activation_type="swish",
+        use_cnn_module=False,
+        zero_triu=False,
+        cnn_module_kernel=31,
+        padding_idx=-1,
+    ):
+        """Construct an Encoder object."""
+        super(Encoder, self).__init__()
+
+        activation = get_activation(activation_type)
+        if pos_enc_layer_type == "abs_pos":
+            pos_enc_class = PositionalEncoding
+        elif pos_enc_layer_type == "scaled_abs_pos":
+            pos_enc_class = ScaledPositionalEncoding
+        elif pos_enc_layer_type == "rel_pos":
+            assert selfattention_layer_type == "rel_selfattn"
+            pos_enc_class = RelPositionalEncoding
+        elif pos_enc_layer_type == "legacy_rel_pos":
+            pos_enc_class = LegacyRelPositionalEncoding
+            assert selfattention_layer_type == "legacy_rel_selfattn"
+        else:
+            raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type)
+
+        self.conv_subsampling_factor = 1
+        if input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(idim, attention_dim),
+                torch.nn.LayerNorm(attention_dim),
+                torch.nn.Dropout(dropout_rate),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d":
+            self.embed = Conv2dSubsampling(
+                idim,
+                attention_dim,
+                dropout_rate,
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+            self.conv_subsampling_factor = 4
+        elif input_layer == "vgg2l":
+            self.embed = VGG2L(idim, attention_dim)
+            self.conv_subsampling_factor = 4
+        elif input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(idim, attention_dim, padding_idx=padding_idx),
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        elif isinstance(input_layer, torch.nn.Module):
+            self.embed = torch.nn.Sequential(
+                input_layer,
+                pos_enc_class(attention_dim, positional_dropout_rate),
+            )
+        elif input_layer is None:
+            self.embed = torch.nn.Sequential(
+                pos_enc_class(attention_dim, positional_dropout_rate)
+            )
+        elif input_layer == "no_pos_emb":
+            self.embed = torch.nn.Identity()
+        else:
+            raise ValueError("unknown input_layer: " + input_layer)
+        self.normalize_before = normalize_before
+
+        # self-attention module definition
+        if selfattention_layer_type == "selfattn":
+            logging.info("encoder self-attention layer type = self-attention")
+            encoder_selfattn_layer = MultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                attention_dim,
+                attention_dropout_rate,
+            )
+        elif selfattention_layer_type == "legacy_rel_selfattn":
+            assert pos_enc_layer_type == "legacy_rel_pos"
+            encoder_selfattn_layer = LegacyRelPositionMultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                attention_dim,
+                attention_dropout_rate,
+            )
+        elif selfattention_layer_type == "rel_selfattn":
+            logging.info("encoder self-attention layer type = relative self-attention")
+            assert pos_enc_layer_type == "rel_pos"
+            encoder_selfattn_layer = RelPositionMultiHeadedAttention
+            encoder_selfattn_layer_args = (
+                attention_heads,
+                attention_dim,
+                attention_dropout_rate,
+                zero_triu,
+            )
+        else:
+            raise ValueError("unknown encoder_attn_layer: " + selfattention_layer_type)
+
+        # feed-forward module definition
+        if positionwise_layer_type == "linear":
+            positionwise_layer = PositionwiseFeedForward
+            positionwise_layer_args = (
+                attention_dim,
+                linear_units,
+                dropout_rate,
+                activation,
+            )
+        elif positionwise_layer_type == "conv1d":
+            positionwise_layer = MultiLayeredConv1d
+            positionwise_layer_args = (
+                attention_dim,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        elif positionwise_layer_type == "conv1d-linear":
+            positionwise_layer = Conv1dLinear
+            positionwise_layer_args = (
+                attention_dim,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        else:
+            raise NotImplementedError("Support only linear or conv1d.")
+
+        # convolution module definition
+        convolution_layer = ConvolutionModule
+        convolution_layer_args = (attention_dim, cnn_module_kernel, activation)
+
+        self.encoders = repeat(
+            num_blocks,
+            lambda lnum: EncoderLayer(
+                attention_dim,
+                encoder_selfattn_layer(*encoder_selfattn_layer_args),
+                positionwise_layer(*positionwise_layer_args),
+                positionwise_layer(*positionwise_layer_args) if macaron_style else None,
+                convolution_layer(*convolution_layer_args) if use_cnn_module else None,
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if self.normalize_before:
+            self.after_norm = LayerNorm(attention_dim)
+
+    def forward(self, xs, masks):
+        """Encode input sequence.
+
+        Args:
+            xs (torch.Tensor): Input tensor (#batch, time, idim).
+            masks (torch.Tensor): Mask tensor (#batch, time).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, attention_dim).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+        if isinstance(self.embed, (Conv2dSubsampling, VGG2L)):
+            xs, masks = self.embed(xs, masks)
+        else:
+            xs = self.embed(xs)
+
+        xs, masks = self.encoders(xs, masks)
+        if isinstance(xs, tuple):
+            xs = xs[0]
+
+        if self.normalize_before:
+            xs = self.after_norm(xs)
+        return xs, masks

lscodec/models/conformer/encoder_layer.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Johns Hopkins University (Shinji Watanabe)
+#                Northwestern Polytechnical University (Pengcheng Guo)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Encoder self-attention layer definition."""
+
+import torch
+
+from torch import nn
+
+from lscodec.models.conformer.layer_norm import LayerNorm
+
+
+class EncoderLayer(nn.Module):
+    """Encoder layer module.
+
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` instance
+            can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        feed_forward_macaron (torch.nn.Module): Additional feed-forward module instance.
+            `PositionwiseFeedForward`, `MultiLayeredConv1d`, or `Conv1dLinear` instance
+            can be used as the argument.
+        conv_module (torch.nn.Module): Convolution module instance.
+            `ConvlutionModule` instance can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool): Whether to use layer_norm before the first block.
+        concat_after (bool): Whether to concat attention layer's input and output.
+            if True, additional linear will be applied.
+            i.e. x -> x + linear(concat(x, att(x)))
+            if False, no additional linear will be applied. i.e. x -> x + att(x)
+
+    """
+
+    def __init__(
+        self,
+        size,
+        self_attn,
+        feed_forward,
+        feed_forward_macaron,
+        conv_module,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+    ):
+        """Construct an EncoderLayer object."""
+        super(EncoderLayer, self).__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.feed_forward_macaron = feed_forward_macaron
+        self.conv_module = conv_module
+        self.norm_ff = LayerNorm(size)  # for the FNN module
+        self.norm_mha = LayerNorm(size)  # for the MHA module
+        if feed_forward_macaron is not None:
+            self.norm_ff_macaron = LayerNorm(size)
+            self.ff_scale = 0.5
+        else:
+            self.ff_scale = 1.0
+        if self.conv_module is not None:
+            self.norm_conv = LayerNorm(size)  # for the CNN module
+            self.norm_final = LayerNorm(size)  # for the final output of the block
+        self.dropout = nn.Dropout(dropout_rate)
+        self.size = size
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear = nn.Linear(size + size, size)
+
+    def forward(self, x_input, mask, cache=None):
+        """Compute encoded features.
+
+        Args:
+            x_input (Union[Tuple, torch.Tensor]): Input tensor w/ or w/o pos emb.
+                - w/ pos emb: Tuple of tensors [(#batch, time, size), (1, time, size)].
+                - w/o pos emb: Tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+        if isinstance(x_input, tuple):
+            x, pos_emb = x_input[0], x_input[1]
+        else:
+            x, pos_emb = x_input, None
+
+        # whether to use macaron style
+        if self.feed_forward_macaron is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_ff_macaron(x)
+            x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
+            if not self.normalize_before:
+                x = self.norm_ff_macaron(x)
+
+        # multi-headed self-attention module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_mha(x)
+
+        if cache is None:
+            x_q = x
+        else:
+            assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size)
+            x_q = x[:, -1:, :]
+            residual = residual[:, -1:, :]
+            mask = None if mask is None else mask[:, -1:, :]
+
+        if pos_emb is not None:
+            x_att = self.self_attn(x_q, x, x, pos_emb, mask)
+        else:
+            x_att = self.self_attn(x_q, x, x, mask)
+
+        if self.concat_after:
+            x_concat = torch.cat((x, x_att), dim=-1)
+            x = residual + self.concat_linear(x_concat)
+        else:
+            x = residual + self.dropout(x_att)
+        if not self.normalize_before:
+            x = self.norm_mha(x)
+
+        # convolution module
+        if self.conv_module is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_conv(x)
+            x = residual + self.dropout(self.conv_module(x))
+            if not self.normalize_before:
+                x = self.norm_conv(x)
+
+        # feed forward module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_ff(x)
+        x = residual + self.ff_scale * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm_ff(x)
+
+        if self.conv_module is not None:
+            x = self.norm_final(x)
+
+        if cache is not None:
+            x = torch.cat([cache, x], dim=1)
+
+        if pos_emb is not None:
+            return (x, pos_emb), mask
+
+        return x, mask

lscodec/models/conformer/layer_norm.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Layer normalization module."""
+
+import torch
+
+
+class LayerNorm(torch.nn.LayerNorm):
+    """Layer normalization module.
+
+    Args:
+        nout (int): Output dim size.
+        dim (int): Dimension to be normalized.
+
+    """
+
+    def __init__(self, nout, dim=-1):
+        """Construct an LayerNorm object."""
+        super(LayerNorm, self).__init__(nout, eps=1e-12)
+        self.dim = dim
+
+    def forward(self, x):
+        """Apply layer normalization.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Normalized tensor.
+
+        """
+        if self.dim == -1:
+            return super(LayerNorm, self).forward(x)
+        return (
+            super(LayerNorm, self)
+            .forward(x.transpose(self.dim, -1))
+            .transpose(self.dim, -1)
+        )

lscodec/models/conformer/multi_layer_conv.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Tomoki Hayashi
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Layer modules for FFT block in FastSpeech (Feed-forward Transformer)."""
+
+import torch
+
+
+class MultiLayeredConv1d(torch.nn.Module):
+    """Multi-layered conv1d for Transformer block.
+
+    This is a module of multi-leyered conv1d designed
+    to replace positionwise feed-forward network
+    in Transforner block, which is introduced in
+    `FastSpeech: Fast, Robust and Controllable Text to Speech`_.
+
+    .. _`FastSpeech: Fast, Robust and Controllable Text to Speech`:
+        https://arxiv.org/pdf/1905.09263.pdf
+
+    """
+
+    def __init__(self, in_chans, hidden_chans, kernel_size, dropout_rate):
+        """Initialize MultiLayeredConv1d module.
+
+        Args:
+            in_chans (int): Number of input channels.
+            hidden_chans (int): Number of hidden channels.
+            kernel_size (int): Kernel size of conv1d.
+            dropout_rate (float): Dropout rate.
+
+        """
+        super(MultiLayeredConv1d, self).__init__()
+        self.w_1 = torch.nn.Conv1d(
+            in_chans,
+            hidden_chans,
+            kernel_size,
+            stride=1,
+            padding=(kernel_size - 1) // 2,
+        )
+        self.w_2 = torch.nn.Conv1d(
+            hidden_chans,
+            in_chans,
+            kernel_size,
+            stride=1,
+            padding=(kernel_size - 1) // 2,
+        )
+        self.dropout = torch.nn.Dropout(dropout_rate)
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (torch.Tensor): Batch of input tensors (B, T, in_chans).
+
+        Returns:
+            torch.Tensor: Batch of output tensors (B, T, hidden_chans).
+
+        """
+        x = torch.relu(self.w_1(x.transpose(-1, 1))).transpose(-1, 1)
+        return self.w_2(self.dropout(x).transpose(-1, 1)).transpose(-1, 1)
+
+
+class Conv1dLinear(torch.nn.Module):
+    """Conv1D + Linear for Transformer block.
+
+    A variant of MultiLayeredConv1d, which replaces second conv-layer to linear.
+
+    """
+
+    def __init__(self, in_chans, hidden_chans, kernel_size, dropout_rate):
+        """Initialize Conv1dLinear module.
+
+        Args:
+            in_chans (int): Number of input channels.
+            hidden_chans (int): Number of hidden channels.
+            kernel_size (int): Kernel size of conv1d.
+            dropout_rate (float): Dropout rate.
+
+        """
+        super(Conv1dLinear, self).__init__()
+        self.w_1 = torch.nn.Conv1d(
+            in_chans,
+            hidden_chans,
+            kernel_size,
+            stride=1,
+            padding=(kernel_size - 1) // 2,
+        )
+        self.w_2 = torch.nn.Linear(hidden_chans, in_chans)
+        self.dropout = torch.nn.Dropout(dropout_rate)
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (torch.Tensor): Batch of input tensors (B, T, in_chans).
+
+        Returns:
+            torch.Tensor: Batch of output tensors (B, T, hidden_chans).
+
+        """
+        x = torch.relu(self.w_1(x.transpose(-1, 1))).transpose(-1, 1)
+        return self.w_2(self.dropout(x))

lscodec/models/conformer/nets_utils.py

@@ -0,0 +1,498 @@
+# -*- coding: utf-8 -*-
+
+"""Network related utility tools."""
+
+import logging
+from typing import Dict
+
+import numpy as np
+import torch
+
+
+def to_device(m, x):
+    """Send tensor into the device of the module.
+
+    Args:
+        m (torch.nn.Module): Torch module.
+        x (Tensor): Torch tensor.
+
+    Returns:
+        Tensor: Torch tensor located in the same place as torch module.
+
+    """
+    if isinstance(m, torch.nn.Module):
+        device = next(m.parameters()).device
+    elif isinstance(m, torch.Tensor):
+        device = m.device
+    else:
+        raise TypeError(
+            "Expected torch.nn.Module or torch.tensor, " f"bot got: {type(m)}"
+        )
+    return x.to(device)
+
+
+def pad_list(xs, pad_value):
+    """Perform padding for the list of tensors.
+
+    Args:
+        xs (List): List of Tensors [(T_1, `*`), (T_2, `*`), ..., (T_B, `*`)].
+        pad_value (float): Value for padding.
+
+    Returns:
+        Tensor: Padded tensor (B, Tmax, `*`).
+
+    Examples:
+        >>> x = [torch.ones(4), torch.ones(2), torch.ones(1)]
+        >>> x
+        [tensor([1., 1., 1., 1.]), tensor([1., 1.]), tensor([1.])]
+        >>> pad_list(x, 0)
+        tensor([[1., 1., 1., 1.],
+                [1., 1., 0., 0.],
+                [1., 0., 0., 0.]])
+
+    """
+    n_batch = len(xs)
+    max_len = max(x.size(0) for x in xs)
+    pad = xs[0].new(n_batch, max_len, *xs[0].size()[1:]).fill_(pad_value)
+
+    for i in range(n_batch):
+        pad[i, : xs[i].size(0)] = xs[i]
+
+    return pad
+
+
+def make_pad_mask(lengths, xs=None, length_dim=-1):
+    """Make mask tensor containing indices of padded part.
+
+    Args:
+        lengths (LongTensor or List): Batch of lengths (B,).
+        xs (Tensor, optional): The reference tensor.
+            If set, masks will be the same shape as this tensor.
+        length_dim (int, optional): Dimension indicator of the above tensor.
+            See the example.
+
+    Returns:
+        Tensor: Mask tensor containing indices of padded part.
+                dtype=torch.uint8 in PyTorch 1.2-
+                dtype=torch.bool in PyTorch 1.2+ (including 1.2)
+
+    Examples:
+        With only lengths.
+
+        >>> lengths = [5, 3, 2]
+        >>> make_non_pad_mask(lengths)
+        masks = [[0, 0, 0, 0 ,0],
+                 [0, 0, 0, 1, 1],
+                 [0, 0, 1, 1, 1]]
+
+        With the reference tensor.
+
+        >>> xs = torch.zeros((3, 2, 4))
+        >>> make_pad_mask(lengths, xs)
+        tensor([[[0, 0, 0, 0],
+                 [0, 0, 0, 0]],
+                [[0, 0, 0, 1],
+                 [0, 0, 0, 1]],
+                [[0, 0, 1, 1],
+                 [0, 0, 1, 1]]], dtype=torch.uint8)
+        >>> xs = torch.zeros((3, 2, 6))
+        >>> make_pad_mask(lengths, xs)
+        tensor([[[0, 0, 0, 0, 0, 1],
+                 [0, 0, 0, 0, 0, 1]],
+                [[0, 0, 0, 1, 1, 1],
+                 [0, 0, 0, 1, 1, 1]],
+                [[0, 0, 1, 1, 1, 1],
+                 [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8)
+
+        With the reference tensor and dimension indicator.
+
+        >>> xs = torch.zeros((3, 6, 6))
+        >>> make_pad_mask(lengths, xs, 1)
+        tensor([[[0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [1, 1, 1, 1, 1, 1]],
+                [[0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1]],
+                [[0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1]]], dtype=torch.uint8)
+        >>> make_pad_mask(lengths, xs, 2)
+        tensor([[[0, 0, 0, 0, 0, 1],
+                 [0, 0, 0, 0, 0, 1],
+                 [0, 0, 0, 0, 0, 1],
+                 [0, 0, 0, 0, 0, 1],
+                 [0, 0, 0, 0, 0, 1],
+                 [0, 0, 0, 0, 0, 1]],
+                [[0, 0, 0, 1, 1, 1],
+                 [0, 0, 0, 1, 1, 1],
+                 [0, 0, 0, 1, 1, 1],
+                 [0, 0, 0, 1, 1, 1],
+                 [0, 0, 0, 1, 1, 1],
+                 [0, 0, 0, 1, 1, 1]],
+                [[0, 0, 1, 1, 1, 1],
+                 [0, 0, 1, 1, 1, 1],
+                 [0, 0, 1, 1, 1, 1],
+                 [0, 0, 1, 1, 1, 1],
+                 [0, 0, 1, 1, 1, 1],
+                 [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8)
+
+    """
+    if length_dim == 0:
+        raise ValueError("length_dim cannot be 0: {}".format(length_dim))
+
+    if not isinstance(lengths, list):
+        lengths = lengths.tolist()
+    bs = int(len(lengths))
+    if xs is None:
+        maxlen = int(max(lengths))
+    else:
+        maxlen = xs.size(length_dim)
+
+    seq_range = torch.arange(0, maxlen, dtype=torch.int64)
+    seq_range_expand = seq_range.unsqueeze(0).expand(bs, maxlen)
+    seq_length_expand = seq_range_expand.new(lengths).unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+
+    if xs is not None:
+        assert xs.size(0) == bs, (xs.size(0), bs)
+
+        if length_dim < 0:
+            length_dim = xs.dim() + length_dim
+        # ind = (:, None, ..., None, :, , None, ..., None)
+        ind = tuple(
+            slice(None) if i in (0, length_dim) else None for i in range(xs.dim())
+        )
+        mask = mask[ind].expand_as(xs).to(xs.device)
+    return mask
+
+
+def make_non_pad_mask(lengths, xs=None, length_dim=-1):
+    """Make mask tensor containing indices of non-padded part.
+
+    Args:
+        lengths (LongTensor or List): Batch of lengths (B,).
+        xs (Tensor, optional): The reference tensor.
+            If set, masks will be the same shape as this tensor.
+        length_dim (int, optional): Dimension indicator of the above tensor.
+            See the example.
+
+    Returns:
+        ByteTensor: mask tensor containing indices of padded part.
+                    dtype=torch.uint8 in PyTorch 1.2-
+                    dtype=torch.bool in PyTorch 1.2+ (including 1.2)
+
+    Examples:
+        With only lengths.
+
+        >>> lengths = [5, 3, 2]
+        >>> make_non_pad_mask(lengths)
+        masks = [[1, 1, 1, 1 ,1],
+                 [1, 1, 1, 0, 0],
+                 [1, 1, 0, 0, 0]]
+
+        With the reference tensor.
+
+        >>> xs = torch.zeros((3, 2, 4))
+        >>> make_non_pad_mask(lengths, xs)
+        tensor([[[1, 1, 1, 1],
+                 [1, 1, 1, 1]],
+                [[1, 1, 1, 0],
+                 [1, 1, 1, 0]],
+                [[1, 1, 0, 0],
+                 [1, 1, 0, 0]]], dtype=torch.uint8)
+        >>> xs = torch.zeros((3, 2, 6))
+        >>> make_non_pad_mask(lengths, xs)
+        tensor([[[1, 1, 1, 1, 1, 0],
+                 [1, 1, 1, 1, 1, 0]],
+                [[1, 1, 1, 0, 0, 0],
+                 [1, 1, 1, 0, 0, 0]],
+                [[1, 1, 0, 0, 0, 0],
+                 [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8)
+
+        With the reference tensor and dimension indicator.
+
+        >>> xs = torch.zeros((3, 6, 6))
+        >>> make_non_pad_mask(lengths, xs, 1)
+        tensor([[[1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [0, 0, 0, 0, 0, 0]],
+                [[1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0]],
+                [[1, 1, 1, 1, 1, 1],
+                 [1, 1, 1, 1, 1, 1],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0],
+                 [0, 0, 0, 0, 0, 0]]], dtype=torch.uint8)
+        >>> make_non_pad_mask(lengths, xs, 2)
+        tensor([[[1, 1, 1, 1, 1, 0],
+                 [1, 1, 1, 1, 1, 0],
+                 [1, 1, 1, 1, 1, 0],
+                 [1, 1, 1, 1, 1, 0],
+                 [1, 1, 1, 1, 1, 0],
+                 [1, 1, 1, 1, 1, 0]],
+                [[1, 1, 1, 0, 0, 0],
+                 [1, 1, 1, 0, 0, 0],
+                 [1, 1, 1, 0, 0, 0],
+                 [1, 1, 1, 0, 0, 0],
+                 [1, 1, 1, 0, 0, 0],
+                 [1, 1, 1, 0, 0, 0]],
+                [[1, 1, 0, 0, 0, 0],
+                 [1, 1, 0, 0, 0, 0],
+                 [1, 1, 0, 0, 0, 0],
+                 [1, 1, 0, 0, 0, 0],
+                 [1, 1, 0, 0, 0, 0],
+                 [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8)
+
+    """
+    return ~make_pad_mask(lengths, xs, length_dim)
+
+
+def mask_by_length(xs, lengths, fill=0):
+    """Mask tensor according to length.
+
+    Args:
+        xs (Tensor): Batch of input tensor (B, `*`).
+        lengths (LongTensor or List): Batch of lengths (B,).
+        fill (int or float): Value to fill masked part.
+
+    Returns:
+        Tensor: Batch of masked input tensor (B, `*`).
+
+    Examples:
+        >>> x = torch.arange(5).repeat(3, 1) + 1
+        >>> x
+        tensor([[1, 2, 3, 4, 5],
+                [1, 2, 3, 4, 5],
+                [1, 2, 3, 4, 5]])
+        >>> lengths = [5, 3, 2]
+        >>> mask_by_length(x, lengths)
+        tensor([[1, 2, 3, 4, 5],
+                [1, 2, 3, 0, 0],
+                [1, 2, 0, 0, 0]])
+
+    """
+    assert xs.size(0) == len(lengths)
+    ret = xs.data.new(*xs.size()).fill_(fill)
+    for i, l in enumerate(lengths):
+        ret[i, :l] = xs[i, :l]
+    return ret
+
+
+def th_accuracy(pad_outputs, pad_targets, ignore_label):
+    """Calculate accuracy.
+
+    Args:
+        pad_outputs (Tensor): Prediction tensors (B * Lmax, D).
+        pad_targets (LongTensor): Target label tensors (B, Lmax, D).
+        ignore_label (int): Ignore label id.
+
+    Returns:
+        float: Accuracy value (0.0 - 1.0).
+
+    """
+    pad_pred = pad_outputs.view(
+        pad_targets.size(0), pad_targets.size(1), pad_outputs.size(1)
+    ).argmax(2)
+    mask = pad_targets != ignore_label
+    numerator = torch.sum(
+        pad_pred.masked_select(mask) == pad_targets.masked_select(mask)
+    )
+    denominator = torch.sum(mask)
+    return float(numerator) / float(denominator)
+
+
+def to_torch_tensor(x):
+    """Change to torch.Tensor or ComplexTensor from numpy.ndarray.
+
+    Args:
+        x: Inputs. It should be one of numpy.ndarray, Tensor, ComplexTensor, and dict.
+
+    Returns:
+        Tensor or ComplexTensor: Type converted inputs.
+
+    Examples:
+        >>> xs = np.ones(3, dtype=np.float32)
+        >>> xs = to_torch_tensor(xs)
+        tensor([1., 1., 1.])
+        >>> xs = torch.ones(3, 4, 5)
+        >>> assert to_torch_tensor(xs) is xs
+        >>> xs = {'real': xs, 'imag': xs}
+        >>> to_torch_tensor(xs)
+        ComplexTensor(
+        Real:
+        tensor([1., 1., 1.])
+        Imag;
+        tensor([1., 1., 1.])
+        )
+
+    """
+    # If numpy, change to torch tensor
+    if isinstance(x, np.ndarray):
+        if x.dtype.kind == "c":
+            # Dynamically importing because torch_complex requires python3
+            from torch_complex.tensor import ComplexTensor
+
+            return ComplexTensor(x)
+        else:
+            return torch.from_numpy(x)
+
+    # If {'real': ..., 'imag': ...}, convert to ComplexTensor
+    elif isinstance(x, dict):
+        # Dynamically importing because torch_complex requires python3
+        from torch_complex.tensor import ComplexTensor
+
+        if "real" not in x or "imag" not in x:
+            raise ValueError("has 'real' and 'imag' keys: {}".format(list(x)))
+        # Relative importing because of using python3 syntax
+        return ComplexTensor(x["real"], x["imag"])
+
+    # If torch.Tensor, as it is
+    elif isinstance(x, torch.Tensor):
+        return x
+
+    else:
+        error = (
+            "x must be numpy.ndarray, torch.Tensor or a dict like "
+            "{{'real': torch.Tensor, 'imag': torch.Tensor}}, "
+            "but got {}".format(type(x))
+        )
+        try:
+            from torch_complex.tensor import ComplexTensor
+        except Exception:
+            # If PY2
+            raise ValueError(error)
+        else:
+            # If PY3
+            if isinstance(x, ComplexTensor):
+                return x
+            else:
+                raise ValueError(error)
+
+
+def get_subsample(train_args, mode, arch):
+    """Parse the subsampling factors from the args for the specified `mode` and `arch`.
+
+    Args:
+        train_args: argument Namespace containing options.
+        mode: one of ('asr', 'mt', 'st')
+        arch: one of ('rnn', 'rnn-t', 'rnn_mix', 'rnn_mulenc', 'transformer')
+
+    Returns:
+        np.ndarray / List[np.ndarray]: subsampling factors.
+    """
+    if arch == "transformer":
+        return np.array([1])
+
+    elif mode == "mt" and arch == "rnn":
+        # +1 means input (+1) and layers outputs (train_args.elayer)
+        subsample = np.ones(train_args.elayers + 1, dtype=np.int)
+        logging.warning("Subsampling is not performed for machine translation.")
+        logging.info("subsample: " + " ".join([str(x) for x in subsample]))
+        return subsample
+
+    elif (
+        (mode == "asr" and arch in ("rnn", "rnn-t"))
+        or (mode == "mt" and arch == "rnn")
+        or (mode == "st" and arch == "rnn")
+    ):
+        subsample = np.ones(train_args.elayers + 1, dtype=np.int)
+        if train_args.etype.endswith("p") and not train_args.etype.startswith("vgg"):
+            ss = train_args.subsample.split("_")
+            for j in range(min(train_args.elayers + 1, len(ss))):
+                subsample[j] = int(ss[j])
+        else:
+            logging.warning(
+                "Subsampling is not performed for vgg*. "
+                "It is performed in max pooling layers at CNN."
+            )
+        logging.info("subsample: " + " ".join([str(x) for x in subsample]))
+        return subsample
+
+    elif mode == "asr" and arch == "rnn_mix":
+        subsample = np.ones(
+            train_args.elayers_sd + train_args.elayers + 1, dtype=np.int
+        )
+        if train_args.etype.endswith("p") and not train_args.etype.startswith("vgg"):
+            ss = train_args.subsample.split("_")
+            for j in range(
+                min(train_args.elayers_sd + train_args.elayers + 1, len(ss))
+            ):
+                subsample[j] = int(ss[j])
+        else:
+            logging.warning(
+                "Subsampling is not performed for vgg*. "
+                "It is performed in max pooling layers at CNN."
+            )
+        logging.info("subsample: " + " ".join([str(x) for x in subsample]))
+        return subsample
+
+    elif mode == "asr" and arch == "rnn_mulenc":
+        subsample_list = []
+        for idx in range(train_args.num_encs):
+            subsample = np.ones(train_args.elayers[idx] + 1, dtype=np.int)
+            if train_args.etype[idx].endswith("p") and not train_args.etype[
+                idx
+            ].startswith("vgg"):
+                ss = train_args.subsample[idx].split("_")
+                for j in range(min(train_args.elayers[idx] + 1, len(ss))):
+                    subsample[j] = int(ss[j])
+            else:
+                logging.warning(
+                    "Encoder %d: Subsampling is not performed for vgg*. "
+                    "It is performed in max pooling layers at CNN.",
+                    idx + 1,
+                )
+            logging.info("subsample: " + " ".join([str(x) for x in subsample]))
+            subsample_list.append(subsample)
+        return subsample_list
+
+    else:
+        raise ValueError("Invalid options: mode={}, arch={}".format(mode, arch))
+
+
+def rename_state_dict(
+    old_prefix: str, new_prefix: str, state_dict: Dict[str, torch.Tensor]
+):
+    """Replace keys of old prefix with new prefix in state dict."""
+    # need this list not to break the dict iterator
+    old_keys = [k for k in state_dict if k.startswith(old_prefix)]
+    if len(old_keys) > 0:
+        logging.warning(f"Rename: {old_prefix} -> {new_prefix}")
+    for k in old_keys:
+        v = state_dict.pop(k)
+        new_k = k.replace(old_prefix, new_prefix)
+        state_dict[new_k] = v
+
+
+def get_activation(act):
+    """Return activation function."""
+    # Lazy load to avoid unused import
+    from lscodec.models.conformer.swish import Swish
+
+    activation_funcs = {
+        "hardtanh": torch.nn.Hardtanh,
+        "tanh": torch.nn.Tanh,
+        "relu": torch.nn.ReLU,
+        "selu": torch.nn.SELU,
+        "swish": Swish,
+    }
+
+    return activation_funcs[act]()

lscodec/models/conformer/positionwise_feed_forward.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Positionwise feed forward layer definition."""
+
+import torch
+
+
+class PositionwiseFeedForward(torch.nn.Module):
+    """Positionwise feed forward layer.
+
+    Args:
+        idim (int): Input dimenstion.
+        hidden_units (int): The number of hidden units.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim, hidden_units, dropout_rate, activation=torch.nn.ReLU()):
+        """Construct an PositionwiseFeedForward object."""
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = torch.nn.Linear(idim, hidden_units)
+        self.w_2 = torch.nn.Linear(hidden_units, idim)
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.activation = activation
+
+    def forward(self, x):
+        """Forward funciton."""
+        return self.w_2(self.dropout(self.activation(self.w_1(x))))

lscodec/models/conformer/repeat.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Repeat the same layer definition."""
+
+import torch
+
+
+class MultiSequential(torch.nn.Sequential):
+    """Multi-input multi-output torch.nn.Sequential."""
+
+    def forward(self, *args):
+        """Repeat."""
+        for m in self:
+            args = m(*args)
+        return args
+
+
+def repeat(N, fn):
+    """Repeat module N times.
+
+    Args:
+        N (int): Number of repeat time.
+        fn (Callable): Function to generate module.
+
+    Returns:
+        MultiSequential: Repeated model instance.
+
+    """
+    return MultiSequential(*[fn(n) for n in range(N)])

lscodec/models/conformer/subsampling.py

@@ -0,0 +1,257 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Subsampling layer definition."""
+
+import torch
+
+from lscodec.models.conformer.embedding import PositionalEncoding
+
+
+class TooShortUttError(Exception):
+    """Raised when the utt is too short for subsampling.
+
+    Args:
+        message (str): Message for error catch
+        actual_size (int): the short size that cannot pass the subsampling
+        limit (int): the limit size for subsampling
+
+    """
+
+    def __init__(self, message, actual_size, limit):
+        """Construct a TooShortUttError for error handler."""
+        super().__init__(message)
+        self.actual_size = actual_size
+        self.limit = limit
+
+
+def check_short_utt(ins, size):
+    """Check if the utterance is too short for subsampling."""
+    if isinstance(ins, Conv2dSubsampling2) and size < 3:
+        return True, 3
+    if isinstance(ins, Conv2dSubsampling) and size < 7:
+        return True, 7
+    if isinstance(ins, Conv2dSubsampling6) and size < 11:
+        return True, 11
+    if isinstance(ins, Conv2dSubsampling8) and size < 15:
+        return True, 15
+    return False, -1
+
+
+class Conv2dSubsampling(torch.nn.Module):
+    """Convolutional 2D subsampling (to 1/4 length).
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+        pos_enc (torch.nn.Module): Custom position encoding layer.
+
+    """
+
+    def __init__(self, idim, odim, dropout_rate, pos_enc=None):
+        """Construct an Conv2dSubsampling object."""
+        super(Conv2dSubsampling, self).__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(1, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 3, 2),
+            torch.nn.ReLU(),
+        )
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(odim * (((idim - 1) // 2 - 1) // 2), odim),
+            pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate),
+        )
+
+    def forward(self, x, x_mask):
+        """Subsample x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, time', odim),
+                where time' = time // 4.
+            torch.Tensor: Subsampled mask (#batch, 1, time'),
+                where time' = time // 4.
+
+        """
+        x = x.unsqueeze(1)  # (b, c, t, f)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+        if x_mask is None:
+            return x, None
+        return x, x_mask[:, :, :-2:2][:, :, :-2:2]
+
+    def __getitem__(self, key):
+        """Get item.
+
+        When reset_parameters() is called, if use_scaled_pos_enc is used,
+            return the positioning encoding.
+
+        """
+        if key != -1:
+            raise NotImplementedError("Support only `-1` (for `reset_parameters`).")
+        return self.out[key]
+
+
+class Conv2dSubsampling2(torch.nn.Module):
+    """Convolutional 2D subsampling (to 1/2 length).
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+        pos_enc (torch.nn.Module): Custom position encoding layer.
+
+    """
+
+    def __init__(self, idim, odim, dropout_rate, pos_enc=None):
+        """Construct an Conv2dSubsampling2 object."""
+        super(Conv2dSubsampling2, self).__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(1, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 3, 1),
+            torch.nn.ReLU(),
+        )
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(odim * (((idim - 1) // 2 - 2)), odim),
+            pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate),
+        )
+
+    def forward(self, x, x_mask):
+        """Subsample x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, time', odim),
+                where time' = time // 2.
+            torch.Tensor: Subsampled mask (#batch, 1, time'),
+                where time' = time // 2.
+
+        """
+        x = x.unsqueeze(1)  # (b, c, t, f)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+        if x_mask is None:
+            return x, None
+        return x, x_mask[:, :, :-2:2][:, :, :-2:1]
+
+    def __getitem__(self, key):
+        """Get item.
+
+        When reset_parameters() is called, if use_scaled_pos_enc is used,
+            return the positioning encoding.
+
+        """
+        if key != -1:
+            raise NotImplementedError("Support only `-1` (for `reset_parameters`).")
+        return self.out[key]
+
+
+class Conv2dSubsampling6(torch.nn.Module):
+    """Convolutional 2D subsampling (to 1/6 length).
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+        pos_enc (torch.nn.Module): Custom position encoding layer.
+
+    """
+
+    def __init__(self, idim, odim, dropout_rate, pos_enc=None):
+        """Construct an Conv2dSubsampling6 object."""
+        super(Conv2dSubsampling6, self).__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(1, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 5, 3),
+            torch.nn.ReLU(),
+        )
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(odim * (((idim - 1) // 2 - 2) // 3), odim),
+            pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate),
+        )
+
+    def forward(self, x, x_mask):
+        """Subsample x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, time', odim),
+                where time' = time // 6.
+            torch.Tensor: Subsampled mask (#batch, 1, time'),
+                where time' = time // 6.
+
+        """
+        x = x.unsqueeze(1)  # (b, c, t, f)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+        if x_mask is None:
+            return x, None
+        return x, x_mask[:, :, :-2:2][:, :, :-4:3]
+
+
+class Conv2dSubsampling8(torch.nn.Module):
+    """Convolutional 2D subsampling (to 1/8 length).
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+        pos_enc (torch.nn.Module): Custom position encoding layer.
+
+    """
+
+    def __init__(self, idim, odim, dropout_rate, pos_enc=None):
+        """Construct an Conv2dSubsampling8 object."""
+        super(Conv2dSubsampling8, self).__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(1, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 3, 2),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(odim, odim, 3, 2),
+            torch.nn.ReLU(),
+        )
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(odim * ((((idim - 1) // 2 - 1) // 2 - 1) // 2), odim),
+            pos_enc if pos_enc is not None else PositionalEncoding(odim, dropout_rate),
+        )
+
+    def forward(self, x, x_mask):
+        """Subsample x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: Subsampled tensor (#batch, time', odim),
+                where time' = time // 8.
+            torch.Tensor: Subsampled mask (#batch, 1, time'),
+                where time' = time // 8.
+
+        """
+        x = x.unsqueeze(1)  # (b, c, t, f)
+        x = self.conv(x)
+        b, c, t, f = x.size()
+        x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f))
+        if x_mask is None:
+            return x, None
+        return x, x_mask[:, :, :-2:2][:, :, :-2:2][:, :, :-2:2]

lscodec/models/conformer/swish.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Johns Hopkins University (Shinji Watanabe)
+#                Northwestern Polytechnical University (Pengcheng Guo)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Swish() activation function for Conformer."""
+
+import torch
+
+
+class Swish(torch.nn.Module):
+    """Construct an Swish object."""
+
+    def forward(self, x):
+        """Return Swich activation function."""
+        return x * torch.sigmoid(x)

lscodec/models/conformer/vgg2l.py

@@ -0,0 +1,89 @@
+"""VGG2L module definition for transformer encoder."""
+
+from typing import Tuple
+from typing import Union
+
+import torch
+
+
+class VGG2L(torch.nn.Module):
+    """VGG2L module for custom encoder.
+
+    Args:
+        idim: Dimension of inputs
+        odim: Dimension of outputs
+        pos_enc: Positional encoding class
+
+    """
+
+    def __init__(self, idim: int, odim: int, pos_enc: torch.nn.Module = None):
+        """Construct a VGG2L object."""
+        super().__init__()
+
+        self.vgg2l = torch.nn.Sequential(
+            torch.nn.Conv2d(1, 64, 3, stride=1, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(64, 64, 3, stride=1, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d((3, 2)),
+            torch.nn.Conv2d(64, 128, 3, stride=1, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(128, 128, 3, stride=1, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d((2, 2)),
+        )
+
+        if pos_enc is not None:
+            self.output = torch.nn.Sequential(
+                torch.nn.Linear(128 * ((idim // 2) // 2), odim), pos_enc
+            )
+        else:
+            self.output = torch.nn.Linear(128 * ((idim // 2) // 2), odim)
+
+    def forward(
+        self, x: torch.Tensor, x_mask: torch.Tensor
+    ) -> Union[
+        Tuple[torch.Tensor, torch.Tensor],
+        Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor],
+    ]:
+        """VGG2L forward for x.
+
+        Args:
+            x: Input tensor (B, T, idim)
+            x_mask: Input mask (B, 1, T)
+
+        Returns:
+            x: Output tensor (B, sub(T), odim)
+                   or ((B, sub(T), odim), (B, sub(T), att_dim))
+            x_mask: Output mask (B, 1, sub(T))
+
+        """
+        x = x.unsqueeze(1)
+        x = self.vgg2l(x)
+
+        b, c, t, f = x.size()
+
+        x = self.output(x.transpose(1, 2).contiguous().view(b, t, c * f))
+
+        if x_mask is not None:
+            x_mask = self.create_new_mask(x_mask)
+
+        return x, x_mask
+
+    def create_new_mask(self, x_mask: torch.Tensor) -> torch.Tensor:
+        """Create a subsampled version of x_mask.
+
+        Args:
+            x_mask: Input mask (B, 1, T)
+
+        Returns:
+            x_mask: Output mask (B, 1, sub(T))
+
+        """
+        x_t1 = x_mask.size(2) - (x_mask.size(2) % 3)
+        x_mask = x_mask[:, :, :x_t1][:, :, ::3]
+
+        x_t2 = x_mask.size(2) - (x_mask.size(2) % 2)
+        x_mask = x_mask[:, :, :x_t2][:, :, ::2]
+
+        return x_mask

lscodec/models/ctx_v2w.py

@@ -0,0 +1,108 @@
+# -*- coding: utf-8 -*-
+# MIT License
+# Copyright (c) 2025 Yiwei Guo
+
+import torch
+from typing import Union
+from lscodec.models.conformer.decoder import Decoder as ConformerDecoder
+
+from lscodec.utils import crop_seq
+
+from lscodec.models.prompt_prenet import ConvPromptPrenet
+from lscodec.models.hifigan import HiFiGANGenerator
+import logging
+
+
+class CTXVEC2WAVFrontend(torch.nn.Module):
+
+    def __init__(self,
+                 prompt_net_type,
+                 num_mels,
+                 vqvec_channels,
+                 prompt_channels,
+                 conformer_params):
+
+        super(CTXVEC2WAVFrontend, self).__init__()
+
+        if prompt_net_type == "ConvPromptPrenet":
+            self.prompt_prenet = ConvPromptPrenet(
+                embed=prompt_channels,
+                conv_layers=[(128, 3, 1, 1), (256, 5, 1, 2), (512, 5, 1, 2), (conformer_params["attention_dim"], 3, 1, 1)],
+                dropout=0.1,
+                skip_connections=True,
+                residual_scale=0.25,
+                non_affine_group_norm=False,
+                conv_bias=True,
+                activation=torch.nn.ReLU()
+            )
+        elif prompt_net_type == "Conv1d":
+            self.prompt_prenet = torch.nn.Conv1d(prompt_channels, conformer_params["attention_dim"], kernel_size=5, padding=2)  # , padding=2)
+        else:
+            raise NotImplementedError
+
+        self.input_dim = vqvec_channels
+        self.encoder1 = ConformerDecoder(vqvec_channels, input_layer='linear', **conformer_params)
+
+        self.hidden_proj = torch.nn.Linear(conformer_params["attention_dim"], conformer_params["attention_dim"])
+
+        self.encoder2 = ConformerDecoder(0, input_layer=None, **conformer_params)
+        self.mel_proj = torch.nn.Linear(conformer_params["attention_dim"], num_mels)
+
+    def forward(self, vqvec, prompt, mask=None, prompt_mask=None):
+        """
+        params:
+            vqvec: sequence of VQ-vectors.
+            prompt: sequence of mel-spectrogram prompt (acoustic context)
+            mask: mask of the vqvec. True or 1 stands for valid values.
+            prompt_mask: mask of the prompt.
+        vqvec and prompt are of shape [B, D, T]. All masks are of shape [B, T].
+        returns:
+            enc_out: the input to the CTX-vec2wav Generator (HifiGAN);
+            mel: the frontend predicted mel spectrogram (for faster convergence);
+        """
+        prompt = self.prompt_prenet(prompt.transpose(1, 2)).transpose(1, 2)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-2)
+        if prompt_mask is not None:
+            prompt_mask = prompt_mask.unsqueeze(-2)
+        enc_out, _ = self.encoder1(vqvec, mask, prompt, prompt_mask)
+
+        h = self.hidden_proj(enc_out)
+
+        enc_out, _ = self.encoder2(h, mask, prompt, prompt_mask)
+        mel = self.mel_proj(enc_out)  # (B, L, 80)
+
+        return enc_out, mel
+
+
+class CTXVEC2WAVGenerator(torch.nn.Module):
+
+    def __init__(self, frontend: CTXVEC2WAVFrontend, backend: HiFiGANGenerator):
+
+        super(CTXVEC2WAVGenerator, self).__init__()
+        self.frontend = frontend
+        self.backend = backend
+
+    def forward(self, vqvec, prompt, mask=None, prompt_mask=None, crop_len=0, crop_offsets=None):
+        """
+        :param vqvec: (torch.Tensor) The shape is (B, L, D). Sequence of VQ-vectors.
+        :param prompt: (torch.Tensor) The shape is (B, L', D'). Sequence of prompt features (acoustic context)
+        :param mask: (torch.Tensor) The dtype is torch.bool. The shape is (B, L). True or 1 stands for valid values in `vqvec`.
+        :param prompt_mask: (torch.Tensor) The dtype is torch.bool. The shape is (B, L'). True or 1 stands for valid values in `prompt`.
+        :return: frontend predicted mel spectrogram; reconstructed waveform.
+        """
+        h, mel, _ = self.frontend(vqvec, prompt, mask=mask, prompt_mask=prompt_mask)  # (B, L, adim), (B, L, 80)
+        if mask is not None:
+            h = h.masked_fill(~mask.unsqueeze(-1), 0)
+        h = h.transpose(1, 2)
+        if crop_len > 0:
+            h = crop_seq(h, crop_offsets, crop_len)
+        wav = self.backend(h)  # (B, C, T)
+        return mel, None, wav
+
+    def inference(self, vqvec, prompt, **kwargs):
+        h, mel = self.frontend(vqvec, prompt)
+        wav = self.backend.inference(h, **kwargs)
+
+        return mel, None, wav

lscodec/models/fairseq_modules/fp32_group_norm.py

@@ -0,0 +1,133 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Layer norm done in fp32 (for fp16 training)
+"""
+
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+from math import ceil, floor
+
+import torch
+
+
+def torch_quantile(
+    input: torch.Tensor,
+    q: float | torch.Tensor,
+    dim: int | None = None,
+    keepdim: bool = False,
+    *,
+    interpolation: str = "nearest",
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Better torch.quantile for one SCALAR quantile.
+
+    Using torch.kthvalue. Better than torch.quantile because:
+        - No 2**24 input size limit (pytorch/issues/67592),
+        - Much faster, at least on big input sizes.
+
+    Arguments:
+        input (torch.Tensor): See torch.quantile.
+        q (float): See torch.quantile. Supports only scalar input
+            currently.
+        dim (int | None): See torch.quantile.
+        keepdim (bool): See torch.quantile. Supports only False
+            currently.
+        interpolation: {"nearest", "lower", "higher"}
+            See torch.quantile.
+        out (torch.Tensor | None): See torch.quantile. Supports only
+            None currently.
+    """
+    # Sanitization: q
+    try:
+        q = float(q)
+        assert 0 <= q <= 1
+    except Exception:
+        raise ValueError(f"Only scalar input 0<=q<=1 is currently supported (got {q})!")
+
+    # Sanitization: dim
+    # Because one cannot pass  `dim=None` to `squeeze()` or `kthvalue()`
+    if dim_was_none := dim is None:
+        dim = 0
+        input = input.reshape((-1,) + (1,) * (input.ndim - 1))
+
+    # Sanitization: inteporlation
+    if interpolation == "nearest":
+        inter = round
+    elif interpolation == "lower":
+        inter = floor
+    elif interpolation == "higher":
+        inter = ceil
+    else:
+        raise ValueError(
+            "Supported interpolations currently are {'nearest', 'lower', 'higher'} "
+            f"(got '{interpolation}')!"
+        )
+
+    # Sanitization: out
+    if out is not None:
+        raise ValueError(f"Only None value is currently supported for out (got {out})!")
+
+    # Logic
+    k = inter(q * (input.shape[dim] - 1)) + 1
+    out = torch.kthvalue(input, k, dim, keepdim=True, out=out)[0]
+
+    # Rectification: keepdim
+    if keepdim:
+        return out
+    if dim_was_none:
+        return out.squeeze()
+    else:
+        return out.squeeze(dim)
+
+    return out
+
+class Fp32GroupNorm(nn.GroupNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, input):
+        # print("In FP32 GroupNorm", input[0].mean(), input[0].std(), input[0].min(), input[0].max())
+        output = F.group_norm(
+            input.float(),
+            self.num_groups,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        )
+        
+        # NOTE: dirty hack
+        # if input.max()<1000:
+        #     output = F.group_norm(
+        #             input.float(),
+        #             self.num_groups,
+        #             self.weight.float() if self.weight is not None else None,
+        #             self.bias.float() if self.bias is not None else None,
+        #             self.eps,
+        #         )
+        # else:
+        #     print("Will remove outliers in a dirty way in FP32GroupNorm")
+        #     q1 = torch_quantile(input.view(-1), 0.25)
+        #     q3 = torch_quantile(input.view(-1), 0.75)
+        #     IQR = q3-q1
+        #     upper_bound = q3 + 1.5 * IQR
+        #     lower_bound = q1 - 1.5 * IQR
+        #     # correct_mean = input[(lower_bound<input) & (input<upper_bound)].mean()
+        #     # correct_std = input[(lower_bound<input) & (input<upper_bound)].std()
+        #     # print(f"After removing outliers, the correct mean and std are {correct_mean}, {correct_std}")
+        #     # output = (input-correct_mean)/(correct_std+self.eps) * self.weight.float().unsqueeze(0).unsqueeze(-1) + self.bias.unsqueeze(0).unsqueeze(-1)
+            
+        #     input[(lower_bound<input) & (input<upper_bound)] = 0.0
+        #     output = F.group_norm(
+        #             input.float(),
+        #             self.num_groups,
+        #             self.weight.float() if self.weight is not None else None,
+        #             self.bias.float() if self.bias is not None else None,
+        #             self.eps,
+        #         )
+            
+        return output.type_as(input)

lscodec/models/fairseq_modules/gumbel_vector_quantizer.py

@@ -0,0 +1,212 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# 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 torch.nn as nn
+import torch.nn.functional as F
+
+
+class GumbelVectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_vars,
+        temp,
+        groups,
+        combine_groups,
+        vq_dim,
+        time_first,
+        activation=nn.GELU(),
+        weight_proj_depth=1,
+        weight_proj_factor=1,
+        hard=True,
+        std=0,
+    ):
+        """Vector quantization using gumbel softmax
+
+        Args:
+            dim: input dimension (channels)
+            num_vars: number of quantized vectors per group
+            temp: temperature for training. this should be a tuple of 3 elements: (start, stop, decay factor)
+            groups: number of groups for vector quantization
+            combine_groups: whether to use the vectors for all groups
+            vq_dim: dimensionality of the resulting quantized vector
+            time_first: if true, expect input in BxTxC format, otherwise in BxCxT
+            activation: what activation to use (should be a module). this is only used if weight_proj_depth is > 1
+            weight_proj_depth: number of layers (with activation in between) to project input before computing logits
+            weight_proj_factor: this is used only if weight_proj_depth is > 1. scales the inner dimensionality of
+                                projections by this factor
+        """
+        super().__init__()
+
+        self.groups = groups
+        self.combine_groups = combine_groups
+        self.input_dim = dim
+        self.num_vars = num_vars
+        self.time_first = time_first
+        self.hard = hard
+
+        assert (
+            vq_dim % groups == 0
+        ), f"dim {vq_dim} must be divisible by groups {groups} for concatenation"
+
+        var_dim = vq_dim // groups
+        num_groups = groups if not combine_groups else 1
+
+        self.vars = nn.Parameter(torch.FloatTensor(1, num_groups * num_vars, var_dim))
+        if std == 0:
+            nn.init.uniform_(self.vars)
+        else:
+            nn.init.normal_(self.vars, mean=0, std=std)
+
+        if weight_proj_depth > 1:
+
+            def block(input_dim, output_dim):
+                return nn.Sequential(nn.Linear(input_dim, output_dim), activation)
+
+            inner_dim = self.input_dim * weight_proj_factor
+            self.weight_proj = nn.Sequential(
+                *[
+                    block(self.input_dim if i == 0 else inner_dim, inner_dim)
+                    for i in range(weight_proj_depth - 1)
+                ],
+                nn.Linear(inner_dim, groups * num_vars),
+            )
+        else:
+            self.weight_proj = nn.Linear(self.input_dim, groups * num_vars)
+            nn.init.normal_(self.weight_proj.weight, mean=0, std=1)
+            nn.init.zeros_(self.weight_proj.bias)
+
+        if isinstance(temp, str):
+            import ast
+
+            temp = ast.literal_eval(temp)
+        assert len(temp) == 3, f"{temp}, {len(temp)}"
+
+        self.max_temp, self.min_temp, self.temp_decay = temp
+        self.curr_temp = self.max_temp
+        self.codebook_indices = None
+
+    def set_num_updates(self, num_updates):
+        self.curr_temp = max(
+            self.max_temp * self.temp_decay**num_updates, self.min_temp
+        )
+
+    def get_codebook_indices(self):
+        if self.codebook_indices is None:
+            from itertools import product
+
+            p = [range(self.num_vars)] * self.groups
+            inds = list(product(*p))
+            self.codebook_indices = torch.tensor(
+                inds, dtype=torch.long, device=self.vars.device
+            ).flatten()
+
+            if not self.combine_groups:
+                self.codebook_indices = self.codebook_indices.view(
+                    self.num_vars**self.groups, -1
+                )
+                for b in range(1, self.groups):
+                    self.codebook_indices[:, b] += self.num_vars * b
+                self.codebook_indices = self.codebook_indices.flatten()
+        return self.codebook_indices
+
+    def codebook(self):
+        indices = self.get_codebook_indices()
+        return (
+            self.vars.squeeze(0)
+            .index_select(0, indices)
+            .view(self.num_vars**self.groups, -1)
+        )
+
+    def sample_from_codebook(self, b, n):
+        indices = self.get_codebook_indices()
+        indices = indices.view(-1, self.groups)
+        cb_size = indices.size(0)
+        assert (
+            n < cb_size
+        ), f"sample size {n} is greater than size of codebook {cb_size}"
+        sample_idx = torch.randint(low=0, high=cb_size, size=(b * n,))
+        indices = indices[sample_idx]
+
+        z = self.vars.squeeze(0).index_select(0, indices.flatten()).view(b, n, -1)
+        return z
+
+    def to_codebook_index(self, indices):
+        res = indices.new_full(indices.shape[:-1], 0)
+        for i in range(self.groups):
+            exponent = self.groups - i - 1
+            res += indices[..., i] * (self.num_vars**exponent)
+        return res
+
+    def forward_idx(self, x):
+        res = self.forward(x, produce_targets=True)
+        return res["x"], res["targets"]
+
+    def forward(self, x, produce_targets=True):
+
+        result = {"num_vars": self.num_vars * self.groups}
+
+        # if not self.time_first:
+            # x = x.transpose(1, 2)
+
+        bsz_tsz, fsz = x.shape
+        # x = x.reshape(-1, fsz)
+        x = self.weight_proj(x)
+        x = x.view(bsz_tsz * self.groups, -1)
+
+        with torch.no_grad():
+            _, k = x.max(-1)
+            hard_x = (
+                x.new_zeros(*x.shape)
+                .scatter_(-1, k.view(-1, 1), 1.0)
+                .view(bsz_tsz, self.groups, -1)
+            )
+            hard_probs = torch.mean(hard_x.float(), dim=0)
+            result["code_perplexity"] = torch.exp(
+                -torch.sum(hard_probs * torch.log(hard_probs + 1e-7), dim=-1)
+            ).sum()
+
+        avg_probs = torch.softmax(
+            x.view(bsz_tsz, self.groups, -1).float(), dim=-1
+        ).mean(dim=0)
+        result["prob_perplexity"] = torch.exp(
+            -torch.sum(avg_probs * torch.log(avg_probs + 1e-7), dim=-1)
+        ).sum()
+
+        result["temp"] = self.curr_temp
+
+        if self.training:
+            x = F.gumbel_softmax(x.float(), tau=self.curr_temp, hard=self.hard).type_as(
+                x
+            )
+        else:
+            x = hard_x
+
+        x = x.view(bsz_tsz, -1)
+
+        vars = self.vars
+        if self.combine_groups:
+            vars = vars.repeat(1, self.groups, 1)
+
+        if produce_targets:
+            result["targets"] = (
+                x.view(bsz_tsz * self.groups, -1)
+                .argmax(dim=-1)
+                .detach()
+            )
+
+        x = x.unsqueeze(-1) * vars
+        x = x.view(bsz_tsz, self.groups, self.num_vars, -1)
+        x = x.sum(-2)
+        x = x.view(bsz_tsz, -1)
+
+        # if not self.time_first:
+            # x = x.transpose(1, 2)  # BTC -> BCT
+
+        result["x"] = x
+
+        # return result
+        return x, result['targets'], None

lscodec/models/fairseq_modules/kmeans_vector_quantizer.py

@@ -0,0 +1,206 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# 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 torch.nn as nn
+from lscodec.models.fairseq_modules.fp32_group_norm import Fp32GroupNorm
+
+
+class KmeansVectorQuantizer(nn.Module):
+    def __init__(
+        self, dim, num_vars, groups, combine_groups, vq_dim, time_first, gamma=0.25
+    ):
+        """Vector quantization using straight pass-through estimator (i.e. kmeans)
+
+        Args:
+            dim: input dimension (channels)
+            num_vars: number of quantized vectors per group
+            groups: number of groups for vector quantization
+            combine_groups: whether to use the vectors for all groups
+            vq_dim: dimensionality of the resulting quantized vector
+            time_first: if true, expect input in BxTxC format, otherwise in BxCxT
+            gamma: commitment loss coefficient
+        """
+        super().__init__()
+
+        self.groups = groups
+        self.combine_groups = combine_groups
+        self.input_dim = dim
+        self.num_vars = num_vars
+        self.vq_dim = vq_dim
+        self.time_first = time_first
+
+        assert (
+            vq_dim % groups == 0
+        ), f"dim {vq_dim} must be divisible by groups {groups} for concatenation"
+
+        self.var_dim = vq_dim // groups
+        num_groups = groups if not combine_groups else 1
+
+        self.embedding = nn.Parameter(
+            0.01 * torch.randn(num_vars, num_groups, self.var_dim)
+        )
+        self.projection = nn.Sequential(
+            nn.Conv1d(dim, dim, kernel_size=1, groups=groups, bias=False),
+            Fp32GroupNorm(groups, dim),
+        )
+        self.gamma = gamma
+        self.mse_mean = nn.MSELoss(reduction="mean")
+
+    def _pass_grad(self, x, y):
+        """Manually set gradient for backward pass.
+        for y = f(x), ensure that during the backward pass,
+        dL/dy = dL/dx regardless of f(x).
+        Returns:
+            y, with the gradient forced to be dL/dy = dL/dx.
+        """
+
+        return y.detach() + (x - x.detach())
+
+    @property
+    def expand_embedding(self):
+        if self.combine_groups:
+            return self.embedding.expand(self.num_vars, self.groups, self.var_dim)
+        return self.embedding
+
+    def forward_idx(self, x):
+        res = self.forward(x, produce_targets=True)
+        return res["x"], res["targets"]
+
+    def forward_idx_limited(self, x, valid_label2vqidx_mat):
+        # mask_mat = convert_valid_label2vqidx_to_mask_mat(valid_label2vqidx)
+        res = self.forward_group2(x, mask_mat=valid_label2vqidx_mat, produce_targets=True)
+        return res['x'], res['targets']
+
+    def forward(self, x, produce_targets=False):
+
+        result = {"num_vars": self.num_vars}
+
+        if self.time_first:
+            x = x.transpose(1, 2)
+
+        bsz, fsz, tsz = x.shape
+
+        ze = self.projection(x)
+        ze_ = ze.view(bsz, self.groups, self.var_dim, tsz).permute(0, 3, 1, 2)
+        d = (
+            (ze_.unsqueeze(0) - self.expand_embedding.unsqueeze(1).unsqueeze(1))
+            .view(self.num_vars, bsz, tsz, self.groups, -1)
+            .norm(dim=-1, p=2)
+        )
+        idx = d.argmin(dim=0)
+        zq = (
+            torch.stack(
+                [
+                    self.expand_embedding[idx[..., group], group]
+                    for group in range(self.groups)
+                ],
+                dim=-2,
+            )
+            .view(bsz, tsz, self.groups * self.var_dim)
+            .permute(0, 2, 1)
+        )
+        assert ze.shape == zq.shape, (ze.shape, zq.shape)
+        x = self._pass_grad(ze, zq)
+
+        with torch.no_grad():
+            hard_x = (
+                idx.new_zeros(bsz * tsz * self.groups, self.num_vars)
+                .scatter_(-1, idx.view(-1, 1), 1.0)
+                .view(bsz * tsz, self.groups, -1)
+            )
+            hard_probs = torch.mean(hard_x.float(), dim=0)
+            result["code_perplexity"] = torch.exp(
+                -torch.sum(hard_probs * torch.log(hard_probs + 1e-7), dim=-1)
+            ).sum()
+
+        if produce_targets:
+            result["targets"] = idx
+
+        if self.time_first:
+            x = x.transpose(1, 2)  # BCT -> BTC
+        result["x"] = x
+
+        ze = ze.float()
+        zq = zq.float()
+        latent_loss = self.mse_mean(zq, ze.detach())
+        commitment_loss = self.mse_mean(ze, zq.detach())
+
+        result["kmeans_loss"] = latent_loss + self.gamma * commitment_loss
+
+        return result
+
+    def forward_group2(self, x, mask_mat=None, produce_targets=False, inf=999999):
+        assert mask_mat is not None
+
+        result = {"num_vars": self.num_vars}
+
+        if self.time_first:
+            x = x.transpose(1, 2)
+
+        bsz, fsz, tsz = x.shape
+
+        ze = self.projection(x)
+
+        ze_ = ze.view(bsz, self.groups, self.var_dim, tsz).permute(0, 3, 1, 2)
+        ze_0 = ze_[:, :, 0, None, :]
+        ze_1 = ze_[:, :, 1, None, :]  # 4 * 100 * 320 * 128
+        cb0_expand = self.expand_embedding[:, 0, :]
+        cb1_expand = self.expand_embedding[:, 1, :]  # 320 * 128
+        dist_0 = ((ze_0 - cb0_expand) ** 2).sum(dim=-1)[:, :, :, None]
+        dist_1 = ((ze_1 - cb1_expand) ** 2).sum(dim=-1)[:, :, None, :]
+        res_0, res_1 = torch.broadcast_tensors(dist_0, dist_1)
+        mask_mat = (1 - mask_mat[None, None, :, :].to(res_0.device) * torch.ones_like(res_0)) * inf
+        # mask_mat = mask_mat.to(x.device)
+        d_flt = (res_0 + res_1 + mask_mat).view(bsz, tsz, -1)
+        idx_flt = torch.argmin(d_flt, dim=-1)
+        idx = torch.stack((idx_flt // self.num_vars, idx_flt % self.num_vars), dim=-1)
+
+        zq = (
+            torch.stack(
+                [
+                    self.expand_embedding[idx[..., group], group]
+                    for group in range(self.groups)
+                ],
+                dim=-2,
+            )
+            .view(bsz, tsz, self.groups * self.var_dim)
+            .permute(0, 2, 1)
+        )
+        assert ze.shape == zq.shape, (ze.shape, zq.shape)
+        x = self._pass_grad(ze, zq)
+
+        with torch.no_grad():
+            hard_x = (
+                idx.new_zeros(bsz * tsz * self.groups, self.num_vars)
+                .scatter_(-1, idx.view(-1, 1), 1.0)
+                .view(bsz * tsz, self.groups, -1)
+            )
+            hard_probs = torch.mean(hard_x.float(), dim=0)
+            result["code_perplexity"] = torch.exp(
+                -torch.sum(hard_probs * torch.log(hard_probs + 1e-7), dim=-1)
+            ).sum()
+
+        if produce_targets:
+            result["targets"] = idx
+
+        if self.time_first:
+            x = x.transpose(1, 2)  # BCT -> BTC
+        result["x"] = x
+
+        ze = ze.float()
+        zq = zq.float()
+        latent_loss = self.mse_mean(zq, ze.detach())
+        commitment_loss = self.mse_mean(ze, zq.detach())
+
+        result["kmeans_loss"] = latent_loss + self.gamma * commitment_loss
+
+        return result
+
+if __name__ == "__main__":
+    quantizer = KmeansVectorQuantizer(dim=256, num_vars=320, groups=2, combine_groups=False, vq_dim=256, time_first=True)
+    x = torch.ones(4, 100, 256)
+    result = quantizer.forward_group2(x, mask_mat=torch.randint(0, 2, (320, 320)))
+    print(result)

lscodec/models/fairseq_modules/layer_norm.py

@@ -0,0 +1,48 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# 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 torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from apex.normalization import FusedLayerNorm as _FusedLayerNorm
+
+    has_fused_layernorm = True
+
+    class FusedLayerNorm(_FusedLayerNorm):
+        @torch.jit.unused
+        def forward(self, x):
+            if not x.is_cuda:
+                return super().forward(x)
+            else:
+                with torch.cuda.device(x.device):
+                    return super().forward(x)
+
+except ImportError:
+    has_fused_layernorm = False
+
+
+def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        export = True
+    if not export and torch.cuda.is_available() and has_fused_layernorm:
+        return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
+    return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)
+
+
+class Fp32LayerNorm(nn.LayerNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, input):
+        output = F.layer_norm(
+            input.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        )
+        return output.type_as(input)

lscodec/models/fairseq_modules/transpose_last.py

@@ -0,0 +1,21 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+transpose last 2 dimensions of the input
+"""
+
+import torch.nn as nn
+
+
+class TransposeLast(nn.Module):
+    def __init__(self, deconstruct_idx=None, tranpose_dim=-2):
+        super().__init__()
+        self.deconstruct_idx = deconstruct_idx
+        self.tranpose_dim = tranpose_dim
+
+    def forward(self, x):
+        if self.deconstruct_idx is not None:
+            x = x[self.deconstruct_idx]
+        return x.transpose(self.tranpose_dim, -1)

lscodec/models/feature_extractor.py

@@ -0,0 +1,178 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Portions of this code are a modification of the fairseq project.
+# The original source code can be found at:
+# https://github.com/facebookresearch/fairseq/blob/d13e14a800bb588e5a77fb4e551f554ff9b24a72/fairseq/models/wav2vec/wav2vec.py
+#
+# Copyright (c) 2025 Yiwei Guo
+# Licensed under the MIT license.
+
+import math
+
+import torch.nn as nn
+from lscodec.models.fairseq_modules.fp32_group_norm import Fp32GroupNorm
+from lscodec.models.fairseq_modules.layer_norm import Fp32LayerNorm
+from lscodec.models.fairseq_modules.transpose_last import TransposeLast
+import torch
+
+
+def norm_block(is_layer_norm, dim, affine=True):
+    if is_layer_norm:
+        mod = nn.Sequential(
+            TransposeLast(),
+            Fp32LayerNorm(dim, elementwise_affine=affine),
+            TransposeLast(),
+        )
+    else:
+        mod = Fp32GroupNorm(1, dim, affine=affine)
+
+    return mod
+
+
+class ZeroPad1d(nn.Module):
+    def __init__(self, pad_left, pad_right):
+        super().__init__()
+        self.pad_left = pad_left
+        self.pad_right = pad_right
+
+    def forward(self, x):
+        return nn.functional.pad(x, (self.pad_left, self.pad_right))
+
+
+class ConvFeatureExtractionModel(nn.Module):
+    def __init__(
+            self,
+            conv_layers,
+            dropout,
+            log_compression,
+            skip_connections,
+            residual_scale,
+            non_affine_group_norm,
+            activation,
+    ):
+        super().__init__()
+
+        def block(n_in, n_out, k, stride):
+            return nn.Sequential(
+                nn.Conv1d(n_in, n_out, k, stride=stride, bias=False),
+                nn.Dropout(p=dropout),
+                norm_block(
+                    is_layer_norm=False, dim=n_out, affine=not non_affine_group_norm
+                ),
+                activation,
+            )
+
+        in_d = 1
+        self.conv_layers = nn.ModuleList()
+        for dim, k, stride in conv_layers:
+            self.conv_layers.append(block(in_d, dim, k, stride))
+            in_d = dim
+
+        self.log_compression = log_compression
+        self.skip_connections = skip_connections
+        self.residual_scale = math.sqrt(residual_scale)
+
+    def forward(self, x):  # x is audio of shape [B, 1, T]
+        # Bx1xT -> BxCxT
+        for conv in self.conv_layers:
+            residual = x
+            x = conv(x)
+            if self.skip_connections and x.size(1) == residual.size(1):
+                tsz = x.size(2)
+                r_tsz = residual.size(2)
+                residual = residual[..., :: r_tsz // tsz][..., :tsz]
+                x = (x + residual) * self.residual_scale
+
+        if self.log_compression:
+            x = x.abs()
+            x = x + 1
+            x = x.log()
+
+        return x
+
+
+class ConvAggregator(nn.Module):
+    def __init__(
+        self,
+        conv_layers,
+        embed,
+        dropout,
+        skip_connections,
+        residual_scale,
+        non_affine_group_norm,
+        conv_bias,
+        zero_pad,
+        activation,
+    ):
+        super().__init__()
+
+        def block(n_in, n_out, k, stride):
+            # padding dims only really make sense for stride = 1
+            ka = k // 2
+            kb = ka - 1 if k % 2 == 0 else ka
+
+            pad = (
+                ZeroPad1d(ka + kb, 0) if zero_pad else nn.ReplicationPad1d((ka + kb, 0))
+            )
+
+            return nn.Sequential(
+                pad,
+                nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias),
+                nn.Dropout(p=dropout),
+                norm_block(False, n_out, affine=not non_affine_group_norm),
+                activation,
+            )
+
+        in_d = embed
+        self.conv_layers = nn.ModuleList()
+        self.residual_proj = nn.ModuleList()
+        for dim, k, stride in conv_layers:
+            if in_d != dim and skip_connections:
+                self.residual_proj.append(nn.Conv1d(in_d, dim, 1, bias=False))
+            else:
+                self.residual_proj.append(None)
+
+            self.conv_layers.append(block(in_d, dim, k, stride))
+            in_d = dim
+        self.conv_layers = nn.Sequential(*self.conv_layers)
+        self.skip_connections = skip_connections
+        self.residual_scale = math.sqrt(residual_scale)
+
+    def forward(self, x):
+        for rproj, conv in zip(self.residual_proj, self.conv_layers):
+            residual = x
+            x = conv(x)
+            if self.skip_connections:
+                if rproj is not None:
+                    residual = rproj(residual)
+                x = (x + residual) * self.residual_scale
+        return x
+
+
+class LSCodecEncoder(nn.Module):
+    def __init__(self, conv_extraction_model: ConvFeatureExtractionModel,
+                 conv_aggregation_model: ConvAggregator,
+                 dropout_p: float = 0.1):
+        super().__init__()
+        self.conv_models = nn.Sequential(conv_extraction_model, conv_aggregation_model)
+
+        self.dropout = nn.Dropout(dropout_p)
+
+    def forward(self, waveform, frame_mask, length_tolerance=3):
+        # waveform: [B, 1, T]. Frame mask is 1 at valid values.
+        conv_features = self.conv_models(waveform)  # [B, 1024, L]
+        conv_features = self.dropout(conv_features)
+        B, D, L = conv_features.shape
+        if frame_mask is None:
+            frame_mask = torch.ones(size=(B, L)).bool().to(conv_features.device)
+        if (frame_mask is not None) and (0 < frame_mask.shape[-1] - L <= length_tolerance):
+            pad_tensor = torch.zeros(size=(B, D, frame_mask.shape[-1] - L)).to(conv_features.device)
+            # NOTE: we don't pad these values with 999, but instead treat those as valid data points
+            # NOTE: this is because the mask is longer than the valid data points, and changing the mask is inconvenient.
+            conv_features = torch.concat([conv_features, pad_tensor], dim=-1)
+
+        return conv_features

lscodec/models/hifigan.py

@@ -0,0 +1,732 @@
+# -*- coding: utf-8 -*-
+
+"""HiFi-GAN Modules.
+
+This code is based on https://github.com/jik876/hifi-gan.
+
+"""
+
+import copy
+import logging
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from lscodec.layers import HiFiGANResidualBlock as ResidualBlock
+from lscodec.utils import read_hdf5
+
+
+class HiFiGANGenerator(torch.nn.Module):
+    """HiFiGAN generator module."""
+
+    def __init__(
+            self,
+            in_channels=80,
+            out_channels=1,
+            channels=512,
+            kernel_size=7,
+            upsample_scales=(8, 8, 2, 2),
+            upsample_kernel_sizes=(16, 16, 4, 4),
+            resblock_kernel_sizes=(3, 7, 11),
+            resblock_dilations=[(1, 3, 5), (1, 3, 5), (1, 3, 5)],
+            use_additional_convs=True,
+            bias=True,
+            nonlinear_activation="LeakyReLU",
+            nonlinear_activation_params={"negative_slope": 0.1},
+            use_weight_norm=True,
+    ):
+        """Initialize HiFiGANGenerator module.
+
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels.
+            channels (int): Number of hidden representation channels.
+            kernel_size (int): Kernel size of initial and final conv layer.
+            upsample_scales (list): List of upsampling scales.
+            upsample_kernel_sizes (list): List of kernel sizes for upsampling layers.
+            resblock_kernel_sizes (list): List of kernel sizes for residual blocks.
+            resblock_dilations (list): List of dilation list for residual blocks.
+            use_additional_convs (bool): Whether to use additional conv layers in residual blocks.
+            bias (bool): Whether to add bias parameter in convolution layers.
+            nonlinear_activation (str): Activation function module name.
+            nonlinear_activation_params (dict): Hyperparameters for activation function.
+            use_weight_norm (bool): Whether to use weight norm.
+                If set to true, it will be applied to all the conv layers.
+
+        """
+        super().__init__()
+
+        # check hyperparameters are valid
+        assert kernel_size % 2 == 1, "Kernel size must be odd number."
+        assert len(upsample_scales) == len(upsample_kernel_sizes)
+        assert len(resblock_dilations) == len(resblock_kernel_sizes)
+
+        # define modules
+        self.num_upsamples = len(upsample_kernel_sizes)
+        self.num_blocks = len(resblock_kernel_sizes)
+        self.input_conv = torch.nn.Conv1d(
+            in_channels,
+            channels,
+            kernel_size,
+            1,
+            padding=(kernel_size - 1) // 2,
+        )
+        self.upsamples = torch.nn.ModuleList()
+        self.blocks = torch.nn.ModuleList()
+        for i in range(len(upsample_kernel_sizes)):
+            assert upsample_kernel_sizes[i] == 2 * upsample_scales[i]
+            self.upsamples += [
+                torch.nn.Sequential(
+                    getattr(torch.nn, nonlinear_activation)(
+                        **nonlinear_activation_params
+                    ),
+                    torch.nn.ConvTranspose1d(
+                        channels // (2 ** i),
+                        channels // (2 ** (i + 1)),
+                        upsample_kernel_sizes[i],
+                        upsample_scales[i],
+                        padding=upsample_scales[i] // 2 + upsample_scales[i] % 2,
+                        output_padding=upsample_scales[i] % 2,
+                    ),
+                )
+            ]
+            for j in range(len(resblock_kernel_sizes)):
+                self.blocks += [
+                    ResidualBlock(
+                        kernel_size=resblock_kernel_sizes[j],
+                        channels=channels // (2 ** (i + 1)),
+                        dilations=resblock_dilations[j],
+                        bias=bias,
+                        use_additional_convs=use_additional_convs,
+                        nonlinear_activation=nonlinear_activation,
+                        nonlinear_activation_params=nonlinear_activation_params,
+                    )
+                ]
+        self.output_conv = torch.nn.Sequential(
+            # NOTE(kan-bayashi): follow official implementation but why
+            #   using different slope parameter here? (0.1 vs. 0.01)
+            torch.nn.LeakyReLU(),
+            torch.nn.Conv1d(
+                channels // (2 ** (i + 1)),
+                out_channels,
+                kernel_size,
+                1,
+                padding=(kernel_size - 1) // 2,
+            ),
+            torch.nn.Tanh(),
+        )
+
+        # apply weight norm
+        if use_weight_norm:
+            self.apply_weight_norm()
+
+        # reset parameters
+        self.reset_parameters()
+
+    def forward(self, c):
+        """Calculate forward propagation.
+
+        Args:
+            c (Tensor): Input tensor (B, in_channels, T).
+
+        Returns:
+            Tensor: Output tensor (B, out_channels, T).
+
+        """
+        c = self.input_conv(c)
+        for i in range(self.num_upsamples):
+            c = self.upsamples[i](c)
+            cs = 0.0  # initialize
+            for j in range(self.num_blocks):
+                cs += self.blocks[i * self.num_blocks + j](c)
+            c = cs / self.num_blocks
+        c = self.output_conv(c)
+
+        return c
+
+    def reset_parameters(self):
+        """Reset parameters.
+
+        This initialization follows the official implementation manner.
+        https://github.com/jik876/hifi-gan/blob/master/models.py
+
+        """
+
+        def _reset_parameters(m):
+            if isinstance(m, (torch.nn.Conv1d, torch.nn.ConvTranspose1d)):
+                m.weight.data.normal_(0.0, 0.01)
+                logging.debug(f"Reset parameters in {m}.")
+
+        self.apply(_reset_parameters)
+
+    def remove_weight_norm(self):
+        """Remove weight normalization module from all the layers."""
+
+        def _remove_weight_norm(m):
+            try:
+                logging.debug(f"Weight norm is removed from {m}.")
+                torch.nn.utils.remove_weight_norm(m)
+            except ValueError:  # this module didn't have weight norm
+                return
+
+        self.apply(_remove_weight_norm)
+
+    def apply_weight_norm(self):
+        """Apply weight normalization module from all the layers."""
+
+        def _apply_weight_norm(m):
+            if isinstance(m, torch.nn.Conv1d) or isinstance(
+                    m, torch.nn.ConvTranspose1d
+            ):
+                torch.nn.utils.weight_norm(m)
+                logging.debug(f"Weight norm is applied to {m}.")
+
+        self.apply(_apply_weight_norm)
+
+    def register_stats(self, stats):
+        """Register stats for de-normalization as buffer.
+
+        Args:
+            stats (str): Path of statistics file (".npy" or ".h5").
+
+        """
+        assert stats.endswith(".h5") or stats.endswith(".npy")
+        if stats.endswith(".h5"):
+            mean = read_hdf5(stats, "mean").reshape(-1)
+            scale = read_hdf5(stats, "scale").reshape(-1)
+        else:
+            mean = np.load(stats)[0].reshape(-1)
+            scale = np.load(stats)[1].reshape(-1)
+        self.register_buffer("mean", torch.from_numpy(mean).float())
+        self.register_buffer("scale", torch.from_numpy(scale).float())
+        logging.info("Successfully registered stats as buffer.")
+
+    def inference(self, c, normalize_before=False):
+        """Perform inference.
+
+        Args:
+            c (Union[Tensor, ndarray]): Input tensor (T, in_channels).
+            normalize_before (bool): Whether to perform normalization.
+
+        Returns:
+            Tensor: Output tensor (T ** prod(upsample_scales), out_channels).
+
+        """
+        # if not isinstance(c, torch.Tensor):
+        #     c = torch.tensor(c, dtype=torch.float).to(next(self.parameters()).device)
+        if normalize_before:
+            c = (c - self.mean) / self.scale
+        c = self.forward(c.transpose(1, 2))
+        return c.squeeze(0).transpose(1, 0)
+
+
+class HiFiGANPeriodDiscriminator(torch.nn.Module):
+    """HiFiGAN period discriminator module."""
+
+    def __init__(
+            self,
+            in_channels=1,
+            out_channels=1,
+            period=3,
+            kernel_sizes=[5, 3],
+            channels=32,
+            downsample_scales=[3, 3, 3, 3, 1],
+            max_downsample_channels=1024,
+            bias=True,
+            nonlinear_activation="LeakyReLU",
+            nonlinear_activation_params={"negative_slope": 0.1},
+            use_weight_norm=True,
+            use_spectral_norm=False,
+    ):
+        """Initialize HiFiGANPeriodDiscriminator module.
+
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels.
+            period (int): Period.
+            kernel_sizes (list): Kernel sizes of initial conv layers and the final conv layer.
+            channels (int): Number of initial channels.
+            downsample_scales (list): List of downsampling scales.
+            max_downsample_channels (int): Number of maximum downsampling channels.
+            use_additional_convs (bool): Whether to use additional conv layers in residual blocks.
+            bias (bool): Whether to add bias parameter in convolution layers.
+            nonlinear_activation (str): Activation function module name.
+            nonlinear_activation_params (dict): Hyperparameters for activation function.
+            use_weight_norm (bool): Whether to use weight norm.
+                If set to true, it will be applied to all the conv layers.
+            use_spectral_norm (bool): Whether to use spectral norm.
+                If set to true, it will be applied to all the conv layers.
+
+        """
+        super().__init__()
+        assert len(kernel_sizes) == 2
+        assert kernel_sizes[0] % 2 == 1, "Kernel size must be odd number."
+        assert kernel_sizes[1] % 2 == 1, "Kernel size must be odd number."
+
+        self.period = period
+        self.convs = torch.nn.ModuleList()
+        in_chs = in_channels
+        out_chs = channels
+        for downsample_scale in downsample_scales:
+            self.convs += [
+                torch.nn.Sequential(
+                    torch.nn.Conv2d(
+                        in_chs,
+                        out_chs,
+                        (kernel_sizes[0], 1),
+                        (downsample_scale, 1),
+                        padding=((kernel_sizes[0] - 1) // 2, 0),
+                    ),
+                    getattr(torch.nn, nonlinear_activation)(
+                        **nonlinear_activation_params
+                    ),
+                )
+            ]
+            in_chs = out_chs
+            # NOTE(kan-bayashi): Use downsample_scale + 1?
+            out_chs = min(out_chs * 4, max_downsample_channels)
+        self.output_conv = torch.nn.Conv2d(
+            out_chs,
+            out_channels,
+            (kernel_sizes[1] - 1, 1),
+            1,
+            padding=((kernel_sizes[1] - 1) // 2, 0),
+        )
+
+        if use_weight_norm and use_spectral_norm:
+            raise ValueError("Either use use_weight_norm or use_spectral_norm.")
+
+        # apply weight norm
+        if use_weight_norm:
+            self.apply_weight_norm()
+
+        # apply spectral norm
+        if use_spectral_norm:
+            self.apply_spectral_norm()
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            c (Tensor): Input tensor (B, in_channels, T).
+
+        Returns:
+            list: List of each layer's tensors.
+
+        """
+        # transform 1d to 2d -> (B, C, T/P, P)
+        b, c, t = x.shape
+        if t % self.period != 0:
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t += n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        # forward conv
+        outs = []
+        for layer in self.convs:
+            x = layer(x)
+            outs += [x]
+        x = self.output_conv(x)
+        x = torch.flatten(x, 1, -1)
+        outs += [x]
+
+        return outs
+
+    def apply_weight_norm(self):
+        """Apply weight normalization module from all the layers."""
+
+        def _apply_weight_norm(m):
+            if isinstance(m, torch.nn.Conv2d):
+                torch.nn.utils.weight_norm(m)
+                logging.debug(f"Weight norm is applied to {m}.")
+
+        self.apply(_apply_weight_norm)
+
+    def apply_spectral_norm(self):
+        """Apply spectral normalization module from all the layers."""
+
+        def _apply_spectral_norm(m):
+            if isinstance(m, torch.nn.Conv2d):
+                torch.nn.utils.spectral_norm(m)
+                logging.debug(f"Spectral norm is applied to {m}.")
+
+        self.apply(_apply_spectral_norm)
+
+
+class HiFiGANMultiPeriodDiscriminator(torch.nn.Module):
+    """HiFiGAN multi-period discriminator module."""
+
+    def __init__(
+            self,
+            periods=[2, 3, 5, 7, 11],
+            discriminator_params={
+                "in_channels": 1,
+                "out_channels": 1,
+                "kernel_sizes": [5, 3],
+                "channels": 32,
+                "downsample_scales": [3, 3, 3, 3, 1],
+                "max_downsample_channels": 1024,
+                "bias": True,
+                "nonlinear_activation": "LeakyReLU",
+                "nonlinear_activation_params": {"negative_slope": 0.1},
+                "use_weight_norm": True,
+                "use_spectral_norm": False,
+            },
+    ):
+        """Initialize HiFiGANMultiPeriodDiscriminator module.
+
+        Args:
+            periods (list): List of periods.
+            discriminator_params (dict): Parameters for hifi-gan period discriminator module.
+                The period parameter will be overwritten.
+
+        """
+        super().__init__()
+        self.discriminators = torch.nn.ModuleList()
+        for period in periods:
+            params = copy.deepcopy(discriminator_params)
+            params["period"] = period
+            self.discriminators += [HiFiGANPeriodDiscriminator(**params)]
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input noise signal (B, 1, T).
+
+        Returns:
+            List: List of list of each discriminator outputs, which consists of each layer output tensors.
+
+        """
+        outs = []
+        for f in self.discriminators:
+            outs += [f(x)]
+
+        return outs
+
+
+class HiFiGANScaleDiscriminator(torch.nn.Module):
+    """HiFi-GAN scale discriminator module."""
+
+    def __init__(
+            self,
+            in_channels=1,
+            out_channels=1,
+            kernel_sizes=[15, 41, 5, 3],
+            channels=128,
+            max_downsample_channels=1024,
+            max_groups=16,
+            bias=True,
+            downsample_scales=[2, 2, 4, 4, 1],
+            nonlinear_activation="LeakyReLU",
+            nonlinear_activation_params={"negative_slope": 0.1},
+            use_weight_norm=True,
+            use_spectral_norm=False,
+    ):
+        """Initilize HiFiGAN scale discriminator module.
+
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels.
+            kernel_sizes (list): List of four kernel sizes. The first will be used for the first conv layer,
+                and the second is for downsampling part, and the remaining two are for output layers.
+            channels (int): Initial number of channels for conv layer.
+            max_downsample_channels (int): Maximum number of channels for downsampling layers.
+            bias (bool): Whether to add bias parameter in convolution layers.
+            downsample_scales (list): List of downsampling scales.
+            nonlinear_activation (str): Activation function module name.
+            nonlinear_activation_params (dict): Hyperparameters for activation function.
+            use_weight_norm (bool): Whether to use weight norm.
+                If set to true, it will be applied to all the conv layers.
+            use_spectral_norm (bool): Whether to use spectral norm.
+                If set to true, it will be applied to all the conv layers.
+
+        """
+        super().__init__()
+        self.layers = torch.nn.ModuleList()
+
+        # check kernel size is valid
+        assert len(kernel_sizes) == 4
+        for ks in kernel_sizes:
+            assert ks % 2 == 1
+
+        # add first layer
+        self.layers += [
+            torch.nn.Sequential(
+                torch.nn.Conv1d(
+                    in_channels,
+                    channels,
+                    # NOTE(kan-bayashi): Use always the same kernel size
+                    kernel_sizes[0],
+                    bias=bias,
+                    padding=(kernel_sizes[0] - 1) // 2,
+                ),
+                getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params),
+            )
+        ]
+
+        # add downsample layers
+        in_chs = channels
+        out_chs = channels
+        # NOTE(kan-bayashi): Remove hard coding?
+        groups = 4
+        for downsample_scale in downsample_scales:
+            self.layers += [
+                torch.nn.Sequential(
+                    torch.nn.Conv1d(
+                        in_chs,
+                        out_chs,
+                        kernel_size=kernel_sizes[1],
+                        stride=downsample_scale,
+                        padding=(kernel_sizes[1] - 1) // 2,
+                        groups=groups,
+                        bias=bias,
+                    ),
+                    getattr(torch.nn, nonlinear_activation)(
+                        **nonlinear_activation_params
+                    ),
+                )
+            ]
+            in_chs = out_chs
+            # NOTE(kan-bayashi): Remove hard coding?
+            out_chs = min(in_chs * 2, max_downsample_channels)
+            # NOTE(kan-bayashi): Remove hard coding?
+            groups = min(groups * 4, max_groups)
+
+        # add final layers
+        out_chs = min(in_chs * 2, max_downsample_channels)
+        self.layers += [
+            torch.nn.Sequential(
+                torch.nn.Conv1d(
+                    in_chs,
+                    out_chs,
+                    kernel_size=kernel_sizes[2],
+                    stride=1,
+                    padding=(kernel_sizes[2] - 1) // 2,
+                    bias=bias,
+                ),
+                getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params),
+            )
+        ]
+        self.layers += [
+            torch.nn.Conv1d(
+                out_chs,
+                out_channels,
+                kernel_size=kernel_sizes[3],
+                stride=1,
+                padding=(kernel_sizes[3] - 1) // 2,
+                bias=bias,
+            ),
+        ]
+
+        if use_weight_norm and use_spectral_norm:
+            raise ValueError("Either use use_weight_norm or use_spectral_norm.")
+
+        # apply weight norm
+        if use_weight_norm:
+            self.apply_weight_norm()
+
+        # apply spectral norm
+        if use_spectral_norm:
+            self.apply_spectral_norm()
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input noise signal (B, 1, T).
+
+        Returns:
+            List: List of output tensors of each layer.
+
+        """
+        outs = []
+        for f in self.layers:
+            x = f(x)
+            outs += [x]
+
+        return outs
+
+    def apply_weight_norm(self):
+        """Apply weight normalization module from all the layers."""
+
+        def _apply_weight_norm(m):
+            if isinstance(m, torch.nn.Conv2d):
+                torch.nn.utils.weight_norm(m)
+                logging.debug(f"Weight norm is applied to {m}.")
+
+        self.apply(_apply_weight_norm)
+
+    def apply_spectral_norm(self):
+        """Apply spectral normalization module from all the layers."""
+
+        def _apply_spectral_norm(m):
+            if isinstance(m, torch.nn.Conv2d):
+                torch.nn.utils.spectral_norm(m)
+                logging.debug(f"Spectral norm is applied to {m}.")
+
+        self.apply(_apply_spectral_norm)
+
+
+class HiFiGANMultiScaleDiscriminator(torch.nn.Module):
+    """HiFi-GAN multi-scale discriminator module."""
+
+    def __init__(
+            self,
+            scales=3,
+            downsample_pooling="AvgPool1d",
+            # follow the official implementation setting
+            downsample_pooling_params={
+                "kernel_size": 4,
+                "stride": 2,
+                "padding": 2,
+            },
+            discriminator_params={
+                "in_channels": 1,
+                "out_channels": 1,
+                "kernel_sizes": [15, 41, 5, 3],
+                "channels": 128,
+                "max_downsample_channels": 1024,
+                "max_groups": 16,
+                "bias": True,
+                "downsample_scales": [2, 2, 4, 4, 1],
+                "nonlinear_activation": "LeakyReLU",
+                "nonlinear_activation_params": {"negative_slope": 0.1},
+            },
+            follow_official_norm=False,
+    ):
+        """Initilize HiFiGAN multi-scale discriminator module.
+
+        Args:
+            scales (int): Number of multi-scales.
+            downsample_pooling (str): Pooling module name for downsampling of the inputs.
+            downsample_pooling_params (dict): Parameters for the above pooling module.
+            discriminator_params (dict): Parameters for hifi-gan scale discriminator module.
+            follow_official_norm (bool): Whether to follow the norm setting of the official
+                implementaion. The first discriminator uses spectral norm and the other
+                discriminators use weight norm.
+
+        """
+        super().__init__()
+        self.discriminators = torch.nn.ModuleList()
+
+        # add discriminators
+        for i in range(scales):
+            params = copy.deepcopy(discriminator_params)
+            if follow_official_norm:
+                if i == 0:
+                    params["use_weight_norm"] = False
+                    params["use_spectral_norm"] = True
+                else:
+                    params["use_weight_norm"] = True
+                    params["use_spectral_norm"] = False
+            self.discriminators += [HiFiGANScaleDiscriminator(**params)]
+        self.pooling = getattr(torch.nn, downsample_pooling)(
+            **downsample_pooling_params
+        )
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input noise signal (B, 1, T).
+
+        Returns:
+            List: List of list of each discriminator outputs, which consists of each layer output tensors.
+
+        """
+        outs = []
+        for f in self.discriminators:
+            outs += [f(x)]
+            x = self.pooling(x)
+
+        return outs
+
+
+class HiFiGANMultiScaleMultiPeriodDiscriminator(torch.nn.Module):
+    """HiFi-GAN multi-scale + multi-period discriminator module."""
+
+    def __init__(
+            self,
+            # Multi-scale discriminator related
+            scales=3,
+            scale_downsample_pooling="AvgPool1d",
+            scale_downsample_pooling_params={
+                "kernel_size": 4,
+                "stride": 2,
+                "padding": 2,
+            },
+            scale_discriminator_params={
+                "in_channels": 1,
+                "out_channels": 1,
+                "kernel_sizes": [15, 41, 5, 3],
+                "channels": 128,
+                "max_downsample_channels": 1024,
+                "max_groups": 16,
+                "bias": True,
+                "downsample_scales": [2, 2, 4, 4, 1],
+                "nonlinear_activation": "LeakyReLU",
+                "nonlinear_activation_params": {"negative_slope": 0.1},
+            },
+            follow_official_norm=True,
+            # Multi-period discriminator related
+            periods=[2, 3, 5, 7, 11],
+            period_discriminator_params={
+                "in_channels": 1,
+                "out_channels": 1,
+                "kernel_sizes": [5, 3],
+                "channels": 32,
+                "downsample_scales": [3, 3, 3, 3, 1],
+                "max_downsample_channels": 1024,
+                "bias": True,
+                "nonlinear_activation": "LeakyReLU",
+                "nonlinear_activation_params": {"negative_slope": 0.1},
+                "use_weight_norm": True,
+                "use_spectral_norm": False,
+            },
+    ):
+        """Initilize HiFiGAN multiscale + multi-period discriminator module.
+
+        Args:
+            scales (int): Number of multi-scales.
+            scale_downsample_pooling (str): Pooling module name for downsampling of the inputs.
+            scale_downsample_pooling_params (dict): Parameters for the above pooling module.
+            scale_discriminator_params (dict): Parameters for hifi-gan scale discriminator module.
+            follow_official_norm (bool): Whether to follow the norm setting of the official
+                implementaion. The first discriminator uses spectral norm and the other
+                discriminators use weight norm.
+            periods (list): List of periods.
+            period_discriminator_params (dict): Parameters for hifi-gan period discriminator module.
+                The period parameter will be overwritten.
+
+        """
+        super().__init__()
+        self.msd = HiFiGANMultiScaleDiscriminator(
+            scales=scales,
+            downsample_pooling=scale_downsample_pooling,
+            downsample_pooling_params=scale_downsample_pooling_params,
+            discriminator_params=scale_discriminator_params,
+            follow_official_norm=follow_official_norm,
+        )
+        self.mpd = HiFiGANMultiPeriodDiscriminator(
+            periods=periods,
+            discriminator_params=period_discriminator_params,
+        )
+
+    def forward(self, x):
+        """Calculate forward propagation.
+
+        Args:
+            x (Tensor): Input noise signal (B, 1, T).
+
+        Returns:
+            List: List of list of each discriminator outputs,
+                which consists of each layer output tensors.
+                Multiscale and multi period ones are concatenated.
+
+        """
+        msd_outs = self.msd(x)
+        mpd_outs = self.mpd(x)
+        return msd_outs + mpd_outs