MyanmarTTS 1.0.0__py3-none-any.whl

myanmar_tts/__init__.py

@@ -0,0 +1,42 @@
+"""MyanmarTTS - Burmese text-to-speech, from scratch, CC0."""
+import os
+from huggingface_hub import hf_hub_download
+
+HF_REPO = "freococo/MyanmarTTS"
+MODEL_FILE = "model_fp16.safetensors"
+VOCODER_FILE = "vocos.pt"
+REF_AUDIO = "samples/sample_0.wav"
+
+__version__ = "1.0.0"
+
+
+class MyanmarTTS:
+    def __init__(self, device="cuda"):
+        self.device = device
+        self._model = None
+        self._ref = None
+
+    def _load(self):
+        if self._model is not None:
+            return
+        print(f"Downloading model files from {HF_REPO}...")
+        model_path = hf_hub_download(repo_id=HF_REPO, filename=MODEL_FILE)
+        vocos_path = hf_hub_download(repo_id=HF_REPO, filename=VOCODER_FILE)
+        self._ref = hf_hub_download(repo_id=HF_REPO, filename=REF_AUDIO)
+
+        from .api import StableTTSAPI
+        print("Loading model...")
+        self._model = StableTTSAPI(model_path, vocos_path, "vocos")
+        self._model.to(self.device)
+        print("Ready.")
+
+    def tts(self, text):
+        """Return a numpy float32 waveform at 44100 Hz."""
+        self._load()
+        audio, _ = self._model.inference(
+            text, self._ref, "burmese", step=32, solver="dopri5", cfg=3.0
+        )
+        return audio.squeeze(0).cpu().numpy()
+
+
+__all__ = ["MyanmarTTS"]

myanmar_tts/api.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+from dataclasses import asdict
+
+from .utils.audio import LogMelSpectrogram
+from .config import ModelConfig, MelConfig
+from .models.model import StableTTS
+
+from .text import symbols
+from .text import cleaned_text_to_sequence
+from .text.burmese import burmese_to_ipa2
+from .datas.dataset import intersperse
+from .utils.audio import load_and_resample_audio
+
+
+def get_vocoder(model_path, model_name='vocos'):
+    if model_name == 'vocos':
+        from .vocoders.vocos.models.model import Vocos
+        from .config import VocosConfig, MelConfig
+        vocoder = Vocos(VocosConfig(), MelConfig())
+        vocoder.load_state_dict(torch.load(model_path, weights_only=True, map_location='cpu'))
+        vocoder.eval()
+    else:
+        raise NotImplementedError(f"Unsupported vocoder: {model_name}")
+    return vocoder
+
+
+class StableTTSAPI(nn.Module):
+    def __init__(self, tts_model_path, vocoder_model_path, vocoder_name='vocos'):
+        super().__init__()
+        self.mel_config = MelConfig()
+        self.tts_model_config = ModelConfig()
+
+        self.mel_extractor = LogMelSpectrogram(**asdict(self.mel_config))
+
+        self.tts_model = StableTTS(len(symbols), self.mel_config.n_mels, **asdict(self.tts_model_config))
+
+        if tts_model_path.endswith(".safetensors"):
+            from safetensors.torch import load_file
+            state = load_file(tts_model_path)
+        else:
+            state = torch.load(tts_model_path, map_location='cpu', weights_only=True)
+        self.tts_model.load_state_dict(state)
+        self.tts_model.eval()
+
+        self.vocoder_model = get_vocoder(vocoder_model_path, vocoder_name)
+        self.vocoder_model.eval()
+
+        self.g2p_mapping = {
+            'burmese': burmese_to_ipa2,
+        }
+        self.supported_languages = self.g2p_mapping.keys()
+
+    @torch.inference_mode()
+    def inference(self, text, ref_audio, language, step, temperature=1.0,
+                  length_scale=1.0, solver=None, cfg=3.0):
+        device = next(self.parameters()).device
+        phonemizer = self.g2p_mapping.get(language)
+        if phonemizer is None:
+            raise ValueError(f"Unsupported language: {language}")
+
+        text = phonemizer(text)
+        text = torch.tensor(
+            intersperse(cleaned_text_to_sequence(text), item=0),
+            dtype=torch.long, device=device
+        ).unsqueeze(0)
+        text_length = torch.tensor([text.size(-1)], dtype=torch.long, device=device)
+
+        ref_audio = load_and_resample_audio(ref_audio, self.mel_config.sample_rate).to(device)
+        ref_audio = self.mel_extractor(ref_audio)
+
+        mel_output = self.tts_model.synthesise(
+            text, text_length, step, temperature, ref_audio,
+            length_scale, solver, cfg
+        )['decoder_outputs']
+        audio_output = self.vocoder_model(mel_output)
+        return audio_output.cpu(), mel_output.cpu()
+
+    def get_params(self):
+        tts_param = sum(p.numel() for p in self.tts_model.parameters()) / 1e6
+        vocoder_param = sum(p.numel() for p in self.vocoder_model.parameters()) / 1e6
+        return tts_param, vocoder_param

