GAASH-Lab/Matcha-TTS-Kashmiri-Demo

1 2import gradio as gr3from pathlib import Path4import torch5import urllib.request6import os  7 8# HuggingFace Spaces GPU support9try:10    import spaces11    SPACES_AVAILABLE = True12except ImportError:13    SPACES_AVAILABLE = False14    print("[!] spaces module not available, running without GPU decorator")15import soundfile as sf16import traceback17import huggingface_hub18from huggingface_hub import hf_hub_download19 20# Patch for older diffusers compatibility with newer huggingface_hub21if not hasattr(huggingface_hub, "cached_download"):22    huggingface_hub.cached_download = hf_hub_download23 24from transformers import AutoTokenizer, AutoModelForCausalLM25from peft import PeftModel26from matcha.models.matcha_tts import MatchaTTS27from matcha.hifigan.models import Generator as HiFiGAN28from matcha.hifigan.config import v129from matcha.hifigan.env import AttrDict30from matcha.text import text_to_sequence31from matcha.utils.utils import intersperse32 33HF_TOKEN = os.getenv("HF_TOKEN")34 35DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")36MODEL_REPO = "GAASH-Lab/Matcha-TTS-Kashmiri"37 38def load_models():39    print("[*] Downloading GAASH-Lab checkpoint...")40    ckpt = hf_hub_download(repo_id=MODEL_REPO, filename="model.ckpt", token=HF_TOKEN)41    model = MatchaTTS.load_from_checkpoint(ckpt, map_location=DEVICE, weights_only=False)42    model.eval()43    44    print("[*] Loading HiFi-GAN vocoder...")45    # The file 'generator_v1' is what the code calls 'hifigan_T2_v1'46    # We download it from the official GitHub release if not found locally47    vocoder_path = Path("hifigan_T2_v1")48    if not vocoder_path.exists():49        url = "https://github.com/shivammehta25/Matcha-TTS-checkpoints/releases/download/v1.0/generator_v1"50        urllib.request.urlretrieve(url, vocoder_path)51 52    vocoder = HiFiGAN(AttrDict(v1)).to(DEVICE)53    state_dict = torch.load(vocoder_path, map_location=DEVICE)54    vocoder.load_state_dict(state_dict['generator'])55    vocoder.eval()56    vocoder.remove_weight_norm()57    58    return model, vocoder59 60# Translation Config61TRANSLATION_BASE_MODEL = "sarvamai/sarvam-translate"62TRANSLATION_ADAPTER = "GAASH-Lab/Sarvam-Kashmiri-finetuned"63 64# Global cache for translation model (loaded lazily when GPU is available)65_trans_cache = {"tokenizer": None, "model": None, "loaded": False}66 67def load_translation_models():68    """Load translation model lazily on first use (CPU deployment)."""69    global _trans_cache70    71    if _trans_cache["loaded"]:72        return _trans_cache["tokenizer"], _trans_cache["model"]73    74    print("[*] Loading Sarvam Translate Adapter (CPU mode)...")75    try:76        # Load the tokenizer with left padding (required for causal LM)77        tokenizer = AutoTokenizer.from_pretrained(TRANSLATION_BASE_MODEL, trust_remote_code=True)78        tokenizer.padding_side = "left"79        if tokenizer.pad_token is None:80            tokenizer.pad_token = tokenizer.eos_token81        82        # Load the base model on CPU with bfloat16 to reduce memory83        # bfloat16 is better supported on CPU than float1684        print("[*] Loading base model on CPU (bfloat16)...")85        base_model = AutoModelForCausalLM.from_pretrained(86            TRANSLATION_BASE_MODEL,87            torch_dtype=torch.bfloat16,88            device_map="cpu",89            low_cpu_mem_usage=True,90            trust_remote_code=True91        )92        93        # Load the LoRA adapter94        print("[*] Loading LoRA adapter...")95        model = PeftModel.from_pretrained(base_model, TRANSLATION_ADAPTER)96        97        # Merge LoRA weights into base model for faster inference98        # This eliminates adapter overhead during generation99        print("[*] Merging LoRA weights for faster inference...")100        model = model.merge_and_unload()101        model.eval()102        103        print(f"[+] Translation model loaded and merged successfully on CPU.")104        _trans_cache["tokenizer"] = tokenizer105        _trans_cache["model"] = model106        _trans_cache["loaded"] = True107        return tokenizer, model108    except Exception as e:109        print(f"[-] Error loading translation model: {e}")110        traceback.print_exc()111        return None, None112 113# Load TTS models at startup (they're smaller)114model, vocoder = load_models()115# Translation model will be loaded lazily when GPU is available116 117def _translate_impl(text):118    """Internal translation implementation - matching evaluate_model.py approach."""119    # Load model lazily (will be cached after first load)120    trans_tokenizer, trans_model = load_translation_models()121    122    if trans_model is None:123        return "Translation model unavailable."124    125    # Build chat messages (matching evaluate_model.py)126    messages = [127        {"role": "system", "content": "Translate the text below to Kashmiri."},128        {"role": "user", "content": text},129    ]130    131    try:132        # Apply chat template (matching evaluate_model.py)133        prompt = trans_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)134        inputs = trans_tokenizer(prompt, padding=True, truncation=True, max_length=512, return_tensors="pt")135        136        # Move inputs to model's device137        inputs = {k: v.to(trans_model.device) for k, v in inputs.items()}138        139        print(f"[DEBUG] Input tokens: {inputs['input_ids'].shape[1]}")140        141    except Exception as e:142        print(f"Chat template error: {e}")143        traceback.print_exc()144        return "Error in translation template."145 146    try:147        import time148        start_time = time.time()149        print("[DEBUG] Starting generation...")150        151        # Generation settings optimized for CPU inference152        # - Greedy decoding (do_sample=False) is faster than sampling153        # - Same quality as temp=0.01 which was near-greedy anyway154        with torch.no_grad():155            generated = trans_model.generate(156                **inputs,157                max_new_tokens=256,  # Keep full length for long texts158                do_sample=False,     # Greedy decoding for speed159                num_beams=1,         # No beam search overhead160            )161        162        elapsed = time.time() - start_time163        print(f"[DEBUG] Generation completed in {elapsed:.2f}s")164        165        # Decode only the new tokens (matching evaluate_model.py)166        input_len = inputs['input_ids'].shape[1]167        output_ids = generated[0][input_len:]168        decoded = trans_tokenizer.decode(output_ids, skip_special_tokens=True).replace("\n", "")169        170        print(f"[DEBUG] New tokens: {len(output_ids)}")171        print(f"[DEBUG] Decoded: '{decoded}'")172        173        return decoded.strip()174        175    except Exception as e:176        print(f"Generation error: {e}")177        traceback.print_exc()178        return "Error during translation generation."179 180# Simple wrapper function for CPU deployment181def translate(text):182    return _translate_impl(text)183 184 185# --- Update the function signature to accept two arguments ---186@torch.inference_mode()187def process(text, speaker_id, n_timesteps=10):188    # 1. Kashmiri script normalization189    text = text.replace("ي", "ی").replace("ك", "ک").replace("۔", "").strip()190    191    # 2. Text to Sequence192    cleaner = "basic_cleaners" 193    sequence, _ = text_to_sequence(text, [cleaner])194    195    # Filter out any non-integer values (unknown characters not in vocabulary)196    # This happens when text contains characters not supported by the TTS model197    filtered_sequence = [s for s in sequence if isinstance(s, int)]198       199    x = torch.tensor(intersperse(filtered_sequence, 0), dtype=torch.long, device=DEVICE)[None]200    x_lengths = torch.tensor([x.shape[-1]], dtype=torch.long, device=DEVICE)201    202    # 3. Use the Speaker ID from the interface203    # Even if you only use one voice, the model requires this tensor204    spks = torch.tensor([int(speaker_id)], device=DEVICE, dtype=torch.long)205    206    # 4. Generate Mel-spectrogram207    output = model.synthesise(208        x, 209        x_lengths, 210        n_timesteps=n_timesteps, 211        temperature=0.667, 212        spks=spks,213        length_scale=1.0214    )215    216    # 5. Generate Waveform217    audio = vocoder(output['mel']).clamp(-1, 1).cpu().squeeze().numpy()218    output_path = "out.wav"219    sf.write(output_path, audio, 22050)220    return output_path221 222# --- Gradio UI with Translation Option ---223with gr.Blocks(title="GAASH-Lab: Kashmiri TTS & Translation") as demo:224    gr.Markdown("# GAASH-Lab: Kashmiri TTS & Translation")225    gr.Markdown("Enter text in English (check the box) or Kashmiri directly.")226    227    with gr.Row():228        with gr.Column():229            input_text = gr.Textbox(label="Input Text", placeholder="Type here...")230            is_english = gr.Checkbox(label="Input is English (Translate first)", value=False)231            speaker_radio = gr.Radio(choices=["Male", "Female"], value="Male", label="Speaker Voice")232            quality_radio = gr.Radio(233                choices=["Low (fast)", "Medium", "High"], 234                value="Low (fast)", 235                label="Quality"236            )237            gen_btn = gr.Button("Generate Speech", variant="primary")238        239        with gr.Column():240            trans_view = gr.Textbox(label="Processed/Translated Kashmiri Text", interactive=False)241            audio_output = gr.Audio(label="Audio", type="filepath")242 243    def pipeline(text, is_eng, spk_voice, quality):244        spk_id = 422 if spk_voice == "Male" else 423245        246        if "Low" in quality:247            steps = 10248        elif "Medium" in quality:249            steps = 50250        else:251            steps = 500252 253        processed_text = text254        if is_eng:255            print(f"Translating input: {text}")256            processed_text = translate(text)257        258        print(f"Synthesizing for: {processed_text}")259        audio_path = process(processed_text, spk_id, steps)260        return processed_text, audio_path261 262    gen_btn.click(263        pipeline, 264        inputs=[input_text, is_english, speaker_radio, quality_radio], 265        outputs=[trans_view, audio_output]266    )267 268demo.launch(ssr_mode=False)