Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| # app.py — Gradio UI + inference (Hugging Face ZeroGPU compatible) | |
| import os | |
| import sys | |
| import time | |
| import json | |
| import torch | |
| import warnings | |
| import numpy as np | |
| import gradio as gr | |
| import yaml | |
| import spaces | |
| import re | |
| import threading | |
| warnings.filterwarnings("ignore") | |
| # HF repo download | |
| from huggingface_hub import snapshot_download | |
| # --- THREAD-SAFE MISHKAL DIACRITIZATION --- | |
| import mishkal.tashkeel | |
| # Create a thread-local storage object to keep Mishkal instances isolated per Gradio thread | |
| _thread_local = threading.local() | |
| def get_vocalizer(): | |
| # If the current thread doesn't have a vocalizer yet, create one | |
| if not hasattr(_thread_local, "vocalizer"): | |
| _thread_local.vocalizer = mishkal.tashkeel.TashkeelClass() | |
| return _thread_local.vocalizer | |
| def diacritize_arabic(text: str) -> str: | |
| # Safely fetch the thread-specific vocalizer and process the text | |
| vocalizer = get_vocalizer() | |
| diacritized = vocalizer.tashkeel(text) | |
| # Scrub out any hidden control characters (like \x01) that Mishkal accidentally leaves behind | |
| cleaned = re.sub(r'[\x00-\x08\x0B-\x1F\x7F]', '', diacritized) | |
| return cleaned.strip() | |
| print("Thread-safe Mishkal diacritizer initialized.") | |
| # ------------------------------------------ | |
| REPO_ID = "FashionFlora/SlopTTS" | |
| HF_TOKEN = os.environ.get("hf_token") or os.environ.get("HF_TOKEN") | |
| LOCAL_REPO_DIR = "./sloptts_runtime" | |
| print("Downloading repository...", REPO_ID) | |
| repo_path = snapshot_download( | |
| repo_id=REPO_ID, | |
| repo_type="model", | |
| token=HF_TOKEN, | |
| local_dir=LOCAL_REPO_DIR, | |
| allow_patterns=["*"], | |
| ) | |
| print("Downloaded to:", repo_path) | |
| if repo_path not in sys.path: | |
| sys.path.insert(0, repo_path) | |
| from fsq_flow_dataset import ( | |
| symbol_to_id, | |
| BOS_ID, | |
| EOS_ID, | |
| PAD_ID, | |
| SIL_ID, | |
| TextCleaner, | |
| ) | |
| _text_cleaner = TextCleaner() | |
| print("Imported TextCleaner from fsq_flow_dataset") | |
| try: | |
| from phonemizer import phonemize as phonemize_func | |
| from phonemizer.separator import Separator | |
| def phonemize_text(text: str, language_code: str = "en-us") -> str: | |
| sep = Separator(phone="", word=" ", syllable="") | |
| ipa = phonemize_func( | |
| text, | |
| backend="espeak", | |
| language=language_code, | |
| with_stress=True, | |
| separator=sep, | |
| strip=True, | |
| preserve_punctuation=True, | |
| njobs=1, | |
| ) | |
| return ipa | |
| print("Phonemizer is available.") | |
| except Exception as e: | |
| print("Phonemizer not available in environment. Error:", e) | |
| def phonemize_text(text: str, language_code: str = "en-us") -> str: | |
| return text | |
| def split_sentences(text: str) -> list[str]: | |
| """ | |
| Split text into sentences. | |
| Handles: | |
| - Hindi/Devanagari : danda । (U+0964) and double-danda ॥ (U+0965) | |
| - Latin punctuation : . ! ? | |
| - Paragraph breaks : blank lines / newlines | |
| The sentence-ending punctuation is kept attached to the preceding sentence. | |
| """ | |
| text = text.strip() | |
| paragraphs = re.split(r'\n{2,}', text) | |
| sentences = [] | |
| for para in paragraphs: | |
| parts = re.split(r'(?<=[।॥.!?])\s+', para.strip()) | |
| for p in parts: | |
| p = p.strip() | |
| if p: | |
| sentences.append(p) | |
| return sentences | |
| # -------------------------------------------------------------------- | |
| # 3) Paths inside downloaded repo | |
| # -------------------------------------------------------------------- | |
| PREDICTOR_CONFIG = os.path.join(repo_path, "Configs", "config_fsq_flow_convnext.yml") | |
| CODEC_CONFIG = os.path.join(repo_path, "Configs", "config_codec_mel_speaker.yml") | |
| PREDICTOR_CKPT = os.path.join(repo_path, "fsq_flow_convnext_training", "model_epoch_30.pt") | |
| CODEC_CKPT = os.path.join(repo_path, "mel_speaker_codec_training", "epoch_mel_speaker_codec_00002.pth") | |
| # -------------------------------------------------------------------- | |
| # 4) LANGUAGE <-> SPEAKER mapping | |
| # -------------------------------------------------------------------- | |
| LANGUAGES = { | |
| "english": {"id": 0, "espeak": "en-us", "speakers": [0, 1], "display": "English"}, | |
| "arabic": {"id": 1, "espeak": "ar", "speakers": [2, 3], "display": "Arabic"}, | |
| "german": {"id": 2, "espeak": "de", "speakers": [4], "display": "German"}, | |
| "russian": {"id": 3, "espeak": "ru", "speakers": [5], "display": "Russian"}, | |
| "polish": {"id": 4, "espeak": "pl", "speakers": [6, 7], "display": "Polish"}, | |
| "spanish": {"id": 5, "espeak": "es", "speakers": [8], "display": "Spanish"}, | |
| "hindi": {"id": 6, "espeak": "hi", "speakers": [9, 10], "display": "Hindi"}, | |
| } | |
| STORIES = { | |
| "arabic": { | |
| "Story 1: Desert Oasis": "في قلب الصحراء الواسعة، ظهرت واحة خضراء كأنها حلم. كانت أشجار النخيل تتمايل بلطف مع نسيم الصباح. تدفق ينبوع مياه صافية يعكس أشعة الشمس الذهبية. اقتربت قافلة متعبة من المكان، فشعر المسافرون بالراحة والأمل بعد أيام من السفر الشاق.", | |
| "Story 2: The Old Market": "كان السوق القديم يعج بالحركة، روائح التوابل والقهوة تملأ الهواء، والناس يتجولون بين المتاجر. في إحدى الزوايا، كان هناك صانع فخار يشكل الطين ببراعة لا مثيل لها. توقف الأطفال لمراقبته بانبهار، بينما تعالت أصوات الباعة وهم يعرضون بضائعهم الملونة للزوار.", | |
| "Story 3: Sea Voyage": "تلاطمت الأمواج بسفينة الصيد الصغيرة تحت سماء مرصعة بالنجوم. كان الصياد ينظر إلى الأفق البعيد. فجأة، التمع ضوء غامض تحت سطح الماء، يكشف عن أسرار المحيط العميقة. شد شباكه بقوة، متسائلاً عما إذا كان هذا الصيد سيغير مجرى حياته إلى الأبد.", | |
| "Story 4: The Wise King": "حكم الملك العادل بلاده بحكمة وسلام، فازدهرت المدن وعاش الناس في أمان وطمأنينة. كان بابه مفتوحاً دائماً للفقير قبل الغني، يستمع إلى شكاواهم بصدر رحب. وفي أحد الأيام، جاءه حكيم غريب حاملاً لغزاً قديماً، واعداً إياه بكنز من المعرفة إذا تمكن من حله." | |
| }, | |
| "english": { | |
| "Story 1: The Magic Forest": "Deep in the ancient woods, glowing mushrooms illuminated the path. A small fox watched carefully as the traveler passed by. The trees seemed to whisper ancient secrets as the wind brushed through their silver leaves. Suddenly, a gentle melody drifted from a hidden glade, inviting the wanderer to discover what lay beyond the mist.", | |
| "Story 2: Space Journey": "The rocket shuddered as it broke through the atmosphere. Outside the window, millions of stars sparkled against the black void. The crew unbuckled their harnesses, floating weightlessly in the cabin as they gazed at the blue marble shrinking behind them. Their long voyage to the outer rim had finally begun, full of unknown perils and extraordinary wonders.", | |
| "Story 3: A Quiet Morning": "Sunlight filtered through the curtains, warming the small room. The smell of fresh coffee drifted from the kitchen. An old orange cat stretched lazily on the rug, purring as a gentle breeze rattled the windowpanes. It was the perfect start to a Sunday, holding the promise of a peaceful, unhurried day reading by the fire.", | |
| "Story 4: The Inventor": "Gears clicked and whirred as the mechanical bird spread its wings. With a gentle push, it took flight across the workshop. The inventor smiled, wiping a smudge of oil from her cheek. After months of failed attempts, her masterpiece was finally soaring, its brass feathers catching the golden light of the afternoon sun." | |
| }, | |
| "german": { | |
| "Story 1: The Old Forest": "Tief im alten Wald leuchteten kleine Pilze den Weg. Ein kleiner Fuchs beobachtete den Wanderer ganz genau. Plötzlich knackte ein Ast, und eine riesige Eule schwebte lautlos durch die Baumkronen. Der Wanderer hielt den Atem an, fasziniert von der magischen Stille, die diesen unberührten Ort umgab.", | |
| "Story 2: A Journey": "Der Zug rollte langsam durch die schneebedeckten Berge. Aus dem Fenster sah man kleine Dörfer und gefrorene Seen. Ein Passagier nippte an seinem heißen Tee und schlug ein altes Buch auf. Die rhythmischen Geräusche der Räder auf den Schienen wirkten so beruhigend, dass er schon bald in einen tiefen Schlaf fiel.", | |
| "Story 3: The Small Village": "Die Morgensonne wärmte die Dächer des kleinen Dorfes. Der Bäcker hatte bereits frisches Brot gebacken. Der herrliche Duft zog durch die engen, gepflasterten Gassen und weckte die Bewohner. Kinder liefen lachend zum Marktplatz, wo das bunte Treiben des Wochenmarktes langsam begann.", | |
| "Story 4: The Cold Winter": "Ein kalter Wind wehte über das Feld. Die Kinder bauten einen großen Schneemann vor ihrem Haus. Sie gaben ihm eine Karotte als Nase und wickelten einen dicken, roten Schal um seinen Hals. Als es allmählich dunkel wurde, rief ihre Mutter sie ins warme Haus, wo schon heißer Kakao auf sie wartete." | |
| }, | |
| "russian": { | |
| "Story 1: Winter Tale": "Глубоко в лесу снег блестел в лунном свете. Маленькая лиса осторожно ступала по белой тропинке. Ветви елей сгибались под тяжестью пушистого снега, создавая сказочные арки. Вдали послышался тихий вой волка, но лисица лишь ускорила шаг, стремясь к своему теплому укрытию под корнями старого дуба.", | |
| "Story 2: Space Flight": "Ракета вздрогнула, прорываясь сквозь атмосферу. За окном сияли миллионы звезд на фоне черной пустоты. Командир корабля нажал несколько кнопок на мерцающей панели, переводя двигатели в режим круиза. Впереди их ждала неизвестная планета, покрытая загадочными фиолетовыми облаками и скрывающая тайны древней цивилизации.", | |
| "Story 3: Old Castle": "Старый замок стоял на вершине холма. Его высокие башни прятались в утреннем тумане. Внутри замка царила тишина, нарушаемая лишь гулким эхом шагов по каменному полу. Местные жители поговаривали, что по ночам в длинных коридорах можно услышать печальную музыку, которую играет невидимый музыкант.", | |
| "Story 4: Ocean Mystery": "Волны мягко бились о берег. На песке лежала красивая ракушка, принесенная ночным штормом. Мальчик подобрал её и приложил к уху, надеясь услышать шум моря. К его удивлению, вместо шума волн он разобрал тихий, мелодичный шепот, рассказывающий истории о затонувших кораблях и русалках." | |
| }, | |
| "polish": { | |
| "Story 1: Wawel Dragon": "Głęboko pod zamkiem wawelskim spał stary smok. Ludzie w mieście opowiadali o nim niesamowite legendy. Niektórzy twierdzili, że jego łuski były ze szczerego złota, a inni, że ział zielonym ogniem. Każdej nocy rycerze pełnili wartę na murach, z niepokojem nasłuchując dudniącego chrapania dochodzącego z ciemnej jaskini nad rzeką.", | |
| "Story 2: Mountain Trip": "Słońce powoli wschodziło nad Tatrami. Wiatr szumiał w koronach drzew, gdy wyruszaliśmy na szlak. Powietrze było rześkie i pachniało sosnowymi igłami. Po kilku godzinach wspinaczki dotarliśmy na szczyt, skąd rozpościerał się zapierający dech w piersiach widok na dolinę spowitą poranną mgłą.", | |
| "Story 3: Golden Autumn": "Liście spadały z drzew jak złoty deszcz. W parku biegały wiewiórki, zbierając zapasy na zimę. Stare dęby przybrały miedziane barwy, a kasztany z głośnym stukotem uderzały o brukowane alejki. Dzieci z radością zbierały najpiękniejsze z nich, by w domach tworzyć urocze, jesienne ludziki.", | |
| "Story 4: Secret Garden": "Za starym murem krył się tajemniczy ogród. Kwitły tam najpiękniejsze róże o zapachu słodkim jak miód. Pośrodku znajdowała się opuszczona fontanna porośnięta mchem, w której woda przestała płynąć lata temu. Kiedy jednak mała dziewczynka dotknęła kamiennej rzeźby ptaka, ze źródła nagle wytrysnął krystalicznie czysty strumień." | |
| }, | |
| "spanish": { | |
| "Story 1: Hidden Treasure": "En lo profundo del bosque antiguo, los árboles susurraban secretos. Un explorador buscaba el tesoro perdido del viejo rey. Después de caminar durante horas, encontró una cueva oculta detrás de una cascada brillante. Al entrar, la luz de su antorcha reveló paredes cubiertas de inscripciones antiguas y, en el centro, un cofre de oro macizo.", | |
| "Story 2: The Lost City": "La ciudad perdida emergió de la niebla. Sus calles empedradas contaban historias de una civilización olvidada. Las grandes pirámides de piedra estaban cubiertas de enredaderas verdes y flores exóticas. A medida que el equipo de arqueólogos avanzaba, sentían que los ojos de las estatuas guardianas los seguían en silencio.", | |
| "Story 3: Walk on the Beach": "Las olas rompían suavemente en la orilla. El sol se ponía, pintando el cielo de colores cálidos y hermosos. Las gaviotas volaban en círculos buscando su última comida del día. Caminando descalzo por la arena húmeda, sintió una paz profunda que le hizo olvidar todas sus preocupaciones.", | |
| "Story 4: Festival of Lights": "El festival de luces iluminó la plaza central. La música sonaba mientras la gente bailaba con alegría. Los vendedores ofrecían dulces tradicionales y faroles de papel de todos los colores. A la medianoche, cientos de linternas flotantes fueron liberadas al cielo, creando un espectáculo mágico que parecía un río de estrellas." | |
| }, | |
| "hindi": { | |
| "Story 1: Lion and Mouse": "जंगल के राजा शेर को एक छोटे चूहे ने जाल से बचाया। उस दिन से दोनों अच्छे दोस्त बन गए। एक दिन जब शेर बीमार पड़ा, तो चूहा उसके लिए दूर पहाड़ से एक दुर्लभ जड़ी-बूटी लेकर आया। शेर ने मुस्कुराकर कहा कि सच्ची दोस्ती आकार या ताकत नहीं, बल्कि दिल की सच्चाई देखती है।", | |
| "Story 2: Magic Lamp": "पुराने बाजार में उसे एक जादुई चिराग मिला। जैसे ही उसने उसे रगड़ा, एक बड़ा जिन्न बाहर आ गया। जिन्न ने गरजती हुई आवाज़ में कहा कि वह केवल तीन इच्छाएँ पूरी कर सकता है। लड़के ने बहुत सोच-समझकर अपनी पहली इच्छा मांगी, जिससे उसके पूरे गाँव की गरीबी हमेशा के लिए दूर हो गई।", | |
| "Story 3: Starry Night": "रात का आसमान तारों से भरा हुआ था। ठंडी हवा चल रही थी और गाँव में सब शांत था। छतों وبيठе लोग धीरे-धीरे कहानियाँ सुना रहे थे। तभी आसमान से एक टूटता तारा गुज़रा, और सभी बच्चों ने आँखें बंद करके अपनी-अपनी गुप्त मन्नतें मांग लीं।", | |
| "Story 4: Old Fort": "वह पुराना किला पहाड़ी की चोटी पर खड़ा था। उसकी दीवारें सदियों पुरानी कहानियाँ बयां कर रही थीं। जब सूरज ढलने लगा, तो किले के खंडहरों में एक अजीब सी चमक आ गई। ऐसा लगता था जैसे पुराने राजा और उनके सैनिक आज भी वहां अपने खोए हुए साम्राज्य की रक्षा कर रहे हों।" | |
| } | |
| } | |
| # -------------------------------------------------------------------- | |
| # 5) Predictor & codec builder/loaders | |
| # -------------------------------------------------------------------- | |
| def build_predictor(pred_cfg: dict, device: torch.device): | |
| from Modules.predictor_fsq_flow_convnext import FSQFlowPredictorConvNeXt | |
| mp = pred_cfg.get("model_params", {}) | |
| return FSQFlowPredictorConvNeXt( | |
| vocab_size = int(mp.get("vocab_size", 183)), | |
| text_dim = int(mp.get("text_dim", 512)), | |
| text_num_layers = int(mp.get("text_num_layers", 4)), | |
| text_kernel_size = int(mp.get("text_kernel_size", 7)), | |
| num_languages = int(mp.get("num_languages", 10)), | |
| language_dim = int(mp.get("language_dim", 64)), | |
| n_speakers = int(mp.get("n_speakers", 11)), | |
| speaker_dim = int(mp.get("speaker_dim", 128)), | |
| fsq_levels = mp.get("fsq_levels", [4, 4, 4, 4, 4, 4]), | |
| hidden_dim = int(mp.get("hidden_dim", 512)), | |
| cond_layers = int(mp.get("flow_num_layers", 6)), | |
| flow_num_heads = int(mp.get("flow_num_heads", 8)), | |
| n_phonemes = int(mp.get("n_phonemes", 176)), | |
| duration_n_flows = int(mp.get("duration_n_flows", 4)), | |
| duration_kernel_size = int(mp.get("duration_kernel_size", 3)), | |
| dropout = 0.0, | |
| codec_stride = int(mp.get("codec_stride", 2)), | |
| cfg_dropout_prob = 0.0, | |
| pad_token_id = PAD_ID if "PAD_ID" in globals() else 3, | |
| ).to(device) | |
| def build_codec(codec_cfg: dict, device: torch.device): | |
| from Modules.codec_decoder_mel_speaker import MelCodecVocoderSpeaker | |
| mp = codec_cfg.get("model_params", {}) | |
| dec = mp.get("decoder", {}) | |
| return MelCodecVocoderSpeaker( | |
| num_speakers = int(mp.get("num_speakers", mp.get("n_speakers", 11))), | |
| speaker_dim = int(mp.get("speaker_dim", 128)), | |
| n_mels = int(mp.get("n_mels", 80)), | |
| latent_dim = int(mp.get("latent_dim", 512)), | |
| hidden_dim = int(mp.get("hidden_dim", 512)), | |
| codec_strides = mp.get("codec_strides", [2, 2]), | |
| fsq_levels = mp.get("fsq_levels", [4] * 6), | |
| upsample_rates = dec.get("upsample_rates", mp.get("upsample_rates", [7, 9])), | |
| gen_istft_n_fft = dec.get("gen_istft_n_fft", mp.get("gen_istft_n_fft", 28)), | |
| gen_istft_hop_size = dec.get("gen_istft_hop_size", mp.get("gen_istft_hop_size", 7)), | |
| sample_rate = int(mp.get("sample_rate", 44100)), | |
| n_phonemes = int(mp.get("n_phonemes", 176)), | |
| num_languages = int(mp.get("num_languages", 10)), | |
| language_dim = int(mp.get("language_dim", 64)), | |
| ).to(device) | |
| def load_predictor(predictor, ckpt_path: str): | |
| print("Loading predictor checkpoint:", ckpt_path) | |
| state = torch.load(ckpt_path, map_location="cpu") | |
| sd = state.get("predictor", state) | |
| sd = {k.replace("module.", ""): v for k, v in sd.items()} | |
| predictor.load_state_dict(sd, strict=False) | |
| predictor.eval() | |
| return predictor | |
| def load_codec(codec, ckpt_path: str): | |
| print("Loading codec checkpoint:", ckpt_path) | |
| state = torch.load(ckpt_path, map_location="cpu") | |
| if "net" in state: | |
| inner = state["net"] | |
| sd = inner.get("codec", inner) | |
| elif "state_dict" in state: | |
| sd = state["state_dict"] | |
| else: | |
| sd = state | |
| sd = {k.replace("module.", ""): v for k, v in sd.items()} | |
| codec.load_state_dict(sd, strict=False) | |
| codec.eval() | |
| return codec | |
| # -------------------------------------------------------------------- | |
| # 6) Synthesis function (single sentence processing) | |
| # -------------------------------------------------------------------- | |
| def synthesize_sentence( | |
| text: str, | |
| predictor, | |
| codec, | |
| device: torch.device, | |
| speaker_id: int, | |
| language_id: int, | |
| prev_text: str = None, | |
| next_text: str = None, | |
| n_steps: int = 64, | |
| cfg_strength: float = 4.0, | |
| temperature: float = 0.8, | |
| dur_noise_scale: float = 0.667, | |
| length_scale: float = 1.0, | |
| prev_z_context: torch.Tensor = None, # NOWE: kontekst z poprzedniego zdania | |
| ): | |
| # STRIP HIDDEN CONTROL CHARACTERS | |
| # This specifically removes \x01 and other invisible non-printable ASCII junk | |
| text = re.sub(r'[\x00-\x08\x0B-\x1F\x7F]', '', text) | |
| espeak_code = "en-us" | |
| for k, v in LANGUAGES.items(): | |
| if v["id"] == language_id: | |
| espeak_code = v["espeak"] | |
| break | |
| # -- ARABIC DIACRITIZATION STEP -- | |
| if espeak_code == "ar": | |
| text = diacritize_arabic(text) | |
| print("Diacritized Arabic text:", text) | |
| ipa = phonemize_text(text, espeak_code) | |
| print("IPA:", ipa) | |
| ids = _text_cleaner(ipa) | |
| if not ids: | |
| raise ValueError(f"Empty phoneme sequence for text: {text}") | |
| def wrap_with_bos_eos(ids): | |
| try: | |
| return [BOS_ID, SIL_ID] + ids + [SIL_ID, EOS_ID] | |
| except Exception: | |
| return [0, 1] + ids + [1, 2] | |
| tgt_wrapped = wrap_with_bos_eos(ids) | |
| prev_wrapped = [] | |
| if prev_text: | |
| # Also clean context text just in case | |
| prev_text = re.sub(r'[\x00-\x08\x0B-\x1F\x7F]', '', prev_text) | |
| prev_ids = _text_cleaner(phonemize_text(prev_text, espeak_code)) | |
| prev_wrapped = wrap_with_bos_eos(prev_ids) | |
| next_wrapped = [] | |
| if next_text: | |
| next_text = re.sub(r'[\x00-\x08\x0B-\x1F\x7F]', '', next_text) | |
| next_ids = _text_cleaner(phonemize_text(next_text, espeak_code)) | |
| next_wrapped = wrap_with_bos_eos(next_ids) | |
| full_ctx = prev_wrapped + tgt_wrapped + next_wrapped | |
| tok_start = len(prev_wrapped) | |
| tok_end = tok_start + len(tgt_wrapped) | |
| tgt_t = torch.tensor(tgt_wrapped, dtype=torch.long, device=device).unsqueeze(0) | |
| tgt_l = torch.tensor([len(tgt_wrapped)], dtype=torch.long, device=device) | |
| ctx_t = torch.tensor(full_ctx, dtype=torch.long, device=device).unsqueeze(0) | |
| ctx_l = torch.tensor([len(full_ctx)], dtype=torch.long, device=device) | |
| spk_t = torch.tensor([speaker_id], dtype=torch.long, device=device) | |
| lang_t = torch.tensor([language_id], dtype=torch.long, device=device) | |
| t_s = torch.tensor([tok_start], dtype=torch.long, device=device) | |
| t_e = torch.tensor([tok_end], dtype=torch.long, device=device) | |
| DUR_CONTEXT_TOKENS = 0 | |
| prev_tail = prev_wrapped[-DUR_CONTEXT_TOKENS:] if DUR_CONTEXT_TOKENS > 0 else [] | |
| next_head = next_wrapped[:DUR_CONTEXT_TOKENS] if DUR_CONTEXT_TOKENS > 0 else [] | |
| ext_phonemes = prev_tail + tgt_wrapped + next_head | |
| dur_curr_start = len(prev_tail) | |
| dur_curr_end = dur_curr_start + len(tgt_wrapped) | |
| ext_t = torch.tensor(ext_phonemes, dtype=torch.long, device=device).unsqueeze(0) | |
| ext_l = torch.tensor([len(ext_phonemes)], dtype=torch.long, device=device) | |
| ext_ts = torch.tensor([tok_start - len(prev_tail)], dtype=torch.long, device=device) | |
| ext_te = torch.tensor([tok_end + len(next_head)], dtype=torch.long, device=device) | |
| dur_out = predictor.infer_durations( | |
| phonemes=ext_t, phoneme_lengths=ext_l, | |
| speaker_ids=spk_t, language_ids=lang_t, | |
| ctx_tokens=ctx_t, ctx_lens=ctx_l, | |
| token_start_indices=ext_ts, token_end_indices=ext_te, | |
| noise_scale=dur_noise_scale, length_scale=length_scale, | |
| ) | |
| pred_dur = dur_out["durations"][:, dur_curr_start:dur_curr_end] | |
| frame_lengths = pred_dur.sum(dim=-1).long() | |
| codec_stride = getattr(predictor, "codec_stride", 2) | |
| pred_codec_len = int(((frame_lengths[0] + codec_stride - 1) // codec_stride).item()) | |
| inf_out = predictor.inference( | |
| phonemes=tgt_t, phoneme_lengths=tgt_l, | |
| durations=pred_dur, | |
| speaker_ids=spk_t, language_ids=lang_t, | |
| n_steps=n_steps, temperature=temperature, cfg_strength=cfg_strength, | |
| ctx_tokens=ctx_t, ctx_lens=ctx_l, | |
| token_start_indices=t_s, token_end_indices=t_e, | |
| ) | |
| z_pred = inf_out["pred_embeddings"][:, :pred_codec_len, :] | |
| # --- NOWA LOGIKA CIĄGŁOŚCI VOCODERA --- | |
| if prev_z_context is not None: | |
| # Łączymy poprzednie 2 sekundy z obecnym wyjściem | |
| z_pred_combined = torch.cat([prev_z_context, z_pred], dim=1) | |
| trim_frames = prev_z_context.shape[1] | |
| else: | |
| z_pred_combined = z_pred | |
| trim_frames = 0 | |
| wav, _, _ = codec.decode_from_continuous(z_pred_combined, spk_t, lang_t, f0=None) | |
| wav_np = wav.squeeze().cpu().float().numpy() | |
| if trim_frames > 0: | |
| # 1 ramka z_pred = 2 (stride) * 441 (hop) = 882 próbki audio | |
| trim_samples = trim_frames * 882 | |
| wav_np = wav_np[trim_samples:] | |
| return wav_np, z_pred | |
| # -------------------------------------------------------------------- | |
| # 7) Load models (ON CPU initially) | |
| # -------------------------------------------------------------------- | |
| startup_device = torch.device("cpu") | |
| print("Initializing models on CPU...") | |
| with open(PREDICTOR_CONFIG, "r") as f: | |
| pred_cfg = yaml.safe_load(f) | |
| with open(CODEC_CONFIG, "r") as f: | |
| codec_cfg = yaml.safe_load(f) | |
| predictor = build_predictor(pred_cfg, startup_device) | |
| predictor = load_predictor(predictor, PREDICTOR_CKPT) | |
| codec = build_codec(codec_cfg, startup_device) | |
| codec = load_codec(codec, CODEC_CKPT) | |
| for p in codec.parameters(): | |
| p.requires_grad_(False) | |
| print("Models initialized on CPU.") | |
| # -------------------------------------------------------------------- | |
| # 8) Gradio UI + API with @spaces.GPU | |
| # -------------------------------------------------------------------- | |
| LANG_OPTIONS = [ (v["display"], k) for k, v in LANGUAGES.items() ] | |
| LANG_KEYS = { k: v for k, v in LANGUAGES.items() } | |
| def speakers_for_lang(lang_key: str): | |
| return [str(x) for x in LANG_KEYS[lang_key]["speakers"]] | |
| def synth_api(text, lang_key, speaker_str, steps, cfg, length_scale, temp, dur_noise): | |
| t0 = time.time() | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| predictor.to(device) | |
| codec.to(device) | |
| if not text or not text.strip(): | |
| return (44100, np.zeros(1)), "No text provided." | |
| try: | |
| speaker_id = int(speaker_str) | |
| lang_id = LANG_KEYS[lang_key]["id"] | |
| except Exception as e: | |
| return (44100, np.zeros(1)), f"Invalid selection: {e}" | |
| # Split text into sentences for sequential processing | |
| sentences = split_sentences(text) | |
| if not sentences: | |
| return (44100, np.zeros(1)), "No valid sentences found to synthesize." | |
| sr = 44100 | |
| all_wavs = [] | |
| # Usunięto sztuczną ciszę | |
| # silence = np.zeros(int(sr * 0.15), dtype=np.float32) | |
| prev_z = None # Inicjalizacja bufora na kody latentne | |
| try: | |
| for i, sent in enumerate(sentences): | |
| prev_text = sentences[i - 1] if i > 0 else None | |
| next_text = sentences[i + 1] if i < len(sentences) - 1 else None | |
| # --- POBIERANIE 2 SEKUND KONTEKSTU (100 RAMEK @ 50Hz) --- | |
| prev_z_context = None | |
| if prev_z is not None: | |
| context_len = min(100, prev_z.shape[1]) | |
| prev_z_context = prev_z[:, -context_len:, :] | |
| wav, current_z = synthesize_sentence( | |
| text=sent, | |
| predictor=predictor, | |
| codec=codec, | |
| device=device, | |
| speaker_id=speaker_id, | |
| language_id=lang_id, | |
| prev_text=prev_text, | |
| next_text=next_text, | |
| n_steps=int(steps), | |
| cfg_strength=float(cfg), | |
| length_scale=float(length_scale), | |
| temperature=float(temp), | |
| dur_noise_scale=float(dur_noise), | |
| prev_z_context=prev_z_context, # Przekazujemy kontekst | |
| ) | |
| # Aktualizujemy prev_z na potrzeby kolejnego zdania | |
| prev_z = current_z | |
| all_wavs.append(wav) | |
| # Nie dodajemy sztucznej ciszy pomiędzy wygenerowanymi zdaniami | |
| # if i < len(sentences) - 1: | |
| # all_wavs.append(silence) | |
| full_wav = np.concatenate(all_wavs) | |
| except Exception as e: | |
| return (44100, np.zeros(1)), f"Synthesis error: {e}" | |
| dur = len(full_wav) / sr | |
| gen_time = time.time() - t0 | |
| rtf = gen_time / dur if dur > 0 else 0.0 | |
| stats = f"⏳ Gen time: {gen_time:.2f}s | 🎵 Duration: {dur:.2f}s | ⚡ RTF: {rtf:.3f} | 📝 Sentences: {len(sentences)}" | |
| return (sr, full_wav), stats | |
| with gr.Blocks(title="SlopTTS Inference") as app: | |
| # --- Custom Readme Header --- | |
| gr.Markdown(""" | |
| # 🎙️ SFlowTTS Inference Playground | |
| ### ℹ️ About Me & The Model | |
| Hi there! I am the developer behind SFlowTTS. This model is using a **44.1 kHz** FSQ vocoder combined with a discrete Flow-Matching architecture. This 2-stage architecture has ~200M parameters. I trained the entire model within 7 days on 2x H100 | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| lang_dd = gr.Dropdown(choices=[k for _, k in LANG_OPTIONS], value="english", label="Language") | |
| speaker_dd = gr.Dropdown(choices=speakers_for_lang("english"), value=speakers_for_lang("english")[0], label="Speaker ID") | |
| story_dd = gr.Dropdown(choices=list(STORIES["english"].keys()), value=list(STORIES["english"].keys())[0], label="Select a Story") | |
| text_in = gr.Textbox(label="Text to synthesize", value=list(STORIES["english"].values())[0], lines=6) | |
| with gr.Accordion("Advanced Settings", open=True): | |
| length_sl = gr.Slider(minimum=0.5, maximum=2.0, step=0.05, value=1.0, label="Multiply Duration (Length Scale)") | |
| temp_sl = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=0.8, label="Flow Matching Temperature") | |
| dur_noise_sl = gr.Slider(minimum=0.1, maximum=2.0, step=0.01, value=0.667, label="Duration Noise Scale") | |
| steps_sl = gr.Slider(minimum=2, maximum=200, step=1, value=64, label="Inference Steps") | |
| cfg_sl = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=4.0, label="CFG Strength") | |
| synth_btn = gr.Button("Synthesize", variant="primary") | |
| audio_out = gr.Audio(label="Generated Audio", type="numpy") | |
| stats_out = gr.Textbox(label="Generation St1ats", interactive=False) | |
| def on_lang_change(lang_key): | |
| speakers = speakers_for_lang(lang_key) | |
| story_choices = list(STORIES[lang_key].keys()) | |
| first_story_title = story_choices[0] | |
| first_story_text = STORIES[lang_key][first_story_title] | |
| return ( | |
| gr.Dropdown(choices=speakers, value=speakers[0]), | |
| gr.Dropdown(choices=story_choices, value=first_story_title), | |
| first_story_text | |
| ) | |
| def on_story_change(lang_key, story_title): | |
| return STORIES[lang_key][story_title] | |
| lang_dd.change(fn=on_lang_change, inputs=lang_dd, outputs=[speaker_dd, story_dd, text_in]) | |
| story_dd.change(fn=on_story_change, inputs=[lang_dd, story_dd], outputs=text_in) | |
| synth_btn.click( | |
| fn=synth_api, | |
| inputs=[text_in, lang_dd, speaker_dd, steps_sl, cfg_sl, length_sl, temp_sl, dur_noise_sl], | |
| outputs=[audio_out, stats_out] | |
| ) | |
| if __name__ == "__main__": | |
| app.launch() |