myanmar_tts/burmese.py

@@ -0,0 +1,42 @@
+"""
+Character-level text frontend for Burmese.
+No real G2P: we clean the text and pass characters through directly.
+Matches the interface of english_to_ipa2 / japanese_to_ipa2 (returns a list).
+"""
+import re
+
+# Burmese Unicode ranges
+_BURMESE_RANGES = (
+    (0x1000, 0x109F),   # Myanmar
+    (0xAA60, 0xAA7F),   # Myanmar Extended-A
+    (0xA9E0, 0xA9FF),   # Myanmar Extended-B
+)
+
+# Allow these through the cleaner; everything else gets dropped
+_ALLOWED_ASCII_PUNCT = set(".,!?'\"-:;()")
+_ALLOWED_BURMESE_PUNCT = set("\u104A\u104B\u104C\u104D\u104E\u104F")
+
+
+def _is_burmese(ch):
+    o = ord(ch)
+    return any(lo <= o <= hi for lo, hi in _BURMESE_RANGES)
+
+
+def clean_burmese(text):
+    """Keep only Burmese chars, spaces, and whitelisted punctuation."""
+    out = []
+    for ch in text:
+        if ch in (" ", "\t", "\n", "\r"):
+            out.append(" ")
+        elif _is_burmese(ch):
+            out.append(ch)
+        elif ch in _ALLOWED_ASCII_PUNCT or ch in _ALLOWED_BURMESE_PUNCT:
+            out.append(ch)
+        # else: drop silently
+    return re.sub(r"\s+", " ", "".join(out)).strip()
+
+
+def burmese_to_ipa2(text):
+    """Return a list of character tokens (matches the IPA-G2P interface)."""
+    cleaned = clean_burmese(text)
+    return list(cleaned)

myanmar_tts/config.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+
+@dataclass
+class MelConfig:
+    sample_rate: int = 44100
+    n_fft: int = 2048
+    win_length: int = 2048
+    hop_length: int = 512
+    f_min: float = 0.0
+    f_max: float = None
+    pad: int = 0
+    n_mels: int = 128
+    center: bool = False
+    pad_mode: str = "reflect"
+    mel_scale: str = "slaney"
+    
+    def __post_init__(self):
+        if self.pad == 0:
+            self.pad = (self.n_fft - self.hop_length) // 2
+            
+@dataclass
+class ModelConfig:
+    hidden_channels: int = 256
+    filter_channels: int = 1024
+    n_heads: int = 4
+    n_enc_layers: int = 3 
+    n_dec_layers: int = 6 
+    kernel_size: int = 3
+    p_dropout: int = 0.1
+    gin_channels: int = 256
+            
+@dataclass
+class TrainConfig:
+    train_dataset_path: str = 'filelists/filelist.json'
+    test_dataset_path: str = 'filelists/filelist.json' # not used
+    batch_size: int = 32
+    learning_rate: float = 1e-4
+    num_epochs: int = 50
+    model_save_path: str = './checkpoints'
+    log_dir: str = './runs'
+    log_interval: int = 5
+    save_interval: int = 1
+    warmup_steps: int = 200
+    
+@dataclass
+class VocosConfig:
+    input_channels: int = 128
+    dim: int = 512
+    intermediate_dim: int = 1536
+    num_layers: int = 8

myanmar_tts/datas/all_sources_dataset.py

