app.py

import glob
import os
import re
import gradio as gr
import torch
import numpy as np
import librosa
import random
import spaces
from verbalizer import Verbalizer
from ipa_uk import ipa
from unicodedata import normalize
from styletts2_inference.models import StyleTTS2
from ukrainian_word_stress import Stressifier, StressSymbol

# --- Ініціалізація ---
stressify = Stressifier()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
prompts_dir = 'voices'
verbalizer = Verbalizer(device=device)

# Завантаження моделі
multi_model = StyleTTS2(hf_path='patriotyk/styletts2_ukrainian_multispeaker', device=device)

# Завантаження пресетів
multi_styles = {}
if not os.path.exists(prompts_dir): os.makedirs(prompts_dir)
prompts_list = sorted(glob.glob(os.path.join(prompts_dir, '*.pt')))
prompts_names = ['.'.join(p.split('/')[-1].split('.')[:-1]) for p in prompts_list]
for name, path in zip(prompts_names, prompts_list):
    multi_styles[name] = torch.load(path)

# --- Допоміжні функції ---

def split_to_parts(text):
    split_symbols = '.?!:'
    parts = ['']
    index = 0
    for s in text:
        parts[index] += s
        if s in split_symbols and len(parts[index]) > 150:
            index += 1
            parts.append('')
    return [p.strip() for p in parts if p.strip()]

@spaces.GPU
def synthesize(text, speed, voice_mode, voice_name, voice_audio, noise_scale, match_duration, progress=gr.Progress()):
    if text.strip() == "": raise gr.Error("Введіть текст")

    # 1. Отримуємо стиль з одного джерела
    if voice_mode == "Пресет":
        style = multi_styles[voice_name]
        target_duration = None
    else:
        if not voice_audio: raise gr.Error("Завантажте аудіо-референс")
        # Витягуємо особливості (тембр + емоцію) з файлу
        style = multi_model.extract_voice_features(voice_audio)
        if isinstance(style, list): style = style[-1]
        style = style.to(device)
        
        y, _ = librosa.load(voice_audio, sr=24000)
        target_duration = librosa.get_duration(y=y, sr=24000)

    parts = split_to_parts(text)
    final_speed = speed
    
    # 2. Smart Duration (підлаштування під довжину референсу)
    if match_duration and target_duration:
        progress(0, desc="Аналіз тривалості...")
        temp_wavs = []
        for t in parts:
            ps = ipa(stressify(normalize('NFKC', t.replace('+', StressSymbol.CombiningAcuteAccent))))
            if ps:
                tokens = multi_model.tokenizer.encode(ps)
                w = multi_model(tokens, speed=1.0, s_prev=style)
                temp_wavs.append(w)
        gen_len = sum(len(w) for w in temp_wavs) / 24000
        calc_speed = gen_len / target_duration
        if 0.6 <= calc_speed <= 1.4:
            final_speed = calc_speed

    # 3. Синтез з Noise Injection для унікальності дублів
    result_wav = []
    for t in progress.tqdm(parts, desc="Генерація дубля"):
        t = normalize('NFKC', t.replace('+', StressSymbol.CombiningAcuteAccent))
        ps = ipa(stressify(t))
        
        if ps:
            tokens = multi_model.tokenizer.encode(ps)
            
            # Кожен раз створюємо мікро-варіацію стилю
            current_style = style.clone()
            if noise_scale > 0:
                noise = torch.randn_like(current_style) * noise_scale
                current_style += noise
            
            # Виклик моделі (forward)
            wav = multi_model(tokens, speed=final_speed, s_prev=current_style)
            result_wav.append(wav.cpu().numpy().flatten())

    if not result_wav: return None
    return 24000, np.concatenate(result_wav)

# --- UI ---
with gr.Blocks(title="StyleTTS2 UA extended") as demo:
    gr.Markdown("# 🎬 StyleTTS2 UA: Розширені налаштування")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label='Текст:', lines=5, placeholder="Введіть репліку...")
            
            mode = gr.Radio(["Пресет", "Власний референс"], value="Пресет", label="Голос та емоція:")
            speaker = gr.Dropdown(choices=prompts_names, value=prompts_names[0] if prompts_names else None, label="Оберіть пресет:", visible=True)
            audio_src = gr.Audio(label="Завантажте зразок (тембр + емоція):", type="filepath", visible=False)
            
            mode.change(lambda m: [gr.update(visible=m=="Пресет"), gr.update(visible=m=="Власний референс")], 
                        inputs=[mode], outputs=[speaker, audio_src])
            
            with gr.Row():
                speed_slider = gr.Slider(label='Швидкість', maximum=1.3, minimum=0.7, value=1.0, step=0.05)
                noise_slider = gr.Slider(label='Варіативність', minimum=0.0, maximum=0.2, value=0.1, step=0.01)
            
            smart_dur = gr.Checkbox(label="Зробити тривалість як в референсі", value=False)
            btn = gr.Button("🚀 ЗГЕНЕРУВАТИ", variant="primary")

        with gr.Column():
            output_audio = gr.Audio(label="Результат:", type="numpy")

    btn.click(synthesize, 
              inputs=[input_text, speed_slider, mode, speaker, audio_src, noise_slider, smart_dur], 
              outputs=[output_audio])

if __name__ == "__main__":
    demo.queue().launch()