""" generate_sanskrit_v2.py — Authentic Sanskrit Chanting TTS via IndicF5 ===================================================================== Architecture: Sanskrit Devanagari → Vagdhenu prep_text.py (Devanagari → Kannada script routing) → IndicF5 DiT (flow-matching, 337M params, zero-shot voice cloning) ← Reference Audio (Vagdhenu reference_bank/*.wav — real Sanskrit chanting) → Vocos Vocoder (mel → 24kHz audio) → Gate + Stitch (Vagdhenu-style head/tail trimming, inter-pada silence) → Output WAV Why Kannada routing? Devanagari triggers Hindi schwa-deletion in IndicF5's frontend. Routing through Kannada orthography suppresses that, giving true Sanskrit phonetics — every inherent short 'a' is preserved. Why reference audio? IndicF5 is a zero-shot voice cloner. The "authentic Indian chanting touch" is entirely sourced from the 5-12 second reference WAV clips in vagdhenu/src/reference_bank/, which are recordings of Prof. Prathosh (IISc) chanting Sanskrit in traditional pārāyaṇa style. The model clones voice timbre, swara (pitch contour), pace, and chanting rhythm. CPU Optimization: Default nfe_step=12 instead of 64 gives ~5x speedup with minimal quality loss. At 12 steps, a 10-second verse takes ~40-60s on CPU. Usage: python generate_sanskrit_v2.py "धर्मक्षेत्रे कुरुक्षेत्रे समवेता युयुत्सवः" python generate_sanskrit_v2.py --text "..." --meter anushtubh --output out.wav python generate_sanskrit_v2.py --text "..." --nfe 16 --voice-model vagdhenu Credits: Vagdhenu (github.com/prathoshap/vagdhenu) — Prof. Prathosh, IISc IndicF5 (github.com/ai4bharat/IndicF5) — AI4Bharat Kokoro (github.com/hexgrad/kokoro) — hexgrad """ import os import sys import json import time import argparse import numpy as np import torch # Force UTF-8 output on Windows (prevents cp1252 UnicodeEncodeError for Devanagari/macrons) if hasattr(sys.stdout, "reconfigure"): sys.stdout.reconfigure(encoding="utf-8", errors="replace") if hasattr(sys.stderr, "reconfigure"): sys.stderr.reconfigure(encoding="utf-8", errors="replace") # ── Repo-relative path setup (must come first) ─────────────────────────── HERE = os.path.dirname(os.path.abspath(__file__)) VAGDHENU_SRC = os.path.join(HERE, "vagdhenu", "src") BANK_DIR = os.path.join(VAGDHENU_SRC, "reference_bank") BANK_JSON = os.path.join(BANK_DIR, "bank.json") VOCAB_FILE = os.path.join(BANK_DIR, "vocab.txt") # Add vagdhenu/src to path so we can import its modules directly if VAGDHENU_SRC not in sys.path: sys.path.insert(0, VAGDHENU_SRC) # Import Vagdhenu's Sanskrit→Kannada text converter (module-level so IDEs can resolve it) import prep_text as PT # noqa: E402 (sys.path must be extended first) # ── Load .env file (HF token, etc.) ──────────────────────────────────── try: from dotenv import load_dotenv load_dotenv(os.path.join(HERE, '.env')) except ImportError: pass # python-dotenv optional — set HF_TOKEN as system env var instead # Inject HF_TOKEN so huggingface_hub picks it up automatically _hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HF_API_KEYS") if _hf_token: os.environ["HF_TOKEN"] = _hf_token os.environ["HUGGING_FACE_HUB_TOKEN"] = _hf_token # ── Constants ───────────────────────────────────────────────── SR = 24_000 FALLBACK_METER = "vasantatilaka" # IndicF5 DiT architecture (must match the pre-trained weights exactly) DIT_CFG = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4) # ── Phonology helpers (ported verbatim from Vagdhenu render_core.py) ────────── def _n_aksharas(s: str) -> int: """Count the number of syllables (akṣaras) in a Devanagari or Kannada string.""" n = 0 L = len(s) for i, c in enumerate(s): o = ord(c) indep = (0x0905 <= o <= 0x0914) or (0x0C85 <= o <= 0x0C94) cons = (0x0915 <= o <= 0x0939) or (0x0C95 <= o <= 0x0CB9) if indep: n += 1 elif cons: nxt = s[i + 1] if i + 1 < L else "" if nxt not in ("्", "್"): n += 1 return n def _ends_halant(txt: str) -> bool: """True if the text ends with a virama (halant consonant cluster closer).""" t = txt.rstrip(" ।॥|.,;:!?​") return bool(t) and t[-1] in "्್" def _gate(au: np.ndarray, voice: float = 0.08, sil: float = 0.012, fin: float = 0.015, fout: float = 0.040, lead: float = 0.03, keep: float = 0.06, fric: bool = False, halant: bool = False) -> np.ndarray: """ Trim head and tail silence from synthesized audio (Vagdhenu gate function). Applies a short fade-in and cosine fade-out for clean transitions. """ win = int(0.02 * SR) r = [float(np.sqrt((au[i:i + win] ** 2).mean())) for i in range(0, len(au) - win, win)] n = len(r) if n == 0: return au vs = next((i for i in range(n - 1) if r[i] > voice and r[i + 1] > sil), int(np.argmax(r))) s = vs while s > 0 and r[s - 1] > sil: s -= 1 ve_thr = 0.012 if halant else 0.035 ve = max((i for i in range(n) if r[i] > ve_thr), default=vs) keep_s = 0.12 if halant else keep start = max(0, s * win - int(lead * SR)) end = min(len(au), ve * win + int(keep_s * SR)) out = au[start:end].copy() fi = 0 if fric else int(fin * SR) fo = int((0.018 if halant else fout) * SR) if fi and len(out) > fi: out[:fi] *= np.linspace(0, 1, fi) if fo and len(out) > fo: out[-fo:] *= (np.cos(np.linspace(0, np.pi, fo)) * 0.5 + 0.5) return out def _split_padas(text: str) -> list[str]: """ Split a Sanskrit śloka into hemistich/pāda segments. Splits on newlines first, then on daṇḍas (।). Empty segments are dropped. """ pieces = [] for line in text.replace("॥", "।").replace("|", "।").splitlines(): for seg in line.split("।"): seg = seg.strip() if seg: pieces.append(seg) return pieces or ([text.strip()] if text.strip() else []) # ── Model class ─────────────────────────────────────────────────────────────── class SanskritChantEngine: """ Full pipeline for authentic Sanskrit chanting TTS. On construction, it: 1. Loads the IndicF5 DiT model (downloaded from HuggingFace or a local path) 2. Loads the Vocos vocoder 3. Parses the Vagdhenu reference bank (bank.json) synthesize() then: 1. Splits text into pādas 2. Converts each pāda from Devanagari → Kannada via prep_text.py 3. Looks up the per-meter reference WAV + text from bank.json 4. Runs IndicF5 inference at nfe_step (default 12) per pāda 5. Gates each segment and stitches with silence gaps 6. Returns (sample_rate, audio_float32) """ def __init__( self, voice_model: str = "indicf5", # "indicf5" or "vagdhenu" device: str = "cpu", nfe_step: int = 12, cfg_strength:float= 3.0, speed: float= 0.90, gap: float= 0.55, gap_halant: float= 0.20, ): try: from f5_tts.infer.utils_infer import load_model, load_vocoder from f5_tts.model import DiT except ImportError: raise ImportError( "\n[EdgeSanskrit] IndicF5 is not installed.\n" "Run: pip install git+https://github.com/ai4bharat/IndicF5.git" "@13f7c4d627cc10111aea8fe9c0039462cacacdc7\n" ) self.device = device self.nfe_step = nfe_step self.cfg_strength = cfg_strength self.speed = speed self.gap = gap self.gap_halant = gap_halant # ── 1. Load reference bank ────────────────────────────────────── if not os.path.exists(BANK_JSON): raise FileNotFoundError( f"Vagdhenu reference bank not found at {BANK_JSON}\n" "Make sure you have cloned vagdhenu/ beside this script:\n" " git clone https://github.com/prathoshap/vagdhenu vagdhenu" ) with open(BANK_JSON, encoding="utf-8") as f: self._bank = json.load(f) self._lut: dict = {} for k, v in self._bank.items(): if k.startswith("_") or not isinstance(v, dict) or "wav" not in v: continue self._lut[k.lower()] = v self._lut[v["wav"].replace(".wav", "").lower()] = v print(f"[EdgeSanskrit] Reference bank loaded: {len(self._lut)//2} meters") # ── 2. Load DiT (IndicF5 base or Vagdhenu fine-tune) ─────────── if not os.path.exists(VOCAB_FILE): raise FileNotFoundError( f"IndicF5 vocab.txt not found at {VOCAB_FILE}\n" "It ships with the Vagdhenu clone in reference_bank/vocab.txt" ) print(f"[EdgeSanskrit] Loading IndicF5 DiT on device='{device}'...") self.cfm = load_model(DiT, DIT_CFG, mel_spec_type="vocos", vocab_file=VOCAB_FILE, device=device) if voice_model == "vagdhenu": self._load_vagdhenu_weights() else: self._load_indicf5_weights() self.cfm.eval() print("[EdgeSanskrit] DiT loaded.") # ── 3. Load Vocos vocoder ─────────────────────────────────────── print("[EdgeSanskrit] Loading Vocos vocoder...") vocos_local = os.path.join(HERE, "models", "vocos") if os.path.exists(vocos_local): from vocos import Vocos self.vocoder = Vocos.from_pretrained(vocos_local).to(device) print("[EdgeSanskrit] Loaded Vocos from local bundled models/vocos") else: self.vocoder = load_vocoder("vocos") self._refcache: dict = {} print("[EdgeSanskrit] Vocos loaded. Engine ready.\n") # ── Weight loaders ──────────────────────────────────────────────── def _load_indicf5_weights(self): """Load base IndicF5 weights from local bundle or HuggingFace.""" from safetensors.torch import load_file local_ckpt = os.path.join(HERE, "models", "IndicF5", "model.safetensors") if os.path.exists(local_ckpt): print(f"[EdgeSanskrit] Loading bundled IndicF5 from {local_ckpt}") ckpt = local_ckpt else: from huggingface_hub import hf_hub_download print("[EdgeSanskrit] Downloading ai4bharat/IndicF5 weights from HuggingFace...") ckpt = hf_hub_download(repo_id="ai4bharat/IndicF5", filename="model.safetensors") print(f"[EdgeSanskrit] Checkpoint: {ckpt}") sd = load_file(ckpt, device="cpu") # Strip the EMA module prefix that IndicF5 uses during training sd = { k.replace("ema_model._orig_mod.", "").replace("ema_model.", ""): v for k, v in sd.items() } missing, unexpected = self.cfm.load_state_dict(sd, strict=False) if missing: print(f"[EdgeSanskrit] Missing keys (expected if small): {len(missing)}") def _load_vagdhenu_weights(self): """Download and load Vagdhenu voice-steered checkpoint from prathoshap/vagdhenu.""" from huggingface_hub import hf_hub_download print("[EdgeSanskrit] Downloading prathoshap/vagdhenu voice-steered weights...") ckpt = hf_hub_download(repo_id="prathoshap/vagdhenu", filename="voice_steer_ema.pt") print(f"[EdgeSanskrit] Checkpoint: {ckpt}") ck = torch.load(ckpt, map_location="cpu", weights_only=True) ema = {k.replace("ema_model.", ""): v for k, v in ck["ema_model_state_dict"].items() if k not in ("initted", "step")} self.cfm.load_state_dict(ema, strict=False) # ── Reference audio cache ───────────────────────────────────────── def _get_ref(self, meter: str): """ Look up reference audio and metadata for a given meter. Results are cached after the first load. Returns (ref_audio_path, ref_text, sec_per_syll, ref_duration_s) """ from f5_tts.infer.utils_infer import preprocess_ref_audio_text import torchaudio as ta key = meter.lower().replace(".wav", "").replace("ā", "a").replace("ṭ", "t").replace("ū", "u") # Try exact key first, then normalized version matched_key = None for k in self._lut: if k == key or k.replace("ā", "a") == key or meter.lower() in k: matched_key = k break if matched_key is None: print(f"[EdgeSanskrit] Meter '{meter}' not found in bank → fallback '{FALLBACK_METER}'") matched_key = FALLBACK_METER if matched_key in self._refcache: return self._refcache[matched_key] e = self._lut[matched_key] wav_path = os.path.join(BANK_DIR, e["wav"]) ref_text = e["ref_text"] sps = float(e.get("sec_per_syll", 0.26)) ref_audio, ref_t = preprocess_ref_audio_text(wav_path, ref_text, clip_short=True) ra, wav_sr = ta.load(ref_audio) ref_len = ra.shape[-1] / wav_sr val = (ref_audio, ref_t, sps, ref_len) self._refcache[matched_key] = val print(f"[EdgeSanskrit] Reference: {e['wav']} ({ref_len:.1f}s) — sec/syl={sps}") return val # ── Main synthesis ──────────────────────────────────────────────── def synthesize( self, text: str, meter: str = "anushtubh", seed: int = 60, speed: float = None, sps: float = None, ) -> tuple[int, np.ndarray]: """ Synthesize Sanskrit Devanagari text with authentic chanting prosody. Args: text: Sanskrit text in Devanagari (supports multi-line, daṇḍas) meter: Vṛtta/meter name from bank.json (e.g. 'anushtubh', 'vasantatilaka') seed: Random seed for reproducibility speed: Speed multiplier override (default: self.speed = 0.90) sps: Seconds-per-syllable override Returns: (sample_rate, audio_float32_array) """ from f5_tts.infer.utils_infer import infer_process padas = _split_padas(text) if not padas: raise ValueError("Input text is empty after splitting.") ref_audio, ref_t, ref_sps, ref_len = self._get_ref(meter) if sps is not None: ref_sps = float(sps) spd = float(speed) if speed is not None else self.speed print(f"[EdgeSanskrit] Synthesizing {len(padas)} pāda(s) nfe={self.nfe_step} device={self.device}") # Convert each pada: Devanagari → Kannada routed text pieces = [PT.model_text(p) for p in padas] nsylls = [_n_aksharas(p) for p in pieces] print(f"[EdgeSanskrit] Pādas:") for i, (orig, kan) in enumerate(zip(padas, pieces)): print(f" [{i+1}] {orig}") print(f" → {kan} ({nsylls[i]} akṣaras)") bseg = [] for i, piece in enumerate(pieces): print(f"\n[EdgeSanskrit] Generating pāda {i+1}/{len(pieces)}...") t0 = time.perf_counter() au = None for attempt in range(4): torch.manual_seed(seed + attempt) fix_dur = (ref_len + nsylls[i] * ref_sps) if ref_sps > 0 else None w, sr, _ = infer_process( ref_audio, ref_t, piece, self.cfm, self.vocoder, mel_spec_type = "vocos", speed = spd, nfe_step = self.nfe_step, cfg_strength = self.cfg_strength, device = self.device, fix_duration = fix_dur, ) w = np.array(w, dtype=np.float32) if np.abs(w).max() > 1.5: w = w / 32768.0 rms = float(np.sqrt((w ** 2).mean())) elapsed = time.perf_counter() - t0 dur = len(w) / SR rtf = elapsed / dur if dur > 0 else float("inf") print(f" attempt {attempt+1}: {len(w)/SR:.2f}s audio, " f"RMS={rms:.4f}, RTF={rtf:.2f}x, elapsed={elapsed:.1f}s") if rms > 0.04: au = w break if au is None: print(f" [WARN] All attempts had low RMS — using last result anyway.") au = w # Vagdhenu-style gate: trim head/tail silence halant = _ends_halant(pieces[i]) au = _gate(au, halant=halant) bseg.append(au) # Stitch pādas with silence gaps gap_samples = int(self.gap * SR) halant_gap_samples = int(self.gap_halant * SR) result = [] for i, seg in enumerate(bseg): result.append(seg) if i < len(bseg) - 1: extra = halant_gap_samples if _ends_halant(pieces[i]) else 0 result.append(np.zeros(gap_samples + extra, dtype=np.float32)) final = np.concatenate(result) # Normalize to 97% peak mx = np.abs(final).max() if mx > 0: final = final / mx * 0.97 print(f"\n[EdgeSanskrit] Final audio: {len(final)/SR:.2f}s") return SR, final # ── CLI ─────────────────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser( description="EdgeSanskrit v2 — Authentic Sanskrit Chanting TTS via IndicF5", formatter_class=argparse.RawTextHelpFormatter, ) parser.add_argument( "text", nargs="?", default="धर्मक्षेत्रे कुरुक्षेत्रे समवेता युयुत्सवः ।\nमामकाः पाण्डवाश्चैव किमकुर्वत संजय ॥", help="Sanskrit Devanagari input text (default: BG 1.1)", ) parser.add_argument("-o", "--output", default="sanskrit_chant_v2.wav", help="Output WAV file path") parser.add_argument("-m", "--meter", default="anushtubh", help="Sanskrit meter/vṛtta (e.g. anushtubh, vasantatilaka)") parser.add_argument("--nfe", type=int, default=12, help="NFE diffusion steps (lower=faster, default=12)") parser.add_argument("--speed", type=float, default=0.90, help="Speech rate multiplier (default=0.90 for chant pacing)") parser.add_argument("--voice-model", default="indicf5", choices=["indicf5", "vagdhenu"], help="'indicf5' = base zero-shot model (default)\n" "'vagdhenu' = Prof. Prathosh's voice-steered model (~5GB)") parser.add_argument("--seed", type=int, default=60, help="Random seed (default=60)") args = parser.parse_args() import soundfile as sf engine = SanskritChantEngine( voice_model = args.voice_model, device = "cpu", nfe_step = args.nfe, speed = args.speed, ) t_start = time.perf_counter() sr, audio = engine.synthesize( text = args.text, meter = args.meter, seed = args.seed, ) total_elapsed = time.perf_counter() - t_start audio_dur = len(audio) / sr rtf = total_elapsed / audio_dur if audio_dur > 0 else float("inf") sf.write(args.output, audio, sr, subtype="PCM_16") print(f"\n{'='*60}") print(f" Output : {args.output}") print(f" Audio : {audio_dur:.2f}s at {sr}Hz") print(f" Elapsed: {total_elapsed:.1f}s") print(f" RTF : {rtf:.2f}x ({'faster' if rtf < 1 else 'slower'} than real-time)") print(f"{'='*60}\n") if __name__ == "__main__": main()