@@ -0,0 +1,160 @@
+import os, random, io, glob, tarfile
+import pyarrow.parquet as pq
+import soundfile as sf
+import torch
+import torchaudio
+from torch.utils.data import IterableDataset, get_worker_info
+
+from ..config import MelConfig
+from ..text import cleaned_text_to_sequence
+from ..text.burmese import burmese_to_ipa2
+
+
+def _intersperse(lst, item):
+    r = [item] * (len(lst) * 2 + 1)
+    r[1::2] = lst
+    return r
+
+
+SCRATCH = "/mnt/local-scratch"
+
+SOURCES = [
+    ("parquet", "/content/shards/data/run*/shard-*.parquet",          "clean_text",     1.0),
+    ("parquet", f"{SCRATCH}/data_all/95k/data/run*/shard-*.parquet",  "original_text",  1.0),
+    ("parquet", f"{SCRATCH}/data_all/bible/data/train-*.parquet",     "transcript_char",1.0),
+    ("parquet", f"{SCRATCH}/data_all/6k/data/shard-*.parquet",        "original_text",  1.0),
+    ("parquet", f"{SCRATCH}/data_all/qtext/data/train-*.parquet",     "original_field", 0.5),
+    ("parquet", f"{SCRATCH}/data_all/synth_mix/data/train-*.parquet", "meta_words",     0.5),
+    ("parquet", f"{SCRATCH}/data_all/synth_6pr/data/train-*.parquet", "meta_words",     0.5),
+    ("parquet", f"{SCRATCH}/data_all/synth_6p/data/train-*.parquet",  "meta_words",     0.5),
+    ("parquet", f"{SCRATCH}/data_all/synth_vox/data/train-*.parquet", "myanmar_text",   0.5),
+    ("tar",     f"{SCRATCH}/data_all/1m/audio_tars/*.tar",            "csv",            1.0),
+]
+
+
+class AllSourcesDataset(IterableDataset):
+    def __init__(self, sources=SOURCES, target_sr=None, min_sec=0.5, max_sec=11.0):
+        self.target_sr = target_sr or MelConfig().sample_rate
+        self.min_sec = min_sec
+        self.max_sec = max_sec
+
+        self.entries = []
+        for kind, pattern, field, weight in sources:
+            paths = sorted(glob.glob(pattern))
+            for p in paths:
+                self.entries.append((kind, p, field, weight))
+        print(f"[AllSources] {len(self.entries)} shards across {len(sources)} sources | target_sr={self.target_sr}")
+
+        self.tar_meta = {}
+        csv_path = f"{SCRATCH}/data_all/1m/train_1m_combined_24k.csv"
+        if os.path.exists(csv_path):
+            print("[AllSources] parsing CSV...")
+            with open(csv_path, "r", encoding="utf-8") as f:
+                next(f, None)
+                for line in f:
+                    parts = line.rstrip("\n").split("|", 1)
+                    if len(parts) == 2:
+                        fname = os.path.basename(parts[0])
+                        self.tar_meta[fname] = parts[1].strip()
+            print(f"[AllSources] indexed {len(self.tar_meta)} CSV entries")
+
+    def _get_text(self, t):
+        if t is None:
+            return ""
+        if isinstance(t, list):
+            return " ".join(str(x) for x in t if x).strip()
+        return str(t).strip()
+
+    def _to_wav_phone(self, wav_bytes, text, given_sr=None):
+        try:
+            if not text:
+                return None
+            data, sr = sf.read(io.BytesIO(wav_bytes), dtype="float32")
+            if data.size == 0:
+                return None
+            wav = torch.from_numpy(data)
+            if wav.dim() > 1:
+                wav = wav.mean(dim=-1)
+            if wav.abs().max().item() < 0.001:
+                return None
+
+            actual_sr = given_sr or sr
+            if actual_sr != self.target_sr:
+                wav = torchaudio.functional.resample(wav, actual_sr, self.target_sr)
+
+            dur = wav.numel() / self.target_sr
+            if dur < self.min_sec or dur > self.max_sec:
+                return None
+
+            phones = burmese_to_ipa2(text)
+            ids = cleaned_text_to_sequence(phones)
+            if len(ids) < 3:
+                return None
+            phone = torch.tensor(_intersperse(ids, 0), dtype=torch.long)
+            return wav, phone
+        except Exception:
+            return None
+
+    def _iter_parquet(self, path, field):
+        try:
+            pf = pq.ParquetFile(path)
+            for batch in pf.iter_batches(batch_size=64):
+                d = batch.to_pydict()
+                for i in range(len(batch)):
+                    text = self._get_text(d[field][i])
+                    if not text:
+                        continue
+                    a = d["audio"][i]
+                    if not isinstance(a, dict):
+                        continue
+                    if a.get("bytes") is not None:
+                        out = self._to_wav_phone(a["bytes"], text)
+                    elif a.get("array") is not None:
+                        arr = a["array"]
+                        sr = a["sampling_rate"]
+                        bio = io.BytesIO()
+                        sf.write(bio, arr, sr, format="WAV")
+                        out = self._to_wav_phone(bio.getvalue(), text, given_sr=sr)
+                    else:
+                        continue
+                    if out is not None:
+                        yield out
+        except Exception as e:
+            print(f"[skip parquet] {path}: {e}")
+            return
+
+    def _iter_tar(self, path):
+        try:
+            with tarfile.open(path, "r:") as tf:  # seekable mode (not stream)
+                for member in tf:
+                    if not member.name.endswith(".wav"):
+                        continue
+                    fname = os.path.basename(member.name)
+                    text = self.tar_meta.get(fname)
+                    if not text:
+                        continue
+                    f = tf.extractfile(member)
+                    if f is None:
+                        continue
+                    out = self._to_wav_phone(f.read(), text)
+                    if out is not None:
+                        yield out
+        except Exception as e:
+            print(f"[skip tar] {path}: {e}")
+            return
+
+    def __iter__(self):
+        info = get_worker_info()
+        wid = info.id if info else 0
+        nw = info.num_workers if info else 1
+        my = self.entries[wid::nw]
+        rng = random.Random(42 + wid)
+        while True:
+            rng.shuffle(my)
+            for kind, path, field, weight in my:
+                if weight < 1.0 and rng.random() > weight:
+                    continue
+                if kind == "parquet":
+                    yield from self._iter_parquet(path, field)
+                else:
+                    yield from self._iter_tar(path)

