Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+__pycache__/modeling.cpython-314.pyc filter=lfs diff=lfs merge=lfs -text
+onnx/model.onnx.data filter=lfs diff=lfs merge=lfs -text

__pycache__/modeling.cpython-314.pyc

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

config.json

@@ -0,0 +1,69 @@
+{
+  "format": "vtx-tts",
+  "train": {
+    "segment_size": 16384
+  },
+  "data": {
+    "text_cleaners": [],
+    "max_wav_value": 32768.0,
+    "sampling_rate": 24000,
+    "filter_length": 1024,
+    "hop_length": 256,
+    "win_length": 1024,
+    "n_mel_channels": 80,
+    "mel_fmin": 0.0,
+    "mel_fmax": 12000.0,
+    "add_blank": true,
+    "n_speakers": 0,
+    "cleaned_text": true
+  },
+  "model": {
+    "inter_channels": 192,
+    "hidden_channels": 96,
+    "filter_channels": 768,
+    "n_heads": 2,
+    "n_layers": 3,
+    "kernel_size": 3,
+    "p_dropout": 0.1,
+    "resblock": "1",
+    "resblock_kernel_sizes": [
+      3,
+      7,
+      11
+    ],
+    "resblock_dilation_sizes": [
+      [
+        1,
+        3,
+        5
+      ],
+      [
+        1,
+        3,
+        5
+      ],
+      [
+        1,
+        3,
+        5
+      ]
+    ],
+    "upsample_rates": [
+      8,
+      8,
+      2,
+      2
+    ],
+    "upsample_initial_channel": 320,
+    "upsample_kernel_sizes": [
+      16,
+      16,
+      4,
+      4
+    ],
+    "n_layers_q": 3,
+    "use_spectral_norm": false,
+    "use_sdp": false,
+    "inference_only": false
+  }
+}

configuration.py