myanmar_tts/datas/collate_wav.py

@@ -0,0 +1,39 @@
+import random
+import torch
+
+
+def _random_slice(x, mn=12, mx=3):
+    L = x.size(-1)
+    if L < 12:
+        return x
+    seg = random.randint(max(1, L // mn), max(1, L // mx))
+    seg = min(seg, L)
+    start = random.randint(0, L - seg)
+    return x[start:start + seg]
+
+
+def collate_fn_wav(batch):
+    wavs, phones = zip(*batch)
+    wav_lens = torch.tensor([w.size(-1) for w in wavs], dtype=torch.long)
+    phone_lens = torch.tensor([p.size(-1) for p in phones], dtype=torch.long)
+
+    sliced = [_random_slice(w) for w in wavs]
+    sliced_lens = torch.tensor([s.size(-1) for s in sliced], dtype=torch.long)
+
+    B = len(wavs)
+    max_w = int(wav_lens.max())
+    max_p = int(phone_lens.max())
+    max_s = int(sliced_lens.max())
+
+    wav_pad = torch.zeros(B, max_w)
+    phone_pad = torch.zeros(B, max_p, dtype=torch.long)
+    sliced_pad = torch.zeros(B, max_s)
+
+    for i, w in enumerate(wavs):
+        wav_pad[i, :w.size(-1)] = w
+    for i, p in enumerate(phones):
+        phone_pad[i, :p.size(-1)] = p
+    for i, s in enumerate(sliced):
+        sliced_pad[i, :s.size(-1)] = s
+
+    return (phone_pad, phone_lens, wav_pad, wav_lens, sliced_pad, sliced_lens)

myanmar_tts/datas/dataset.py

@@ -0,0 +1,69 @@
+import os
+import random
+
+import json
+import torch
+from torch.utils.data import Dataset
+
+from ..text import cleaned_text_to_sequence
+
+def intersperse(lst: list, item: int):
+    """
+    putting a blank token between any two input tokens to improve pronunciation
+    see https://github.com/jaywalnut310/glow-tts/issues/43 for more details
+    """
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+    
+class StableDataset(Dataset):
+    def __init__(self, filelist_path, hop_length):
+        self.filelist_path = filelist_path     
+        self.hop_length = hop_length  
+        
+        self._load_filelist(filelist_path)
+
+    def _load_filelist(self, filelist_path):
+        filelist, lengths = [], []
+        with open(filelist_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = json.loads(line.strip())
+                filelist.append((line['mel_path'], line['phone']))
+                lengths.append(line['mel_length'])
+            
+        self.filelist = filelist
+        self.lengths = lengths # length is used for DistributedBucketSampler
+    
+    def __len__(self):
+        return len(self.filelist)
+
+    def __getitem__(self, idx):
+        mel_path, phone = self.filelist[idx]
+        mel = torch.load(mel_path, map_location='cpu', weights_only=True)
+        phone = torch.tensor(intersperse(cleaned_text_to_sequence(phone), 0), dtype=torch.long)
+        return mel, phone
+    
+def collate_fn(batch):
+    texts = [item[1] for item in batch]
+    mels = [item[0] for item in batch]
+    mels_sliced = [random_slice_tensor(mel) for mel in mels]
+    
+    text_lengths = torch.tensor([text.size(-1) for text in texts], dtype=torch.long)
+    mel_lengths = torch.tensor([mel.size(-1) for mel in mels], dtype=torch.long)
+    mels_sliced_lengths = torch.tensor([mel_sliced.size(-1) for mel_sliced in mels_sliced], dtype=torch.long)
+    
+    # pad to the same length
+    texts_padded = torch.nested.to_padded_tensor(torch.nested.nested_tensor(texts), padding=0)
+    mels_padded = torch.nested.to_padded_tensor(torch.nested.nested_tensor(mels), padding=0)
+    mels_sliced_padded = torch.nested.to_padded_tensor(torch.nested.nested_tensor(mels_sliced), padding=0)
+
+    return texts_padded, text_lengths, mels_padded, mel_lengths, mels_sliced_padded, mels_sliced_lengths
+
+# random slice mel for reference encoder to prevent overfitting
+def random_slice_tensor(x: torch.Tensor):
+    length = x.size(-1)
+    if length < 12:
+        return x 
+    segmnt_size = random.randint(length // 12, length // 3)
+    start = random.randint(0, length - segmnt_size)
+    return x[..., start : start + segmnt_size]

myanmar_tts/datas/local_shards_dataset.py

@@ -0,0 +1,73 @@
+import os, random, io
+import soundfile as sf
+import pyarrow.parquet as pq
+import torch
+import torchaudio
+from torch.utils.data import IterableDataset, get_worker_info
+from ..text import cleaned_text_to_sequence
+from ..text.burmese import burmese_to_ipa2
+
+def _intersperse(lst, item):
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+
+class LocalShardsDataset(IterableDataset):
+    def __init__(self, shard_paths, target_sr=44100, min_sec=0.6, max_sec=18.0):
+        self.shard_paths = list(shard_paths)
+        self.target_sr = target_sr
+        self.min_sec = min_sec
+        self.max_sec = max_sec
+
+    def _process_sample(self, row):
+        try:
+            text = (row.get("clean_text") or row.get("original_text") or "").strip()
+            if not text:
+                return None
+            a = row["audio"]
+            if isinstance(a, dict) and a.get("bytes") is not None:
+                data, sr = sf.read(io.BytesIO(a["bytes"]), dtype="float32")
+            elif isinstance(a, dict) and a.get("array") is not None:
+                import numpy as np
+                data = np.asarray(a["array"], dtype="float32")
+                sr = a["sampling_rate"]
+            else:
+                return None
+            wav = torch.from_numpy(data)
+            if wav.dim() > 1:
+                wav = wav.mean(dim=-1)
+            if sr != self.target_sr:
+                wav = torchaudio.functional.resample(wav, sr, self.target_sr)
+            dur = wav.numel() / self.target_sr
+            if dur < self.min_sec or dur > self.max_sec:
+                return None
+            phones = burmese_to_ipa2(text)
+            ids = cleaned_text_to_sequence(phones)
+            if len(ids) < 3:
+                return None
+            phone = torch.tensor(_intersperse(ids, 0), dtype=torch.long)
+            return wav, phone
+        except Exception:
+            return None
+
+    def __iter__(self):
+        info = get_worker_info()
+        wid = info.id if info else 0
+        nw = info.num_workers if info else 1
+        my = self.shard_paths[wid::nw]
+        rng = random.Random(42 + wid)
+        while True:
+            rng.shuffle(my)
+            for path in my:
+                try:
+                    pf = pq.ParquetFile(path)
+                    for batch in pf.iter_batches(batch_size=64):
+                        d = batch.to_pydict()
+                        for i in range(len(batch)):
+                            row = {k: d[k][i] for k in d}
+                            out = self._process_sample(row)
+                            if out is not None:
+                                yield out
+                except Exception as e:
+                    print(f"[skip] {path}: {e}")
+                    continue