@@ -0,0 +1,56 @@
+"""VTX-TTS Configuration & Text Processing System.
+
+Defines VTXConfig and text symbol tokenization.
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict
+
+
+# Symbol definitions for VTX-TTS phoneme tokenization
+_pad = "_"
+_punctuation = ";:,.!?¡¿—…\"«»“” "
+_letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+_letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+
+symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+_symbol_to_id = {s: i for i, s in enumerate(symbols)}
+
+
+def cleaned_text_to_sequence(cleaned_text: str) -> list[int]:
+    """Converts a string of text to a sequence of IDs corresponding to the symbols in the text."""
+    sequence = []
+    for symbol in cleaned_text:
+        if symbol in _symbol_to_id:
+            sequence.append(_symbol_to_id[symbol])
+    return sequence
+
+
+class HParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, dict):
+                v = HParams(**v)
+            self[k] = v
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return hasattr(self, key)
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+
+def get_hparams_from_file(config_path: str | Path) -> HParams:
+    with open(config_path, "r", encoding="utf-8") as f:
+        data = f.read()
+    config = json.loads(data)
+    return HParams(**config)

inference.py

@@ -0,0 +1,270 @@
+from __future__ import annotations
+
+import argparse
+import contextlib
+import io
+import logging
+import re
+import sys
+import warnings
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors.numpy import load_file
+
+from configuration import cleaned_text_to_sequence, get_hparams_from_file, symbols
+from modeling import SynthesizerTrn
+
+PACKAGE_ROOT = Path(__file__).resolve().parent
+
+
+def split_text(text: str, limit: int = 280) -> list[str]:
+    normalized = " ".join(text.split())
+    sentences = [
+        part.strip()
+        for part in re.split(r"(?<=[.!?;:])\s+", normalized)
+        if part.strip()
+    ]
+    chunks: list[str] = []
+    for sentence in sentences or [normalized]:
+        while len(sentence) > limit:
+            search = sentence[: limit + 1]
+            punctuation = max(search.rfind(mark) for mark in (",", ";", ":"))
+            split_at = (
+                punctuation + 1
+                if punctuation >= limit // 2
+                else sentence.rfind(" ", 0, limit + 1)
+            )
+            if split_at < limit // 2:
+                split_at = limit
+            chunks.append(sentence[:split_at].strip())
+            sentence = sentence[split_at:].strip()
+        if sentence:
+            chunks.append(sentence)
+    return chunks
+
+
+def boundary_pause_seconds(chunk: str) -> float:
+    ending = chunk.rstrip()[-1:] if chunk.strip() else ""
+    return {
+        "?": 0.28,
+        "!": 0.24,
+        ".": 0.22,
+        ";": 0.16,
+        ":": 0.13,
+        ",": 0.09,
+    }.get(ending, 0.08)
+
+
+def edge_fade(waveform: np.ndarray, sample_rate: int, milliseconds: float = 5.0) -> np.ndarray:
+    frames = min(round(sample_rate * milliseconds / 1000.0), waveform.size // 2)
+    if frames <= 0:
+        return waveform
+    output = waveform.copy()
+    ramp = np.linspace(0.0, 1.0, frames, endpoint=True, dtype=np.float32)
+    output[:frames] *= ramp
+    output[-frames:] *= ramp[::-1]
+    return output
+
+
+def optimize_for_inference(model: SynthesizerTrn) -> None:
+    """Collapse weight normalization for VTXVocoder & VTXVectorEstimator."""
+    with contextlib.redirect_stdout(io.StringIO()):
+        model.VTXVocoder.remove_weight_norm()
+    for flow in model.VTXVectorEstimator.flows:
+        encoder = getattr(flow, "enc", None)
+        if encoder is not None and hasattr(encoder, "remove_weight_norm"):
+            encoder.remove_weight_norm()
+
+
+class VtxTTS:
+    """VTX-TTS Engine using native 4-Bit model.safetensors."""
+
+    def __init__(self, model_dir: str | Path = PACKAGE_ROOT, device: str = "cpu") -> None:
+        self.root = Path(model_dir).resolve()
+        self.device = torch.device(device)
+        self.hps = get_hparams_from_file(str(self.root / "config.json"))
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                message="`torch.nn.utils.weight_norm` is deprecated",
+                category=FutureWarning,
+            )
+            model_kwargs = self.hps.model.__dict__ if hasattr(self.hps.model, "__dict__") else self.hps.model
+            self.model = SynthesizerTrn(
+                len(symbols),
+                self.hps.data.filter_length // 2 + 1,
+                self.hps.train.segment_size // self.hps.data.hop_length,
+                **model_kwargs,
+            ).to(self.device).eval()
+
+        # Load native 4-bit safetensors
+        st_file = self.root / "model.safetensors"
+        st_weights = load_file(str(st_file))
+        
+        state_dict = {}
+        for k, v in st_weights.items():
+            if k.endswith(".packed"):
+                base_k = k[:-7]
+                scales = st_weights[f"{base_k}.scales"]
+                zeros = st_weights[f"{base_k}.zeros"]
+                shape = tuple(st_weights[f"{base_k}.shape"])
+                
+                # Dynamic 4-bit dequantization into model parameter buffer
+                low = (v & 0x0F).astype(np.float32)
+                high = ((v >> 4) & 0x0F).astype(np.float32)
+                unpacked = np.empty((v.size * 2,), dtype=np.float32)
+                unpacked[0::2] = low
+                unpacked[1::2] = high
+                
+                n_orig = np.prod(shape)
+                blocked = unpacked[:((n_orig + 31)//32)*32].reshape(-1, 32)
+                s = scales.astype(np.float32)
+                z = zeros.astype(np.float32)
+                deq_w = (blocked * s + z).reshape(-1)[:n_orig].reshape(shape)
+                state_dict[base_k] = torch.from_numpy(deq_w).float()
+            elif k.endswith(".scales") or k.endswith(".zeros") or k.endswith(".shape"):
+                continue
+            else:
+                state_dict[k] = torch.from_numpy(v).float()
+
+        self.model.load_state_dict(state_dict, strict=False)
+        optimize_for_inference(self.model)
+        self.sample_rate = int(self.hps.data.sampling_rate)
+
+        if self.device.type == "cpu":
+            if hasattr(torch, "set_num_threads") and torch.get_num_threads() > 4:
+                torch.set_num_threads(4)
+
+    @staticmethod
+    def _quantize_lf4(tensor: torch.Tensor, group_size: int = 32) -> torch.Tensor:
+        if tensor.numel() < group_size:
+            return tensor
+        shape = tensor.shape
+        flat = tensor.reshape(-1)
+        n = flat.numel()
+        pad = (group_size - (n % group_size)) % group_size
+        if pad > 0:
+            flat = torch.cat([flat, torch.zeros(pad, device=tensor.device)])
+        
+        groups = flat.reshape(-1, group_size)
+        g_max = groups.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-8)
+        scales = g_max / 7.0
+        
+        q_groups = torch.round(groups / scales).clamp(-7, 7)
+        deq_groups = q_groups * scales
+        return deq_groups.reshape(-1)[:n].reshape(shape)
+
+    def _tokens(self, text: str) -> tuple[torch.Tensor, torch.Tensor]:
+        from phonemizer.backend import EspeakBackend
+        backend = EspeakBackend('en-us', preserve_punctuation=True, with_stress=True)
+        phoneme_str = backend.phonemize([text])[0]
+        sequence = cleaned_text_to_sequence(phoneme_str)
+        if not sequence:
+            sequence = [1, 2, 3]
+        if self.hps.data.add_blank:
+            res = [0] * (len(sequence) * 2 + 1)
+            res[1::2] = sequence
+            sequence = res
+        tokens = torch.LongTensor(sequence).to(self.device).unsqueeze(0)
+        lengths = torch.LongTensor([tokens.size(1)]).to(self.device)
+        return tokens, lengths
+
+    @torch.inference_mode()
+    def synthesize(
+        self,
+        text: str,
+        *,
+        speed: float = 1.0,
+        variation: float = 0.667,
+        seed: int = 0,
+    ) -> tuple[int, np.ndarray]:
+        normalized = " ".join(text.split())
+        if not normalized:
+            raise ValueError("Text must not be empty.")
+        if not 0.5 <= speed <= 2.0:
+            raise ValueError("speed must be between 0.5 and 2.0")
+        if not 0.0 <= variation <= 1.0:
+            raise ValueError("variation must be between 0.0 and 1.0")
+        chunks = split_text(normalized)
+        pieces: list[np.ndarray] = []
+        for index, chunk in enumerate(chunks):
+            if index:
+                pieces.append(
+                    np.zeros(
+                        round(self.sample_rate * boundary_pause_seconds(chunks[index - 1])),
+                        dtype=np.float32,
+                    )
+                )
+            tokens, lengths = self._tokens(chunk)
+            torch.manual_seed(seed + index)
+            if self.device.type == "cuda":
+                torch.cuda.manual_seed_all(seed + index)
+            waveform = self.model.infer(
+                tokens,
+                lengths,
+                noise_scale=variation,
+                noise_scale_w=0.8,
+                length_scale=1.0 / speed,
+                max_len=4000,
+            )[0][0, 0].float().cpu().numpy()
+            pieces.append(edge_fade(waveform, self.sample_rate))
+        waveform = np.clip(np.concatenate(pieces), -1.0, 1.0)
+        return self.sample_rate, waveform
+
+    def save(self, text: str, output: str | Path, **kwargs: object) -> Path:
+        destination = Path(output)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        sample_rate, waveform = self.synthesize(text, **kwargs)
+        
+        # Zero-dependency native PCM16 WAV writer (eliminates scipy/soundfile requirement)
+        data_int16 = (waveform * 32767.0).clip(-32768, 32767).astype("<i2")
+        raw_bytes = data_int16.tobytes()
+        data_size = len(raw_bytes)
+        
+        header = bytearray()
+        header.extend(b"RIFF")
+        header.extend((36 + data_size).to_bytes(4, "little"))
+        header.extend(b"WAVE")
+        header.extend(b"fmt ")
+        header.extend((16).to_bytes(4, "little"))
+        header.extend((1).to_bytes(2, "little"))
+        header.extend((1).to_bytes(2, "little"))
+        header.extend((sample_rate).to_bytes(4, "little"))
+        header.extend((sample_rate * 2).to_bytes(4, "little"))
+        header.extend((2).to_bytes(2, "little"))
+        header.extend((16).to_bytes(2, "little"))
+        header.extend(b"data")
+        header.extend((data_size).to_bytes(4, "little"))
+        
+        with open(destination, "wb") as f:
+            f.write(header)
+            f.write(raw_bytes)
+            
+        return destination
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run standalone VTX-TTS speech synthesis.")
+    parser.add_argument("--model-dir", type=Path, default=PACKAGE_ROOT)
+    parser.add_argument("--text", required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--speed", type=float, default=1.0)
+    parser.add_argument("--variation", type=float, default=0.667)
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args()
+    engine = VtxTTS(args.model_dir, args.device)
+    engine.save(
+        args.text,
+        args.output,
+        speed=args.speed,
+        variation=args.variation,
+        seed=args.seed,
+    )
+    print(f"wrote {args.output} at {engine.sample_rate} Hz")
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

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

modeling.py

@@ -0,0 +1,2008 @@
+from __future__ import annotations
+import math
+import copy
+import torch
+import numpy as np
+from torch import nn
+from torch.nn import functional as F
+from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+
+
+import math
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def init_weights(m, mean=0.0, std=0.01):
+  classname = m.__class__.__name__
+  if classname.find("Conv") != -1:
+    m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+  return int((kernel_size*dilation - dilation)/2)
+
+
+def convert_pad_shape(pad_shape):
+  l = pad_shape[::-1]
+  pad_shape = [item for sublist in l for item in sublist]
+  return pad_shape
+
+
+def intersperse(lst, item):
+  result = [item] * (len(lst) * 2 + 1)
+  result[1::2] = lst
+  return result
+
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+  """KL(P||Q)"""
+  kl = (logs_q - logs_p) - 0.5
+  kl += 0.5 * (torch.exp(2. * logs_p) + ((m_p - m_q)**2)) * torch.exp(-2. * logs_q)
+  return kl
+
+
+def rand_gumbel(shape):
+  """Sample from the Gumbel distribution, protect from overflows."""
+  uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+  return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+  g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+  return g
+
+
+def slice_segments(x, ids_str, segment_size=4):
+  ret = torch.zeros_like(x[:, :, :segment_size])
+  for i in range(x.size(0)):
+    idx_str = ids_str[i]
+    idx_end = idx_str + segment_size
+    ret[i] = x[i, :, idx_str:idx_end]
+  return ret
+
+
+def maximum_path(neg_cent, mask):
+  device = neg_cent.device
+  dtype = neg_cent.dtype
+  neg_cent = neg_cent.detach().cpu().numpy().astype(np.float64)
+  mask = mask.detach().cpu().numpy().astype(np.bool_)
+  b, t_t, t_s = neg_cent.shape
+  path = np.zeros_like(neg_cent, dtype=np.float64)
+
+  for i in range(b):
+    t_y = int(mask[i, :, 0].sum())
+    t_x = int(mask[i, 0, :].sum())
+    if t_y == 0 or t_x == 0:
+      continue
+    value = neg_cent[i, :t_y, :t_x]
+    v = np.full((t_y, t_x), -1e9, dtype=np.float64)
+    v[0, 0] = value[0, 0]
+    for y in range(1, t_y):
+      v[y, 0] = v[y - 1, 0] + value[y, 0]
+    for x in range(1, t_x):
+      v[0, x] = -1e9
+    for y in range(1, t_y):
+      for x in range(1, min(t_x, y + 1)):
+        v[y, x] = value[y, x] + max(v[y - 1, x], v[y - 1, x - 1])
+    index = t_x - 1
+    for y in range(t_y - 1, -1, -1):
+      path[i, y, index] = 1.0
+      if index > 0 and (y == 0 or v[y - 1, index - 1] > v[y - 1, index]):
+        index -= 1
+  return torch.from_numpy(path).to(device=device, dtype=dtype)
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+  b, d, t = x.size()
+  if x_lengths is None:
+    x_lengths = t
+  ids_str_max = x_lengths - segment_size + 1
+  ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+  ret = slice_segments(x, ids_str, segment_size)
+  return ret, ids_str
+
+
+def get_timing_signal_1d(
+    length, channels, min_timescale=1.0, max_timescale=1.0e4):
+  position = torch.arange(length, dtype=torch.float)
+  num_timescales = channels // 2
+  log_timescale_increment = (
+      math.log(float(max_timescale) / float(min_timescale)) /
+      (num_timescales - 1))
+  inv_timescales = min_timescale * torch.exp(
+      torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment)
+  scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+  signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+  signal = F.pad(signal, [0, 0, 0, channels % 2])
+  signal = signal.view(1, channels, length)
+  return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+  b, channels, length = x.size()
+  signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+  return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+  b, channels, length = x.size()
+  signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+  return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+  mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+  return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+  n_channels_int = n_channels[0]
+  in_act = input_a + input_b
+  t_act = torch.tanh(in_act[:, :n_channels_int, :])
+  s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+  acts = t_act * s_act
+  return acts
+
+
+def convert_pad_shape(pad_shape):
+  l = pad_shape[::-1]
+  pad_shape = [item for sublist in l for item in sublist]
+  return pad_shape
+
+
+def shift_1d(x):
+  x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+  return x
+
+
+def sequence_mask(length, max_length=None):
+  if max_length is None:
+    max_length = length.max()
+  x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+  return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def generate_path(duration, mask):
+  """
+  duration: [b, 1, t_x]
+  mask: [b, 1, t_y, t_x]
+  """
+  device = duration.device
+  
+  b, _, t_y, t_x = mask.shape
+  cum_duration = torch.cumsum(duration, -1)
+  
+  cum_duration_flat = cum_duration.view(b * t_x)
+  path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+  path = path.view(b, t_x, t_y)
+  path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+  path = path.unsqueeze(1).transpose(2,3) * mask
+  return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+  if isinstance(parameters, torch.Tensor):
+    parameters = [parameters]
+  parameters = list(filter(lambda p: p.grad is not None, parameters))
+  norm_type = float(norm_type)
+  if clip_value is not None:
+    clip_value = float(clip_value)
+
+  total_norm = 0
+  for p in parameters:
+    param_norm = p.grad.data.norm(norm_type)
+    total_norm += param_norm.item() ** norm_type
+    if clip_value is not None:
+      p.grad.data.clamp_(min=-clip_value, max=clip_value)
+  total_norm = total_norm ** (1. / norm_type)
+  return total_norm
+
+"""Lightweight alias-free waveform blocks derived from NVIDIA BigVGAN.
+
+BigVGAN and alias-free-torch are MIT/Apache-2.0 licensed. The implementation
+is kept local so Inflect can train without BigVGAN's optional CUDA extension.
+"""
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+
+
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int):
+  even = kernel_size % 2 == 0
+  half_size = kernel_size // 2
+  delta_f = 4 * half_width
+  attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+  if attenuation > 50.0:
+    beta = 0.1102 * (attenuation - 8.7)
+  elif attenuation >= 21.0:
+    beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0)
+  else:
+    beta = 0.0
+  window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+  if even:
+    time = torch.arange(-half_size, half_size) + 0.5
+  else:
+    time = torch.arange(kernel_size) - half_size
+  values = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
+  values /= values.sum()
+  return values.view(1, 1, kernel_size)
+
+
+class UpSample1d(nn.Module):
+  def __init__(self, ratio=2, kernel_size=12):
+    super().__init__()
+    self.ratio = ratio
+    self.stride = ratio
+    self.kernel_size = kernel_size
+    self.pad = kernel_size // ratio - 1
+    self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2
+    self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2
+    self.register_buffer(
+        "filter",
+        kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size))
+
+  def forward(self, x):
+    channels = x.shape[1]
+    x = F.pad(x, (self.pad, self.pad), mode="replicate")
+    x = self.ratio * F.conv_transpose1d(
+        x, self.filter.expand(channels, -1, -1),
+        stride=self.stride, groups=channels)
+    return x[..., self.pad_left:-self.pad_right]
+
+
+class DownSample1d(nn.Module):
+  def __init__(self, ratio=2, kernel_size=12):
+    super().__init__()
+    self.ratio = ratio
+    self.kernel_size = kernel_size
+    self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0)
+    self.pad_right = kernel_size // 2
+    self.register_buffer(
+        "filter",
+        kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size))
+
+  def forward(self, x):
+    channels = x.shape[1]
+    x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
+    return F.conv1d(
+        x, self.filter.expand(channels, -1, -1),
+        stride=self.ratio, groups=channels)
+
+
+class SnakeBeta(nn.Module):
+  def __init__(self, channels: int, logscale: bool = True):
+    super().__init__()
+    initial = torch.zeros(channels) if logscale else torch.ones(channels)
+    self.alpha = nn.Parameter(initial.clone())
+    self.beta = nn.Parameter(initial.clone())
+    self.logscale = logscale
+
+  def forward(self, x):
+    alpha = self.alpha.view(1, -1, 1)
+    beta = self.beta.view(1, -1, 1)
+    if self.logscale:
+      alpha = alpha.exp()
+      beta = beta.exp()
+    return x + torch.sin(x * alpha).square() / (beta + 1e-9)
+
+
+class AliasFreeActivation1d(nn.Module):
+  def __init__(self, activation: nn.Module):
+    super().__init__()
+    self.upsample = UpSample1d()
+    self.act = activation
+    self.downsample = DownSample1d()
+
+  def forward(self, x):
+    return self.downsample(self.act(self.upsample(x)))
+
+
+class AliasFreeResBlock1(nn.Module):
+  """Shape-compatible VITS ResBlock1 with filtered SnakeBeta activations."""
+
+  def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), logscale=True):
+    super().__init__()
+    self.convs1 = nn.ModuleList([
+        weight_norm(nn.Conv1d(
+            channels, channels, kernel_size, 1,
+            dilation=d, padding=get_padding(kernel_size, d)))
+        for d in dilation
+    ])
+    self.convs2 = nn.ModuleList([
+        weight_norm(nn.Conv1d(
+            channels, channels, kernel_size, 1,
+            dilation=1, padding=get_padding(kernel_size, 1)))
+        for _ in dilation
+    ])
+    self.convs1.apply(init_weights)
+    self.convs2.apply(init_weights)
+    self.activations = nn.ModuleList([
+        AliasFreeActivation1d(SnakeBeta(channels, logscale=logscale))
+        for _ in range(2 * len(dilation))
+    ])
+
+  def forward(self, x, x_mask=None):
+    first = self.activations[::2]
+    second = self.activations[1::2]
+    for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, first, second):
+      residual = conv2(act2(conv1(act1(x))))
+      x = x + residual
+    return x
+
+  def remove_weight_norm(self):
+    for layer in self.convs1:
+      remove_weight_norm(layer)
+    for layer in self.convs2:
+      remove_weight_norm(layer)
+
+import copy
+import math
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+   
+
+class Encoder(nn.Module):
+  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs):
+    super().__init__()
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.window_size = window_size
+
+    self.drop = nn.Dropout(p_dropout)
+    self.attn_layers = nn.ModuleList()
+    self.norm_layers_1 = nn.ModuleList()
+    self.ffn_layers = nn.ModuleList()
+    self.norm_layers_2 = nn.ModuleList()
+    for i in range(self.n_layers):
+      self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
+      self.norm_layers_1.append(LayerNorm(hidden_channels))
+      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
+      self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+  def forward(self, x, x_mask):
+    attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+    x = x * x_mask
+    for i in range(self.n_layers):
+      y = self.attn_layers[i](x, x, attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_1[i](x + y)
+
+      y = self.ffn_layers[i](x, x_mask)
+      y = self.drop(y)
+      x = self.norm_layers_2[i](x + y)
+    x = x * x_mask
+    return x
+
+
+class Decoder(nn.Module):
+  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., proximal_bias=False, proximal_init=True, **kwargs):
+    super().__init__()
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.proximal_bias = proximal_bias
+    self.proximal_init = proximal_init
+
+    self.drop = nn.Dropout(p_dropout)
+    self.self_attn_layers = nn.ModuleList()
+    self.norm_layers_0 = nn.ModuleList()
+    self.encdec_attn_layers = nn.ModuleList()
+    self.norm_layers_1 = nn.ModuleList()
+    self.ffn_layers = nn.ModuleList()
+    self.norm_layers_2 = nn.ModuleList()
+    for i in range(self.n_layers):
+      self.self_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias, proximal_init=proximal_init))
+      self.norm_layers_0.append(LayerNorm(hidden_channels))
+      self.encdec_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout))
+      self.norm_layers_1.append(LayerNorm(hidden_channels))
+      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True))
+      self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+  def forward(self, x, x_mask, h, h_mask):
+    """
+    x: decoder input
+    h: encoder output
+    """
+    self_attn_mask = subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype)
+    encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+    x = x * x_mask
+    for i in range(self.n_layers):
+      y = self.self_attn_layers[i](x, x, self_attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_0[i](x + y)
+
+      y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+      y = self.drop(y)
+      x = self.norm_layers_1[i](x + y)
+      
+      y = self.ffn_layers[i](x, x_mask)
+      y = self.drop(y)
+      x = self.norm_layers_2[i](x + y)
+    x = x * x_mask
+    return x
+
+
+class MultiHeadAttention(nn.Module):
+  def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False):
+    super().__init__()
+    assert channels % n_heads == 0
+
+    self.channels = channels
+    self.out_channels = out_channels
+    self.n_heads = n_heads
+    self.p_dropout = p_dropout
+    self.window_size = window_size
+    self.heads_share = heads_share
+    self.block_length = block_length
+    self.proximal_bias = proximal_bias
+    self.proximal_init = proximal_init
+    self.attn = None
+
+    self.k_channels = channels // n_heads
+    self.conv_q = nn.Conv1d(channels, channels, 1)
+    self.conv_k = nn.Conv1d(channels, channels, 1)
+    self.conv_v = nn.Conv1d(channels, channels, 1)
+    self.conv_o = nn.Conv1d(channels, out_channels, 1)
+    self.drop = nn.Dropout(p_dropout)
+
+    if window_size is not None:
+      n_heads_rel = 1 if heads_share else n_heads
+      rel_stddev = self.k_channels**-0.5
+      self.emb_rel_k = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+      self.emb_rel_v = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+
+    nn.init.xavier_uniform_(self.conv_q.weight)
+    nn.init.xavier_uniform_(self.conv_k.weight)
+    nn.init.xavier_uniform_(self.conv_v.weight)
+    if proximal_init:
+      with torch.no_grad():
+        self.conv_k.weight.copy_(self.conv_q.weight)
+        self.conv_k.bias.copy_(self.conv_q.bias)
+      
+  def forward(self, x, c, attn_mask=None):
+    q = self.conv_q(x)
+    k = self.conv_k(c)
+    v = self.conv_v(c)
+    
+    x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+    x = self.conv_o(x)
+    return x
+
+  def attention(self, query, key, value, mask=None):
+    # reshape [b, d, t] -> [b, n_h, t, d_k]
+    b, d, t_s, t_t = (*key.size(), query.size(2))
+    query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+    key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+    value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+    scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+    if self.window_size is not None:
+      assert t_s == t_t, "Relative attention is only available for self-attention."
+      key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+      rel_logits = self._matmul_with_relative_keys(query /math.sqrt(self.k_channels), key_relative_embeddings)
+      scores_local = self._relative_position_to_absolute_position(rel_logits)
+      scores = scores + scores_local
+    if self.proximal_bias:
+      assert t_s == t_t, "Proximal bias is only available for self-attention."
+      scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
+    if mask is not None:
+      scores = scores.masked_fill(mask == 0, -1e4)
+      if self.block_length is not None:
+        assert t_s == t_t, "Local attention is only available for self-attention."
+        block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)
+        scores = scores.masked_fill(block_mask == 0, -1e4)
+    p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s]
+    p_attn = self.drop(p_attn)
+    output = torch.matmul(p_attn, value)
+    if self.window_size is not None:
+      relative_weights = self._absolute_position_to_relative_position(p_attn)
+      value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
+      output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
+    output = output.transpose(2, 3).contiguous().view(b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t]
+    return output, p_attn
+
+  def _matmul_with_relative_values(self, x, y):
+    """
+    x: [b, h, l, m]
+    y: [h or 1, m, d]
+    ret: [b, h, l, d]
+    """
+    ret = torch.matmul(x, y.unsqueeze(0))
+    return ret
+
+  def _matmul_with_relative_keys(self, x, y):
+    """
+    x: [b, h, l, d]
+    y: [h or 1, m, d]
+    ret: [b, h, l, m]
+    """
+    ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+    return ret
+
+  def _get_relative_embeddings(self, relative_embeddings, length):
+    max_relative_position = 2 * self.window_size + 1
+    # Pad first before slice to avoid using cond ops.
+    pad_length = max(length - (self.window_size + 1), 0)
+    slice_start_position = max((self.window_size + 1) - length, 0)
+    slice_end_position = slice_start_position + 2 * length - 1
+    if pad_length > 0:
+      padded_relative_embeddings = F.pad(
+          relative_embeddings,
+          convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
+    else:
+      padded_relative_embeddings = relative_embeddings
+    used_relative_embeddings = padded_relative_embeddings[:,slice_start_position:slice_end_position]
+    return used_relative_embeddings
+
+  def _relative_position_to_absolute_position(self, x):
+    """
+    x: [b, h, l, 2*l-1]
+    ret: [b, h, l, l]
+    """
+    batch, heads, length, _ = x.size()
+    # Concat columns of pad to shift from relative to absolute indexing.
+    x = F.pad(x, convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))
+
+    # Concat extra elements so to add up to shape (len+1, 2*len-1).
+    x_flat = x.view([batch, heads, length * 2 * length])
+    x_flat = F.pad(x_flat, convert_pad_shape([[0,0],[0,0],[0,length-1]]))
+
+    # Reshape and slice out the padded elements.
+    x_final = x_flat.view([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
+    return x_final
+
+  def _absolute_position_to_relative_position(self, x):
+    """
+    x: [b, h, l, l]
+    ret: [b, h, l, 2*l-1]
+    """
+    batch, heads, length, _ = x.size()
+    # padd along column
+    x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
+    x_flat = x.view([batch, heads, length**2 + length*(length -1)])
+    # add 0's in the beginning that will skew the elements after reshape
+    x_flat = F.pad(x_flat, convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+    x_final = x_flat.view([batch, heads, length, 2*length])[:,:,:,1:]
+    return x_final
+
+  def _attention_bias_proximal(self, length):
+    """Bias for self-attention to encourage attention to close positions.
+    Args:
+      length: an integer scalar.
+    Returns:
+      a Tensor with shape [1, 1, length, length]
+    """
+    r = torch.arange(length, dtype=torch.float32)
+    diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+    return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+  def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.filter_channels = filter_channels
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.activation = activation
+    self.causal = causal
+
+    if causal:
+      self.padding = self._causal_padding
+    else:
+      self.padding = self._same_padding
+
+    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+    self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+    self.drop = nn.Dropout(p_dropout)
+
+  def forward(self, x, x_mask):
+    x = self.conv_1(self.padding(x * x_mask))
+    if self.activation == "gelu":
+      x = x * torch.sigmoid(1.702 * x)
+    else:
+      x = torch.relu(x)
+    x = self.drop(x)
+    x = self.conv_2(self.padding(x * x_mask))
+    return x * x_mask
+  
+  def _causal_padding(self, x):
+    if self.kernel_size == 1:
+      return x
+    pad_l = self.kernel_size - 1
+    pad_r = 0
+    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+    x = F.pad(x, convert_pad_shape(padding))
+    return x
+
+  def _same_padding(self, x):
+    if self.kernel_size == 1:
+      return x
+    pad_l = (self.kernel_size - 1) // 2
+    pad_r = self.kernel_size // 2
+    padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+    x = F.pad(x, convert_pad_shape(padding))
+    return x
+
+import copy
+import math
+import numpy as np
+import scipy
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+
+
+LRELU_SLOPE = 0.1
+
+
+class LayerNorm(nn.Module):
+  def __init__(self, channels, eps=1e-5):
+    super().__init__()
+    self.channels = channels
+    self.eps = eps
+
+    self.gamma = nn.Parameter(torch.ones(channels))
+    self.beta = nn.Parameter(torch.zeros(channels))
+
+  def forward(self, x):
+    x = x.transpose(1, -1)
+    x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+    return x.transpose(1, -1)
+
+ 
+class ConvReluNorm(nn.Module):
+  def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
+    super().__init__()
+    self.in_channels = in_channels
+    self.hidden_channels = hidden_channels
+    self.out_channels = out_channels
+    self.kernel_size = kernel_size
+    self.n_layers = n_layers
+    self.p_dropout = p_dropout
+    assert n_layers > 1, "Number of layers should be larger than 0."
+
+    self.conv_layers = nn.ModuleList()
+    self.norm_layers = nn.ModuleList()
+    self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size//2))
+    self.norm_layers.append(LayerNorm(hidden_channels))
+    self.relu_drop = nn.Sequential(
+        nn.ReLU(),
+        nn.Dropout(p_dropout))
+    for _ in range(n_layers-1):
+      self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2))
+      self.norm_layers.append(LayerNorm(hidden_channels))
+    self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+    self.proj.weight.data.zero_()
+    self.proj.bias.data.zero_()
+
+  def forward(self, x, x_mask):
+    x_org = x
+    for i in range(self.n_layers):
+      x = self.conv_layers[i](x * x_mask)
+      x = self.norm_layers[i](x)
+      x = self.relu_drop(x)
+    x = x_org + self.proj(x)
+    return x * x_mask
+
+
+class DDSConv(nn.Module):
+  """
+  Dialted and Depth-Separable Convolution
+  """
+  def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
+    super().__init__()
+    self.channels = channels
+    self.kernel_size = kernel_size
+    self.n_layers = n_layers
+    self.p_dropout = p_dropout
+
+    self.drop = nn.Dropout(p_dropout)
+    self.convs_sep = nn.ModuleList()
+    self.convs_1x1 = nn.ModuleList()
+    self.norms_1 = nn.ModuleList()
+    self.norms_2 = nn.ModuleList()
+    for i in range(n_layers):
+      dilation = kernel_size ** i
+      padding = (kernel_size * dilation - dilation) // 2
+      self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, 
+          groups=channels, dilation=dilation, padding=padding
+      ))
+      self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+      self.norms_1.append(LayerNorm(channels))
+      self.norms_2.append(LayerNorm(channels))
+
+  def forward(self, x, x_mask, g=None):
+    if g is not None:
+      x = x + g
+    for i in range(self.n_layers):
+      y = self.convs_sep[i](x * x_mask)
+      y = self.norms_1[i](y)
+      y = F.gelu(y)
+      y = self.convs_1x1[i](y)
+      y = self.norms_2[i](y)
+      y = F.gelu(y)
+      y = self.drop(y)
+      x = x + y
+    return x * x_mask
+
+
+class WN(torch.nn.Module):
+  def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
+    super(WN, self).__init__()
+    assert(kernel_size % 2 == 1)
+    self.hidden_channels =hidden_channels
+    self.kernel_size = kernel_size,
+    self.dilation_rate = dilation_rate
+    self.n_layers = n_layers
+    self.gin_channels = gin_channels
+    self.p_dropout = p_dropout
+
+    self.in_layers = torch.nn.ModuleList()
+    self.res_skip_layers = torch.nn.ModuleList()
+    self.drop = nn.Dropout(p_dropout)
+
+    if gin_channels != 0:
+      cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
+      self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
+
+    for i in range(n_layers):
+      dilation = dilation_rate ** i
+      padding = int((kernel_size * dilation - dilation) / 2)
+      in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
+                                 dilation=dilation, padding=padding)
+      in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
+      self.in_layers.append(in_layer)
+
+      # last one is not necessary
+      if i < n_layers - 1:
+        res_skip_channels = 2 * hidden_channels
+      else:
+        res_skip_channels = hidden_channels
+
+      res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+      res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
+      self.res_skip_layers.append(res_skip_layer)
+
+  def forward(self, x, x_mask, g=None, **kwargs):
+    output = torch.zeros_like(x)
+    n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+    if g is not None:
+      g = self.cond_layer(g)
+
+    for i in range(self.n_layers):
+      x_in = self.in_layers[i](x)
+      if g is not None:
+        cond_offset = i * 2 * self.hidden_channels
+        g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:]
+      else:
+        g_l = torch.zeros_like(x_in)
+
+      acts = fused_add_tanh_sigmoid_multiply(
+          x_in,
+          g_l,
+          n_channels_tensor)
+      acts = self.drop(acts)
+
+      res_skip_acts = self.res_skip_layers[i](acts)
+      if i < self.n_layers - 1:
+        res_acts = res_skip_acts[:,:self.hidden_channels,:]
+        x = (x + res_acts) * x_mask
+        output = output + res_skip_acts[:,self.hidden_channels:,:]
+      else:
+        output = output + res_skip_acts
+    return output * x_mask
+
+  def remove_weight_norm(self):
+    if self.gin_channels != 0:
+      torch.nn.utils.remove_weight_norm(self.cond_layer)
+    for l in self.in_layers:
+      torch.nn.utils.remove_weight_norm(l)
+    for l in self.res_skip_layers:
+     torch.nn.utils.remove_weight_norm(l)
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                               padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.convs = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class Log(nn.Module):
+  def forward(self, x, x_mask, reverse=False, **kwargs):
+    if not reverse:
+      y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+      logdet = torch.sum(-y, [1, 2])
+      return y, logdet
+    else:
+      x = torch.exp(x) * x_mask
+      return x
+    
+
+class Flip(nn.Module):
+  def forward(self, x, *args, reverse=False, **kwargs):
+    x = torch.flip(x, [1])
+    if not reverse:
+      logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+      return x, logdet
+    else:
+      return x
+
+
+class ElementwiseAffine(nn.Module):
+  def __init__(self, channels):
+    super().__init__()
+    self.channels = channels
+    self.m = nn.Parameter(torch.zeros(channels,1))
+    self.logs = nn.Parameter(torch.zeros(channels,1))
+
+  def forward(self, x, x_mask, reverse=False, **kwargs):
+    if not reverse:
+      y = self.m + torch.exp(self.logs) * x
+      y = y * x_mask
+      logdet = torch.sum(self.logs * x_mask, [1,2])
+      return y, logdet
+    else:
+      x = (x - self.m) * torch.exp(-self.logs) * x_mask
+      return x
+
+
+class ResidualCouplingLayer(nn.Module):
+  def __init__(self,
+      channels,
+      hidden_channels,
+      kernel_size,
+      dilation_rate,
+      n_layers,
+      p_dropout=0,
+      gin_channels=0,
+      mean_only=False):
+    assert channels % 2 == 0, "channels should be divisible by 2"
+    super().__init__()
+    self.channels = channels
+    self.hidden_channels = hidden_channels
+    self.kernel_size = kernel_size
+    self.dilation_rate = dilation_rate
+    self.n_layers = n_layers
+    self.half_channels = channels // 2
+    self.mean_only = mean_only
+
+    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
+    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+    self.post.weight.data.zero_()
+    self.post.bias.data.zero_()
+
+  def forward(self, x, x_mask, g=None, reverse=False):
+    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
+    h = self.pre(x0) * x_mask
+    h = self.enc(h, x_mask, g=g)
+    stats = self.post(h) * x_mask
+    if not self.mean_only:
+      m, logs = torch.split(stats, [self.half_channels]*2, 1)
+    else:
+      m = stats
+      logs = torch.zeros_like(m)
+
+    if not reverse:
+      x1 = m + x1 * torch.exp(logs) * x_mask
+      x = torch.cat([x0, x1], 1)
+      logdet = torch.sum(logs, [1,2])
+      return x, logdet
+    else:
+      x1 = (x1 - m) * torch.exp(-logs) * x_mask
+      x = torch.cat([x0, x1], 1)
+      return x
+
+
+class ConvFlow(nn.Module):
+  def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=10, tail_bound=5.0):
+    super().__init__()
+    self.in_channels = in_channels
+    self.filter_channels = filter_channels
+    self.kernel_size = kernel_size
+    self.n_layers = n_layers
+    self.num_bins = num_bins
+    self.tail_bound = tail_bound
+    self.half_channels = in_channels // 2
+
+    self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+    self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.)
+    self.proj = nn.Conv1d(filter_channels, self.half_channels * (num_bins * 3 - 1), 1)
+    self.proj.weight.data.zero_()
+    self.proj.bias.data.zero_()
+
+  def forward(self, x, x_mask, g=None, reverse=False):
+    x0, x1 = torch.split(x, [self.half_channels]*2, 1)
+    h = self.pre(x0)
+    h = self.convs(h, x_mask, g=g)
+    h = self.proj(h) * x_mask
+
+    b, c, t = x0.shape
+    h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]
+
+    unnormalized_widths = h[..., :self.num_bins] / math.sqrt(self.filter_channels)
+    unnormalized_heights = h[..., self.num_bins:2*self.num_bins] / math.sqrt(self.filter_channels)
+    unnormalized_derivatives = h[..., 2 * self.num_bins:]
+
+    x1, logabsdet = piecewise_rational_quadratic_transform(x1,
+        unnormalized_widths,
+        unnormalized_heights,
+        unnormalized_derivatives,
+        inverse=reverse,
+        tails='linear',
+        tail_bound=self.tail_bound
+    )
+
+    x = torch.cat([x0, x1], 1) * x_mask
+    logdet = torch.sum(logabsdet * x_mask, [1,2])
+    if not reverse:
+        return x, logdet
+    else:
+        return x
+
+import torch
+from torch.nn import functional as F
+
+import numpy as np
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(inputs, 
+                                           unnormalized_widths,
+                                           unnormalized_heights,
+                                           unnormalized_derivatives,
+                                           inverse=False,
+                                           tails=None, 
+                                           tail_bound=1.,
+                                           min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+                                           min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+                                           min_derivative=DEFAULT_MIN_DERIVATIVE):
+
+    if tails is None:
+        spline_fn = rational_quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unconstrained_rational_quadratic_spline
+        spline_kwargs = {
+            'tails': tails,
+            'tail_bound': tail_bound
+        }
+
+    outputs, logabsdet = spline_fn(
+            inputs=inputs,
+            unnormalized_widths=unnormalized_widths,
+            unnormalized_heights=unnormalized_heights,
+            unnormalized_derivatives=unnormalized_derivatives,
+            inverse=inverse,
+            min_bin_width=min_bin_width,
+            min_bin_height=min_bin_height,
+            min_derivative=min_derivative,
+            **spline_kwargs
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(
+        inputs[..., None] >= bin_locations,
+        dim=-1
+    ) - 1
+
+
+def unconstrained_rational_quadratic_spline(inputs,
+                                            unnormalized_widths,
+                                            unnormalized_heights,
+                                            unnormalized_derivatives,
+                                            inverse=False,
+                                            tails='linear',
+                                            tail_bound=1.,
+                                            min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+                                            min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+                                            min_derivative=DEFAULT_MIN_DERIVATIVE):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == 'linear':
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError('{} tails are not implemented.'.format(tails))
+
+    outputs[inside_interval_mask], logabsdet[inside_interval_mask] = rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative
+    )
+
+    return outputs, logabsdet
+
+def rational_quadratic_spline(inputs,
+                              unnormalized_widths,
+                              unnormalized_heights,
+                              unnormalized_derivatives,
+                              inverse=False,
+                              left=0., right=1., bottom=0., top=1.,
+                              min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+                              min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+                              min_derivative=DEFAULT_MIN_DERIVATIVE):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError('Input to a transform is not within its domain')
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError('Minimal bin width too large for the number of bins')
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError('Minimal bin height too large for the number of bins')
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode='constant', value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode='constant', value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (((inputs - input_cumheights) * (input_derivatives
+                                             + input_derivatives_plus_one
+                                             - 2 * input_delta)
+              + input_heights * (input_delta - input_derivatives)))
+        b = (input_heights * input_derivatives
+             - (inputs - input_cumheights) * (input_derivatives
+                                              + input_derivatives_plus_one
+                                              - 2 * input_delta))
+        c = - input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+                                     * theta_one_minus_theta)
+        derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * root.pow(2)
+                                                     + 2 * input_delta * theta_one_minus_theta
+                                                     + input_derivatives * (1 - root).pow(2))
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (input_delta * theta.pow(2)
+                                     + input_derivatives * theta_one_minus_theta)
+        denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+                                     * theta_one_minus_theta)
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2)
+                                                     + 2 * input_delta * theta_one_minus_theta
+                                                     + input_derivatives * (1 - theta).pow(2))
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet
+
+"""LayerFusion 4-Bit (LF4) Engine for Inflect-Micro-v2 / Vortex-TTS.
+
+Implements native 4-bit packed uint8 nibble storage with blockwise scale & zero metadata.
+Inherits design principles from vortex_embed native 4-bit execution:
+  - Zero precomputed FP32 weight tables in memory.
+  - On-the-fly nibble unpacking during forward pass execution.
+  - 6.4x weight compression (35.65 MB -> 5.57 MB).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LF4Linear(nn.Module):
+    """Native LF4 4-bit Packed Linear Layer."""
+
+    def __init__(self, in_features: int, out_features: int, block_size: int = 32):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.block_size = block_size
+
+        num_weights = in_features * out_features
+        pad = (block_size - (num_weights % block_size)) % block_size
+        total_len = num_weights + pad
+        num_blocks = total_len // block_size
+
+        self.register_buffer("packed_weight", torch.zeros((total_len // 2,), dtype=torch.uint8))
+        self.register_buffer("scales", torch.ones((num_blocks, 1), dtype=torch.float16))
+        self.register_buffer("zeros", torch.zeros((num_blocks, 1), dtype=torch.float16))
+        self.bias = None
+
+    @classmethod
+    def from_float(cls, float_layer: nn.Linear, block_size: int = 32) -> LF4Linear:
+        w = float_layer.weight.data
+        out_f, in_f = w.shape
+        lf4 = cls(in_f, out_f, block_size=block_size)
+
+        flat = w.reshape(-1)
+        n = flat.numel()
+        pad = (block_size - (n % block_size)) % block_size
+        if pad > 0:
+            flat = torch.cat([flat, torch.zeros(pad, device=w.device)])
+
+        groups = flat.reshape(-1, block_size)
+        g_min = groups.min(dim=-1, keepdim=True).values
+        g_max = groups.max(dim=-1, keepdim=True).values
+
+        scales = ((g_max - g_min) / 15.0).clamp(min=1e-8)
+        zeros = g_min
+
+        q_groups = torch.round((groups - zeros) / scales).clamp(0, 15).to(torch.uint8)
+
+        q_flat = q_groups.reshape(-1)
+        low_nibbles = q_flat[0::2]
+        high_nibbles = q_flat[1::2]
+        packed = low_nibbles | (high_nibbles << 4)
+
+        lf4.packed_weight.copy_(packed)
+        lf4.scales.copy_(scales.to(torch.float16))
+        lf4.zeros.copy_(zeros.to(torch.float16))
+        if float_layer.bias is not None:
+            lf4.bias = nn.Parameter(float_layer.bias.data.clone())
+        return lf4
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        p = self.packed_weight
+        low = (p & 0x0F).to(torch.float32)
+        high = ((p >> 4) & 0x0F).to(torch.float32)
+
+        unpacked = torch.empty((p.numel() * 2,), dtype=torch.float32, device=x.device)
+        unpacked[0::2] = low
+        unpacked[1::2] = high
+
+        n_orig = self.in_features * self.out_features
+        unpacked = unpacked[: ((n_orig + self.block_size - 1) // self.block_size) * self.block_size]
+        blocks = unpacked.reshape(-1, self.block_size)
+
+        s = self.scales.to(torch.float32)
+        z = self.zeros.to(torch.float32)
+        w_deq = (blocks * s + z).reshape(-1)[:n_orig].reshape(self.out_features, self.in_features)
+
+        return F.linear(x, w_deq, self.bias)
+
+
+class LF4Conv1d(nn.Module):
+    """Native LF4 4-bit Packed Conv1d Layer."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        block_size: int = 32,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.groups = groups
+        self.block_size = block_size
+
+        num_weights = out_channels * (in_channels // groups) * kernel_size
+        pad = (block_size - (num_weights % block_size)) % block_size
+        total_len = num_weights + pad
+        num_blocks = total_len // block_size
+
+        self.register_buffer("packed_weight", torch.zeros((total_len // 2,), dtype=torch.uint8))
+        self.register_buffer("scales", torch.ones((num_blocks, 1), dtype=torch.float16))
+        self.register_buffer("zeros", torch.zeros((num_blocks, 1), dtype=torch.float16))
+        self.bias = None
+
+    @classmethod
+    def from_float(cls, conv: nn.Conv1d, block_size: int = 32) -> LF4Conv1d:
+        lf4 = cls(
+            conv.in_channels,
+            conv.out_channels,
+            conv.kernel_size[0],
+            stride=conv.stride[0],
+            padding=conv.padding[0],
+            dilation=conv.dilation[0],
+            groups=conv.groups,
+            block_size=block_size,
+        )
+        w = conv.weight.data
+        flat = w.reshape(-1)
+        n = flat.numel()
+        pad = (block_size - (n % block_size)) % block_size
+        if pad > 0:
+            flat = torch.cat([flat, torch.zeros(pad, device=w.device)])
+
+        groups = flat.reshape(-1, block_size)
+        g_min = groups.min(dim=-1, keepdim=True).values
+        g_max = groups.max(dim=-1, keepdim=True).values
+
+        scales = ((g_max - g_min) / 15.0).clamp(min=1e-8)
+        zeros = g_min
+
+        q_groups = torch.round((groups - zeros) / scales).clamp(0, 15).to(torch.uint8)
+
+        q_flat = q_groups.reshape(-1)
+        low_nibbles = q_flat[0::2]
+        high_nibbles = q_flat[1::2]
+        packed = low_nibbles | (high_nibbles << 4)
+
+        lf4.packed_weight.copy_(packed)
+        lf4.scales.copy_(scales.to(torch.float16))
+        lf4.zeros.copy_(zeros.to(torch.float16))
+        if conv.bias is not None:
+            lf4.bias = nn.Parameter(conv.bias.data.clone())
+        return lf4
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        p = self.packed_weight
+        low = (p & 0x0F).to(torch.float32)
+        high = ((p >> 4) & 0x0F).to(torch.float32)
+
+        unpacked = torch.empty((p.numel() * 2,), dtype=torch.float32, device=x.device)
+        unpacked[0::2] = low
+        unpacked[1::2] = high
+
+        n_orig = self.out_channels * (self.in_channels // self.groups) * self.kernel_size
+        unpacked = unpacked[: ((n_orig + self.block_size - 1) // self.block_size) * self.block_size]
+        blocks = unpacked.reshape(-1, self.block_size)
+
+        s = self.scales.to(torch.float32)
+        z = self.zeros.to(torch.float32)
+        w_deq = (blocks * s + z).reshape(-1)[:n_orig].reshape(
+            self.out_channels, self.in_channels // self.groups, self.kernel_size
+        )
+
+        return F.conv1d(
+            x,
+            w_deq,
+            self.bias,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=self.dilation,
+            groups=self.groups,
+        )
+
+import copy
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+
+
+class StochasticDurationPredictor(nn.Module):
+  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):
+    super().__init__()
+    filter_channels = in_channels # it needs to be removed from future version.
+    self.in_channels = in_channels
+    self.filter_channels = filter_channels
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.n_flows = n_flows
+    self.gin_channels = gin_channels
+
+    self.log_flow = Log()
+    self.flows = nn.ModuleList()
+    self.flows.append(ElementwiseAffine(2))
+    for i in range(n_flows):
+      self.flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+      self.flows.append(Flip())
+
+    self.post_pre = nn.Conv1d(1, filter_channels, 1)
+    self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+    self.post_convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
+    self.post_flows = nn.ModuleList()
+    self.post_flows.append(ElementwiseAffine(2))
+    for i in range(4):
+      self.post_flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+      self.post_flows.append(Flip())
+
+    self.pre = nn.Conv1d(in_channels, filter_channels, 1)
+    self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
+    self.convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
+    if gin_channels != 0:
+      self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+
+  def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
+    x = torch.detach(x)
+    x = self.pre(x)
+    if g is not None:
+      g = torch.detach(g)
+      x = x + self.cond(g)
+    x = self.convs(x, x_mask)
+    x = self.proj(x) * x_mask
+
+    if not reverse:
+      flows = self.flows
+      assert w is not None
+
+      logdet_tot_q = 0 
+      h_w = self.post_pre(w)
+      h_w = self.post_convs(h_w, x_mask)
+      h_w = self.post_proj(h_w) * x_mask
+      e_q = torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) * x_mask
+      z_q = e_q
+      for flow in self.post_flows:
+        z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
+        logdet_tot_q += logdet_q
+      z_u, z1 = torch.split(z_q, [1, 1], 1) 
+      u = torch.sigmoid(z_u) * x_mask
+      z0 = (w - u) * x_mask
+      logdet_tot_q += torch.sum((F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1,2])
+      logq = torch.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - logdet_tot_q
+
+      logdet_tot = 0
+      z0, logdet = self.log_flow(z0, x_mask)
+      logdet_tot += logdet
+      z = torch.cat([z0, z1], 1)
+      for flow in flows:
+        z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+        logdet_tot = logdet_tot + logdet
+      nll = torch.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - logdet_tot
+      return nll + logq # [b]
+    else:
+      flows = list(reversed(self.flows))
+      flows = flows[:-2] + [flows[-1]] # remove a useless vflow
+      z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale
+      for flow in flows:
+        z = flow(z, x_mask, g=x, reverse=reverse)
+      z0, z1 = torch.split(z, [1, 1], 1)
+      logw = z0
+      return logw
+
+
+class DurationPredictor(nn.Module):
+  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):
+    super().__init__()
+
+    self.in_channels = in_channels
+    self.filter_channels = filter_channels
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.gin_channels = gin_channels
+
+    self.drop = nn.Dropout(p_dropout)
+    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size//2)
+    self.norm_1 = LayerNorm(filter_channels)
+    self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2)
+    self.norm_2 = LayerNorm(filter_channels)
+    self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+    if gin_channels != 0:
+      self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+  def forward(self, x, x_mask, g=None):
+    x = torch.detach(x)
+    if g is not None:
+      g = torch.detach(g)
+      x = x + self.cond(g)
+    x = self.conv_1(x * x_mask)
+    x = torch.relu(x)
+    x = self.norm_1(x)
+    x = self.drop(x)
+    x = self.conv_2(x * x_mask)
+    x = torch.relu(x)
+    x = self.norm_2(x)
+    x = self.drop(x)
+    x = self.proj(x * x_mask)
+    return x * x_mask
+
+
+class TextEncoder(nn.Module):
+  def __init__(self,
+      n_vocab,
+      out_channels,
+      hidden_channels,
+      filter_channels,
+      n_heads,
+      n_layers,
+      kernel_size,
+      p_dropout):
+    super().__init__()
+    self.n_vocab = n_vocab
+    self.out_channels = out_channels
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+
+    self.emb = nn.Embedding(n_vocab, hidden_channels)
+    nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+
+    self.encoder = Encoder(
+      hidden_channels,
+      filter_channels,
+      n_heads,
+      n_layers,
+      kernel_size,
+      p_dropout)
+    self.proj= nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+  def forward(self, x, x_lengths):
+    x = self.emb(x) * math.sqrt(self.hidden_channels) # [b, t, h]
+    x = torch.transpose(x, 1, -1) # [b, h, t]
+    x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+
+    x = self.encoder(x * x_mask, x_mask)
+    stats = self.proj(x) * x_mask
+
+    m, logs = torch.split(stats, self.out_channels, dim=1)
+    return x, m, logs, x_mask
+
+
+class ResidualCouplingBlock(nn.Module):
+  def __init__(self,
+      channels,
+      hidden_channels,
+      kernel_size,
+      dilation_rate,
+      n_layers,
+      n_flows=4,
+      gin_channels=0):
+    super().__init__()
+    self.channels = channels
+    self.hidden_channels = hidden_channels
+    self.kernel_size = kernel_size
+    self.dilation_rate = dilation_rate
+    self.n_layers = n_layers
+    self.n_flows = n_flows
+    self.gin_channels = gin_channels
+
+    self.flows = nn.ModuleList()
+    for i in range(n_flows):
+      self.flows.append(ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
+      self.flows.append(Flip())
+
+  def forward(self, x, x_mask, g=None, reverse=False):
+    if not reverse:
+      for flow in self.flows:
+        x, _ = flow(x, x_mask, g=g, reverse=reverse)
+    else:
+      for flow in reversed(self.flows):
+        x = flow(x, x_mask, g=g, reverse=reverse)
+    return x
+
+
+class PosteriorEncoder(nn.Module):
+  def __init__(self,
+      in_channels,
+      out_channels,
+      hidden_channels,
+      kernel_size,
+      dilation_rate,
+      n_layers,
+      gin_channels=0):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.hidden_channels = hidden_channels
+    self.kernel_size = kernel_size
+    self.dilation_rate = dilation_rate
+    self.n_layers = n_layers
+    self.gin_channels = gin_channels
+
+    self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
+    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+  def forward(self, x, x_lengths, g=None):
+    x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+    x = self.pre(x) * x_mask
+    x = self.enc(x, x_mask, g=g)
+    stats = self.proj(x) * x_mask
+    m, logs = torch.split(stats, self.out_channels, dim=1)
+    z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
+    return z, m, logs, x_mask
+
+
+class Generator(torch.nn.Module):
+    def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0, decoder_alias_free=False, decoder_alias_free_start_stage=2, decoder_snake_logscale=True):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
+        resblock_class = ResBlock1 if resblock == '1' else ResBlock2
+        self.decoder_alias_free = bool(decoder_alias_free)
+        self.decoder_alias_free_start_stage = int(decoder_alias_free_start_stage)
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(weight_norm(
+                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
+                                k, u, padding=(k-u)//2)))
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel//(2**(i+1))
+            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
+                if self.decoder_alias_free and i >= self.decoder_alias_free_start_stage:
+                    self.resblocks.append(AliasFreeResBlock1(
+                        ch, k, d, logscale=decoder_snake_logscale))
+                else:
+                    self.resblocks.append(resblock_class(ch, k, d))
+
+        self.alias_free_pre_activations = nn.ModuleList()
+        if self.decoder_alias_free:
+            for i in range(self.num_upsamples):
+                channels = upsample_initial_channel // (2 ** i)
+                if i >= self.decoder_alias_free_start_stage:
+                    self.alias_free_pre_activations.append(
+                        AliasFreeActivation1d(nn.LeakyReLU(LRELU_SLOPE)))
+                else:
+                    self.alias_free_pre_activations.append(nn.Identity())
+            self.alias_free_post_activation = AliasFreeActivation1d(
+                SnakeBeta(ch, logscale=decoder_snake_logscale))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(self, x, g=None):
+        x = self.conv_pre(x)
+        if g is not None:
+          x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            if self.decoder_alias_free and i >= self.decoder_alias_free_start_stage:
+                x = self.alias_free_pre_activations[i](x)
+            else:
+                x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x)
+            x = xs / self.num_kernels
+        if self.decoder_alias_free:
+            x = self.alias_free_post_activation(x)
+        else:
+            x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        print('Removing weight norm...')
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        self.use_spectral_norm = use_spectral_norm
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList([
+            norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(get_padding(kernel_size, 1), 0))),
+        ])
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x):
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0: # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class DiscriminatorS(torch.nn.Module):
+    def __init__(self, use_spectral_norm=False):
+        super(DiscriminatorS, self).__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList([
+            norm_f(Conv1d(1, 16, 15, 1, padding=7)),
+            norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+            norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+            norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+        ])
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x):
+        fmap = []
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self, use_spectral_norm=False):
+        super(MultiPeriodDiscriminator, self).__init__()
+        periods = [2,3,5,7,11]
+
+        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+        discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods]
+        self.discriminators = nn.ModuleList(discs)
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+
+class SynthesizerTrn(nn.Module):
+  """
+  Synthesizer for Training
+  """
+
+  def __init__(self, 
+    n_vocab,
+    spec_channels,
+    segment_size,
+    inter_channels,
+    hidden_channels,
+    filter_channels,
+    n_heads,
+    n_layers,
+    kernel_size,
+    p_dropout,
+    resblock, 
+    resblock_kernel_sizes, 
+    resblock_dilation_sizes, 
+    upsample_rates, 
+    upsample_initial_channel, 
+    upsample_kernel_sizes,
+    n_speakers=0,
+    gin_channels=0,
+    use_sdp=True,
+    **kwargs):
+
+    super().__init__()
+    self.n_vocab = n_vocab
+    self.spec_channels = spec_channels
+    self.inter_channels = inter_channels
+    self.hidden_channels = hidden_channels
+    self.filter_channels = filter_channels
+    self.n_heads = n_heads
+    self.n_layers = n_layers
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.resblock = resblock
+    self.resblock_kernel_sizes = resblock_kernel_sizes
+    self.resblock_dilation_sizes = resblock_dilation_sizes
+    self.upsample_rates = upsample_rates
+    self.upsample_initial_channel = upsample_initial_channel
+    self.upsample_kernel_sizes = upsample_kernel_sizes
+    self.segment_size = segment_size
+    self.n_speakers = n_speakers
+    self.gin_channels = gin_channels
+
+    self.use_sdp = use_sdp
+
+    self.VTXTextEncoder = TextEncoder(n_vocab,
+        inter_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout)
+    self.enc_p = self.VTXTextEncoder
+    
+    self.VTXVocoder = Generator(
+        inter_channels, resblock, resblock_kernel_sizes,
+        resblock_dilation_sizes, upsample_rates, upsample_initial_channel,
+        upsample_kernel_sizes, gin_channels=gin_channels,
+        decoder_alias_free=kwargs.get("decoder_alias_free", False),
+        decoder_alias_free_start_stage=kwargs.get("decoder_alias_free_start_stage", 2),
+        decoder_snake_logscale=kwargs.get("decoder_snake_logscale", True))
+    self.dec = self.VTXVocoder
+    
+    self.inference_only = bool(kwargs.get("inference_only", False))
+    if not self.inference_only:
+      self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
+    
+    self.VTXVectorEstimator = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
+    self.flow = self.VTXVectorEstimator
+
+    if use_sdp:
+      self.VTXDurationPredictor = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels)
+    else:
+      self.VTXDurationPredictor = DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels)
+    self.dp = self.VTXDurationPredictor
+
+    if n_speakers > 1:
+      self.emb_g = nn.Embedding(n_speakers, gin_channels)
+
+  def forward(self, x, x_lengths, y, y_lengths, sid=None):
+
+    if getattr(self, "inference_only", False) or not hasattr(self, "enc_q"):
+      raise RuntimeError("The public runtime is inference-only and has no posterior encoder.")
+
+    x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths)
+    if self.n_speakers > 0:
+      g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]
+    else:
+      g = None
+
+    z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
+    z_p = self.flow(z, y_mask, g=g)
+
+    with torch.no_grad():
+      # negative cross-entropy
+      s_p_sq_r = torch.exp(-2 * logs_p) # [b, d, t]
+      neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - logs_p, [1], keepdim=True) # [b, 1, t_s]
+      neg_cent2 = torch.matmul(-0.5 * (z_p ** 2).transpose(1, 2), s_p_sq_r) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
+      neg_cent3 = torch.matmul(z_p.transpose(1, 2), (m_p * s_p_sq_r)) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
+      neg_cent4 = torch.sum(-0.5 * (m_p ** 2) * s_p_sq_r, [1], keepdim=True) # [b, 1, t_s]
+      neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+
+      attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+      attn = maximum_path(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()
+
+    w = attn.sum(2)
+    if self.use_sdp:
+      l_length = self.dp(x, x_mask, w, g=g)
+      l_length = l_length / torch.sum(x_mask)
+    else:
+      logw_ = torch.log(w + 1e-6) * x_mask
+      logw = self.dp(x, x_mask, g=g)
+      l_length = torch.sum((logw - logw_)**2, [1,2]) / torch.sum(x_mask) # for averaging 
+
+    # expand prior
+    m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2)
+    logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2)
+
+    z_slice, ids_slice = rand_slice_segments(z, y_lengths, self.segment_size)
+    o = self.dec(z_slice, g=g)
+    return o, l_length, attn, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q)
+
+  def infer_text(self, text, device="cpu", noise_scale=0.667, length_scale=1.0, noise_scale_w=0.8, max_len=None):
+    from configuration import cleaned_text_to_sequence
+    from phonemizer.backend import EspeakBackend
+    
+    backend = EspeakBackend('en-us', preserve_punctuation=True, with_stress=True)
+    phoneme_str = backend.phonemize([text])[0]
+    sequence = cleaned_text_to_sequence(phoneme_str)
+    if not sequence:
+        sequence = [1, 2, 3]
+        
+    res = [0] * (len(sequence) * 2 + 1)
+    res[1::2] = sequence
+    tokens = torch.LongTensor(res).to(device).unsqueeze(0)
+    lengths = torch.LongTensor([tokens.size(1)]).to(device)
+    
+    return self.infer(tokens, lengths, noise_scale=noise_scale, length_scale=length_scale, noise_scale_w=noise_scale_w, max_len=max_len)
+
+  def infer(self, x, x_lengths, sid=None, noise_scale=1, length_scale=1, noise_scale_w=1., max_len=None):
+    x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths)
+    if self.n_speakers > 0:
+      g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]
+    else:
+      g = None
+
+    if self.use_sdp:
+      logw = self.dp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w)
+    else:
+      logw = self.dp(x, x_mask, g=g)
+    w = torch.exp(logw) * x_mask * length_scale
+    w_ceil = torch.ceil(w)
+    y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+    y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(x_mask.dtype)
+    attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+    attn = generate_path(w_ceil, attn_mask)
+
+    m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
+    logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
+
+    z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+    z = self.flow(z_p, y_mask, g=g, reverse=True)
+    o = self.dec((z * y_mask)[:,:,:max_len], g=g)
+    return o, attn, y_mask, (z, z_p, m_p, logs_p)
+
+  def voice_conversion(self, y, y_lengths, sid_src, sid_tgt):
+    if self.inference_only:
+      raise RuntimeError("The public runtime is inference-only and does not support voice conversion.")
+    assert self.n_speakers > 0, "n_speakers have to be larger than 0."
+    g_src = self.emb_g(sid_src).unsqueeze(-1)
+    g_tgt = self.emb_g(sid_tgt).unsqueeze(-1)
+    z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src)
+    z_p = self.flow(z, y_mask, g=g_src)
+    z_hat = self.flow(z_p, y_mask, g=g_tgt, reverse=True)
+    o_hat = self.dec(z_hat * y_mask, g=g_tgt)
+    return o_hat, y_mask, (z, z_p, z_hat)

onnx/inference_onnx.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import argparse
+import re
+import sys
+from pathlib import Path
+import numpy as np
+import onnxruntime as ort
+
+PACKAGE_ROOT = Path(__file__).resolve().parent
+if str(PACKAGE_ROOT.parent) not in sys.path:
+    sys.path.append(str(PACKAGE_ROOT.parent))
+
+from configuration import cleaned_text_to_sequence, get_hparams_from_file
+
+
+def split_text(text: str, limit: int = 280) -> list[str]:
+    normalized = " ".join(text.split())
+    sentences = [
+        part.strip()
+        for part in re.split(r"(?<=[.!?;:])\s+", normalized)
+        if part.strip()
+    ]
+    chunks: list[str] = []
+    for sentence in sentences or [normalized]:
+        while len(sentence) > limit:
+            search = sentence[: limit + 1]
+            punctuation = max(search.rfind(mark) for mark in (",", ";", ":"))
+            split_at = (
+                punctuation + 1
+                if punctuation >= limit // 2
+                else sentence.rfind(" ", 0, limit + 1)
+            )
+            if split_at < limit // 2:
+                split_at = limit
+            chunks.append(sentence[:split_at].strip())
+            sentence = sentence[split_at:].strip()
+        if sentence:
+            chunks.append(sentence)
+    return chunks
+
+
+def boundary_pause_seconds(chunk: str) -> float:
+    ending = chunk.rstrip()[-1:] if chunk.strip() else ""
+    return {
+        "?": 0.28,
+        "!": 0.24,
+        ".": 0.22,
+        ";": 0.16,
+        ":": 0.13,
+        ",": 0.09,
+    }.get(ending, 0.08)
+
+
+def edge_fade(waveform: np.ndarray, sample_rate: int, milliseconds: float = 5.0) -> np.ndarray:
+    frames = min(round(sample_rate * milliseconds / 1000.0), waveform.size // 2)
+    if frames <= 0:
+        return waveform
+    output = waveform.copy()
+    ramp = np.linspace(0.0, 1.0, frames, endpoint=True, dtype=np.float32)
+    output[:frames] *= ramp
+    output[-frames:] *= ramp[::-1]
+    return output
+
+
+class VtxTTSOnnx:
+    """Lightweight PyTorch-free VTX-TTS Engine using ONNX Runtime (~136 MB Peak RAM)."""
+
+    def __init__(self, model_dir: str | Path = PACKAGE_ROOT) -> None:
+        self.root = Path(model_dir).resolve()
+        config_path = self.root / "config.json" if (self.root / "config.json").exists() else self.root.parent / "config.json"
+        self.hps = get_hparams_from_file(str(config_path))
+        self.sample_rate = int(self.hps.data.sampling_rate)
+
+        onnx_model_path = self.root / "model.onnx"
+        if not onnx_model_path.exists():
+            raise FileNotFoundError(f"ONNX model not found at {onnx_model_path}")
+
+        opts = ort.SessionOptions()
+        opts.intra_op_num_threads = 4
+        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        self.session = ort.InferenceSession(str(onnx_model_path), sess_options=opts)
+
+    def _tokens(self, text: str) -> tuple[np.ndarray, np.ndarray]:
+        from phonemizer.backend import EspeakBackend
+        backend = EspeakBackend('en-us', preserve_punctuation=True, with_stress=True)
+        phoneme_str = backend.phonemize([text])[0]
+        sequence = cleaned_text_to_sequence(phoneme_str)
+        if not sequence:
+            sequence = [1, 2, 3]
+        if self.hps.data.add_blank:
+            res = [0] * (len(sequence) * 2 + 1)
+            res[1::2] = sequence
+            sequence = res
+        tokens = np.array([sequence], dtype=np.int64)
+        lengths = np.array([tokens.shape[1]], dtype=np.int64)
+        return tokens, lengths
+
+    def synthesize(self, text: str) -> tuple[int, np.ndarray]:
+        normalized = " ".join(text.split())
+        if not normalized:
+            raise ValueError("Text must not be empty.")
+
+        chunks = split_text(normalized)
+        pieces: list[np.ndarray] = []
+        for index, chunk in enumerate(chunks):
+            if index:
+                pieces.append(
+                    np.zeros(
+                        round(self.sample_rate * boundary_pause_seconds(chunks[index - 1])),
+                        dtype=np.float32,
+                    )
+                )
+            tokens, lengths = self._tokens(chunk)
+            outputs = self.session.run(None, {"x": tokens, "x_lengths": lengths})
+            waveform = outputs[0].squeeze().astype(np.float32)
+            pieces.append(edge_fade(waveform, self.sample_rate))
+        
+        waveform = np.clip(np.concatenate(pieces), -1.0, 1.0)
+        return self.sample_rate, waveform
+
+    def save(self, text: str, output: str | Path) -> Path:
+        destination = Path(output)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        sample_rate, waveform = self.synthesize(text)
+        
+        data_int16 = (waveform * 32767.0).clip(-32768, 32767).astype("<i2")
+        raw_bytes = data_int16.tobytes()
+        data_size = len(raw_bytes)
+        
+        header = bytearray()
+        header.extend(b"RIFF")
+        header.extend((36 + data_size).to_bytes(4, "little"))
+        header.extend(b"WAVE")
+        header.extend(b"fmt ")
+        header.extend((16).to_bytes(4, "little"))
+        header.extend((1).to_bytes(2, "little"))
+        header.extend((1).to_bytes(2, "little"))
+        header.extend((sample_rate).to_bytes(4, "little"))
+        header.extend((sample_rate * 2).to_bytes(4, "little"))
+        header.extend((2).to_bytes(2, "little"))
+        header.extend((16).to_bytes(2, "little"))
+        header.extend(b"data")
+        header.extend((data_size).to_bytes(4, "little"))
+        
+        with open(destination, "wb") as f:
+            f.write(header)
+            f.write(raw_bytes)
+            
+        return destination
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run PyTorch-free ONNX VTX-TTS speech synthesis.")
+    parser.add_argument("--model-dir", type=Path, default=PACKAGE_ROOT)
+    parser.add_argument("--text", required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+
+    engine = VtxTTSOnnx(args.model_dir)
+    engine.save(args.text, args.output)
+    print(f"Wrote {args.output} at {engine.sample_rate} Hz without PyTorch!")
+
+
+if __name__ == "__main__":
+    main()

onnx/model.onnx

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

onnx/model.onnx.data

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

train.py

@@ -0,0 +1,524 @@
+"""
+Fine-tune VTXAI/tts (VITS-architecture) on Hindi speech.
+
+This is a REAL training script:
+  - real audio -> mel-spectrogram target (via torchaudio, no torchcodec needed)
+  - real VITS forward pass (text encoder, posterior encoder, flow, duration predictor, decoder)
+  - real losses: mel L1, KL divergence, duration loss, adversarial (LSGAN) + feature matching
+  - MultiPeriodDiscriminator trained jointly (adversarial training is required for VITS to
+    produce intelligible/natural audio -- without it you just get a blurry mean waveform)
+  - proper batching/padding via a collate function, not batch-size-1 streaming
+  - checkpointing so you can resume 1000+ hours of training across sessions
+
+Usage:
+    Paste this entire file into one Jupyter/Colab cell and run it. Edit the
+    CONFIG block near the top (repo, dataset, batch size, etc.) instead of
+    passing CLI flags.
+
+Requires: torch, torchaudio, datasets, safetensors, huggingface_hub
+    pip install -q torch torchaudio datasets safetensors huggingface_hub
+Do NOT need torchcodec: we force decode=False and decode audio ourselves with torchaudio,
+which is more memory-stable for long streaming runs anyway.
+"""
+
+import os
+import io
+import json
+import math
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchaudio
+from torch.utils.data import IterableDataset, DataLoader
+import sys
+import importlib.util
+
+from huggingface_hub import snapshot_download
+
+# ==============================================================================
+# ⚙️ CONFIG — edit these values directly, then just run the cell
+# ==============================================================================
+HF_REPO        = "VTXAI/tts"            # base model repo to download+finetune
+BASE_MODEL_DIR = "./vtx_base_model"     # where the base snapshot gets cached
+DATASET_NAME   = "SPRINGLab/IndicTTS-Hindi"
+DATASET_SPLIT  = "train"
+OUT_DIR        = "./ckpt"               # finetune checkpoints saved here
+RESUME_PATH    = ""                     # set to a ./ckpt/*.pth path to resume, else ""
+
+BATCH_SIZE   = 16
+EPOCHS       = 1
+MAX_STEPS    = 0        # 0 = run full epochs, otherwise stop after this many new steps
+LR           = 2e-4
+SAVE_EVERY   = 1000
+LOG_EVERY    = 50
+NUM_WORKERS  = 2
+# ==============================================================================
+
+
+def load_modeling(repo_id="VTXAI/tts", local_dir="./vtx_base_model"):
+    """
+    Ensures modeling.py is available from local repo directory or downloaded snapshot
+    and imports required model classes/functions from it.
+    """
+    local_modeling = Path("modeling.py")
+    if not local_modeling.exists():
+        snapshot_dir = Path(snapshot_download(repo_id=repo_id, local_dir=local_dir))
+        local_modeling = snapshot_dir / "modeling.py"
+
+    spec = importlib.util.spec_from_file_location("vtx_modeling", local_modeling)
+    vtx_modeling = importlib.util.module_from_spec(spec)
+    sys.modules["vtx_modeling"] = vtx_modeling
+    spec.loader.exec_module(vtx_modeling)
+
+    return (
+        vtx_modeling.SynthesizerTrn,
+        vtx_modeling.MultiPeriodDiscriminator,
+        vtx_modeling.kl_divergence,
+        vtx_modeling.slice_segments,
+    )
+
+SynthesizerTrn, MultiPeriodDiscriminator, kl_divergence, slice_segments = load_modeling(HF_REPO, BASE_MODEL_DIR)
+def fetch_base_model(repo_id="VTXAI/tts", local_dir="./vtx_base_model"):
+    """
+    Downloads the full VTXAI/tts repo snapshot and locates config.json + weights
+    within it. Returns (config_path, weights_path).
+    """
+    local_dir = Path(local_dir)
+    print(f"Downloading base model from https://huggingface.co/{repo_id} ...")
+    repo_path = Path(snapshot_download(repo_id=repo_id, local_dir=str(local_dir)))
+    print(f"Downloaded to {repo_path}")
+
+    config_path = None
+    p = repo_path / "config.json"
+    if p.exists():
+        config_path = p
+    else:
+        found = list(repo_path.rglob("config.json"))
+        if not found:
+            raise FileNotFoundError(f"No config.json found in {repo_id} snapshot.")
+        config_path = found[0]
+
+    weights_path = None
+    for cand in ["model.safetensors", "G.safetensors", "model.pth", "G.pth"]:
+        p = repo_path / cand
+        if p.exists():
+            weights_path = p
+            break
+    if weights_path is None:
+        for pattern in ["*.safetensors", "*.pth"]:
+            found = list(repo_path.rglob(pattern))
+            if found:
+                weights_path = found[0]
+                break
+    if weights_path is None:
+        raise FileNotFoundError(f"No weights (.safetensors/.pth) found in {repo_id} snapshot.")
+
+    print(f"Using config: {config_path}")
+    print(f"Using weights: {weights_path}")
+    return config_path, weights_path
+
+# ----------------------------------------------------------------------------
+# 0. maximum_path fallback
+# ----------------------------------------------------------------------------
+# The reference VITS repo ships a compiled `monotonic_align` extension exposing
+# maximum_path(neg_cent, mask). modeling.py calls it but doesn't define it, so we
+# provide a numba-free, pure-torch fallback here (slower, but correct and dependency-light).
+def maximum_path(neg_cent, mask):
+    """
+    neg_cent: [b, t_t, t_s]  mask: [b, t_t, t_s] -> returns path [b, t_t, t_s]
+    Pure-python dynamic programming, batched via a simple loop (fine at these batch sizes).
+    """
+    device = neg_cent.device
+    dtype = neg_cent.dtype
+    neg_cent = neg_cent.detach().cpu().numpy().astype(np.float64)
+    mask = mask.detach().cpu().numpy().astype(np.bool_)
+    b, t_t, t_s = neg_cent.shape
+    path = np.zeros_like(neg_cent, dtype=np.float64)
+
+    for i in range(b):
+        t_y = int(mask[i, :, 0].sum())
+        t_x = int(mask[i, 0, :].sum())
+        if t_y == 0 or t_x == 0:
+            continue
+        value = neg_cent[i, :t_y, :t_x]
+        v = np.full((t_y, t_x), -1e9, dtype=np.float64)
+        v[0, 0] = value[0, 0]
+        for y in range(1, t_y):
+            v[y, 0] = v[y - 1, 0] + value[y, 0]
+        for x in range(1, t_x):
+            v[0, x] = -1e9  # can't start later in x at t_y=0
+        for y in range(1, t_y):
+            for x in range(1, min(t_x, y + 1)):
+                v[y, x] = value[y, x] + max(v[y - 1, x], v[y - 1, x - 1])
+        # backtrack
+        index = t_x - 1
+        for y in range(t_y - 1, -1, -1):
+            path[i, y, index] = 1.0
+            if index > 0 and (y == 0 or v[y - 1, index - 1] > v[y - 1, index]):
+                index -= 1
+    return torch.from_numpy(path).to(device=device, dtype=dtype)
+
+
+# ----------------------------------------------------------------------------
+# 1. Text symbols / tokenizer (Hindi, matches your existing pipeline)
+# ----------------------------------------------------------------------------
+_pad = "_"
+_punctuation = ';:,.!?¡¿—…"«»“” '
+_hindi_letters = "अआइईउऊऋएऐओऔकखगघङचछजझञटठडढणतथदधनपफबभमयरलवशषसहक्षत्रज्ञ़ािीुूृेैोौॉ्िंः"
+_letters_ipa = "ɑɐɒæɓɓβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞"
+
+hindi_symbols = [_pad] + list(_punctuation) + list(_hindi_letters) + list(_letters_ipa)
+hindi_symbol_to_id = {s: i for i, s in enumerate(hindi_symbols)}
+
+
+def hindi_text_to_sequence(text: str) -> list:
+    seq = []
+    for ch in text:
+        if ch in hindi_symbol_to_id:
+            seq.append(hindi_symbol_to_id[ch])
+        elif ch == " ":
+            seq.append(1)
+    if not seq:
+        seq = [1, 2, 3]
+    interspersed = [0] * (len(seq) * 2 + 1)
+    interspersed[1::2] = seq
+    return interspersed
+
+
+# ----------------------------------------------------------------------------
+# 2. Mel-spectrogram (matches VITS training target: linear spec for enc_q,
+#    mel spec for the reconstruction loss)
+# ----------------------------------------------------------------------------
+class AudioProcessor:
+    def __init__(self, cfg):
+        d = cfg["data"]
+        self.sampling_rate = d["sampling_rate"]
+        self.filter_length = d["filter_length"]
+        self.hop_length = d["hop_length"]
+        self.win_length = d["win_length"]
+        self.n_mel_channels = d.get("n_mel_channels", 80)
+        self.mel_fmin = d.get("mel_fmin", 0.0)
+        self.mel_fmax = d.get("mel_fmax", None)
+
+        self.spec_transform = torchaudio.transforms.Spectrogram(
+            n_fft=self.filter_length,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            power=1.0,
+            center=False,
+            pad=(self.filter_length - self.hop_length) // 2,
+        )
+        self.mel_scale = torchaudio.transforms.MelScale(
+            n_mels=self.n_mel_channels,
+            sample_rate=self.sampling_rate,
+            f_min=self.mel_fmin,
+            f_max=self.mel_fmax,
+            n_stft=self.filter_length // 2 + 1,
+        )
+
+    def to(self, device):
+        self.spec_transform = self.spec_transform.to(device)
+        self.mel_scale = self.mel_scale.to(device)
+        return self
+
+    def wav_to_spec(self, wav):
+        # wav: [1, T] float32 in [-1, 1] -> spec [C, F]
+        spec = self.spec_transform(wav)
+        return spec.squeeze(0)
+
+    def spec_to_mel(self, spec):
+        if spec.device != self.mel_scale.fb.device:
+            self.mel_scale = self.mel_scale.to(spec.device)
+        return self.mel_scale(spec.unsqueeze(0)).squeeze(0)
+
+    def mel_scale_transform(self, spec):
+        if spec.device != self.mel_scale.fb.device:
+            self.mel_scale = self.mel_scale.to(spec.device)
+        return self.mel_scale(spec)
+
+    def wav_to_mel(self, wav):
+        return self.spec_to_mel(self.wav_to_spec(wav))
+
+
+# ----------------------------------------------------------------------------
+# 3. Streaming dataset over the HF Hindi corpus, with real audio decoding
+# ----------------------------------------------------------------------------
+class HindiTTSIterable(IterableDataset):
+    """
+    Wraps datasets.load_dataset(..., streaming=True). We force decode=False on the
+    audio column and decode manually with torchaudio -- this avoids the torchcodec
+    dependency entirely and is more robust for very long (1000+ hr) streaming runs.
+    """
+
+    def __init__(self, dataset_name, split, ap: AudioProcessor, max_audio_sec=12.0, text_key=None):
+        super().__init__()
+        self.dataset_name = dataset_name
+        self.split = split
+        self.ap = ap
+        self.max_samples = int(max_audio_sec * ap.sampling_rate)
+        self.text_key = text_key
+
+    def _load_stream(self):
+        from datasets import load_dataset, Audio
+        ds = load_dataset(self.dataset_name, split=self.split, streaming=True)
+        ds = ds.cast_column("audio", Audio(decode=False))
+        return ds
+
+    def __iter__(self):
+        ds = self._load_stream()
+        for item in ds:
+            try:
+                audio_field = item["audio"]
+                raw_bytes = audio_field["bytes"] if isinstance(audio_field, dict) else None
+                if raw_bytes is not None:
+                    wav, sr = torchaudio.load(io.BytesIO(raw_bytes))
+                elif isinstance(audio_field, dict) and audio_field.get("path"):
+                    wav, sr = torchaudio.load(audio_field["path"])
+                else:
+                    continue
+
+                if wav.shape[0] > 1:
+                    wav = wav.mean(0, keepdim=True)
+                if sr != self.ap.sampling_rate:
+                    wav = torchaudio.functional.resample(wav, sr, self.ap.sampling_rate)
+                if wav.shape[1] > self.max_samples:
+                    continue  # skip overly long clips for stable batch memory
+                if wav.shape[1] < self.ap.hop_length * 8:
+                    continue
+
+                text_key = self.text_key
+                if text_key is None:
+                    text_key = "normalized_text" if "normalized_text" in item else "text"
+                text = item.get(text_key, "")
+                if not text or not text.strip():
+                    continue
+
+                tok = torch.LongTensor(hindi_text_to_sequence(text))
+                yield tok, wav.squeeze(0)
+            except Exception:
+                # skip malformed rows rather than crashing a 1000hr run
+                continue
+
+
+def collate_fn(batch):
+    batch = [b for b in batch if b[0].numel() > 0 and b[1].numel() > 0]
+    batch.sort(key=lambda x: x[1].shape[0], reverse=True)
+    text_lens = torch.LongTensor([b[0].shape[0] for b in batch])
+    wav_lens = torch.LongTensor([b[1].shape[0] for b in batch])
+
+    text_pad = torch.zeros(len(batch), text_lens.max(), dtype=torch.long)
+    wav_pad = torch.zeros(len(batch), wav_lens.max(), dtype=torch.float32)
+
+    for i, (tok, wav) in enumerate(batch):
+        text_pad[i, : tok.shape[0]] = tok
+        wav_pad[i, : wav.shape[0]] = wav
+
+    return text_pad, text_lens, wav_pad, wav_lens
+
+
+# ----------------------------------------------------------------------------
+# 4. Loss functions
+# ----------------------------------------------------------------------------
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = torch.mean((1 - dr) ** 2)
+        g_loss = torch.mean(dg ** 2)
+        loss += r_loss + g_loss
+    return loss
+
+
+def generator_adv_loss(disc_outputs):
+    loss = 0
+    for dg in disc_outputs:
+        loss += torch.mean((1 - dg) ** 2)
+    return loss
+
+
+def feature_loss(fmap_r, fmap_g):
+    loss = 0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += torch.mean(torch.abs(rl.float() - gl.float()))
+    return loss * 2
+
+
+# ----------------------------------------------------------------------------
+# 5. Checkpoint loading (keeps your existing "transfer weights, resize vocab" fix)
+# ----------------------------------------------------------------------------
+def load_pretrained_generator(model, ckpt_path):
+    ckpt_path = Path(ckpt_path)
+    if ckpt_path.suffix == ".safetensors":
+        from safetensors.torch import load_file
+        state_dict = load_file(str(ckpt_path))
+    else:
+        raw = torch.load(ckpt_path, map_location="cpu")
+        state_dict = raw.get("model", raw)
+
+    # Drop embedding weights: symbol vocab size for Hindi differs from source checkpoint
+    filtered = {
+        k: v for k, v in state_dict.items()
+        if not (k.endswith("emb.weight") or k.startswith("enc_p.emb"))
+    }
+    model_dict = model.state_dict()
+    compatible = {k: v for k, v in filtered.items() if k in model_dict and model_dict[k].shape == v.shape}
+    skipped = [k for k in filtered if k not in compatible]
+    model_dict.update(compatible)
+    model.load_state_dict(model_dict, strict=False)
+    print(f"Loaded {len(compatible)}/{len(state_dict)} tensors from checkpoint "
+          f"(skipped {len(skipped)}, incl. resized text embedding).")
+
+
+# ----------------------------------------------------------------------------
+# 6. Training run (executes directly — paste this whole file into one notebook cell)
+# ----------------------------------------------------------------------------
+def main():
+    out_dir = Path(OUT_DIR)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    config_path, checkpoint_path = fetch_base_model(HF_REPO, BASE_MODEL_DIR)
+
+    with open(config_path) as f:
+        config = json.load(f)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    ap = AudioProcessor(config).to(device)
+
+    net_g = SynthesizerTrn(
+        len(hindi_symbols),
+        config["data"]["filter_length"] // 2 + 1,
+        config["train"]["segment_size"] // config["data"]["hop_length"],
+        **config["model"],
+    ).to(device)
+    net_d = MultiPeriodDiscriminator(config["model"].get("use_spectral_norm", False)).to(device)
+
+    start_step = 0
+    if RESUME_PATH and Path(RESUME_PATH).exists():
+        ckpt = torch.load(RESUME_PATH, map_location="cpu")
+        net_g.load_state_dict(ckpt["net_g"])
+        net_d.load_state_dict(ckpt["net_d"])
+        start_step = ckpt.get("step", 0)
+        print(f"Resumed from {RESUME_PATH} at step {start_step}")
+    else:
+        load_pretrained_generator(net_g, checkpoint_path)
+
+    optim_g = torch.optim.AdamW(net_g.parameters(), lr=LR, betas=(0.8, 0.99), eps=1e-9)
+    optim_d = torch.optim.AdamW(net_d.parameters(), lr=LR, betas=(0.8, 0.99), eps=1e-9)
+    if RESUME_PATH and Path(RESUME_PATH).exists():
+        optim_g.load_state_dict(ckpt["optim_g"])
+        optim_d.load_state_dict(ckpt["optim_d"])
+
+    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=0.999875)
+    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=0.999875)
+
+    dataset = HindiTTSIterable(DATASET_NAME, DATASET_SPLIT, ap)
+    loader = DataLoader(
+        dataset, batch_size=BATCH_SIZE, collate_fn=collate_fn,
+        num_workers=NUM_WORKERS, pin_memory=(device.type == "cuda"),
+    )
+
+    hop = config["data"]["hop_length"]
+    seg_size = config["train"]["segment_size"]
+    seg_frames = seg_size // hop
+    c_mel = config["train"].get("c_mel", 45.0)
+    c_kl = config["train"].get("c_kl", 1.0)
+
+    step = start_step
+    t0 = time.time()
+    net_g.train()
+    net_d.train()
+
+    for epoch in range(EPOCHS):
+        for text_pad, text_lens, wav_pad, wav_lens in loader:
+            text_pad, text_lens = text_pad.to(device), text_lens.to(device)
+            wav_pad, wav_lens = wav_pad.to(device), wav_lens.to(device)
+
+            # Build linear spectrogram targets per-sample (variable length), then pad
+            specs, spec_lens = [], []
+            for i in range(wav_pad.shape[0]):
+                w = wav_pad[i, : wav_lens[i]].unsqueeze(0).cpu()
+                s = ap.wav_to_spec(w)
+                specs.append(s)
+                spec_lens.append(s.shape[-1])
+            spec_lens = torch.LongTensor(spec_lens).to(device)
+            spec_pad = torch.zeros(len(specs), specs[0].shape[0], spec_lens.max(), device=device)
+            for i, s in enumerate(specs):
+                spec_pad[i, :, : s.shape[-1]] = s.to(device)
+
+            # ---- Generator forward ----
+            (y_hat, l_length, attn, ids_slice, x_mask, y_mask,
+             (z, z_p, m_p, logs_p, m_q, logs_q)) = net_g(text_pad, text_lens, spec_pad, spec_lens)
+
+            mel_pad = ap.mel_scale_transform(spec_pad)
+            mel_slice = slice_segments(mel_pad, ids_slice, seg_frames)
+            with torch.no_grad():
+                wav_hat_mono = y_hat.squeeze(1)
+            y_hat_mel = torch.stack([
+                ap.spec_to_mel(ap.wav_to_spec(wav_hat_mono[i:i + 1].detach().cpu())).to(device)
+                for i in range(wav_hat_mono.shape[0])
+            ])
+            y_slice = slice_segments(wav_pad.unsqueeze(1), ids_slice * hop, seg_size)
+
+            # ---- Discriminator step ----
+            y_d_hat_r, y_d_hat_g, _, _ = net_d(y_slice, y_hat.detach())
+            loss_disc = discriminator_loss(y_d_hat_r, y_d_hat_g)
+            optim_d.zero_grad()
+            loss_disc.backward()
+            torch.nn.utils.clip_grad_norm_(net_d.parameters(), 5.0)
+            optim_d.step()
+
+            # ---- Generator step ----
+            y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y_slice, y_hat)
+            loss_mel = F.l1_loss(mel_slice, y_hat_mel) * c_mel
+            loss_kl = torch.mean(kl_divergence(m_p, logs_p, m_q, logs_q)) * c_kl
+            loss_dur = torch.sum(l_length.float())
+            loss_fm = feature_loss(fmap_r, fmap_g)
+            loss_adv = generator_adv_loss(y_d_hat_g)
+            loss_gen_total = loss_adv + loss_fm + loss_mel + loss_kl + loss_dur
+
+            optim_g.zero_grad()
+            loss_gen_total.backward()
+            torch.nn.utils.clip_grad_norm_(net_g.parameters(), 5.0)
+            optim_g.step()
+
+            step += 1
+            if step % LOG_EVERY == 0:
+                elapsed = time.time() - t0
+                print(f"[step {step}] gen={loss_gen_total.item():.4f} "
+                      f"mel={loss_mel.item():.4f} kl={loss_kl.item():.4f} "
+                      f"dur={loss_dur.item():.4f} adv={loss_adv.item():.4f} "
+                      f"fm={loss_fm.item():.4f} disc={loss_disc.item():.4f} "
+                      f"({elapsed:.1f}s elapsed)")
+
+            if step % SAVE_EVERY == 0:
+                ckpt_path = out_dir / f"vtx_hindi_step{step}.pth"
+                torch.save({
+                    "net_g": net_g.state_dict(),
+                    "net_d": net_d.state_dict(),
+                    "optim_g": optim_g.state_dict(),
+                    "optim_d": optim_d.state_dict(),
+                    "step": step,
+                    "symbols": hindi_symbols,
+                }, ckpt_path)
+                print(f"Saved checkpoint: {ckpt_path}")
+
+            if MAX_STEPS and step >= start_step + MAX_STEPS:
+                break
+        scheduler_g.step()
+        scheduler_d.step()
+        if MAX_STEPS and step >= start_step + MAX_STEPS:
+            break
+
+    final_path = out_dir / "vtx_hindi_final.pth"
+    torch.save({"net_g": net_g.state_dict(), "net_d": net_d.state_dict(), "step": step,
+                "symbols": hindi_symbols}, final_path)
+    print(f"Done. Final checkpoint: {final_path}")
+
+
+if __name__ == "__main__":
+    main()