Update app.py

app.py

@@ -1,92 +1,109 @@
 import spaces
 import gradio as gr
 import torch
-from huggingface_hub import hf_hub_download 
+from huggingface_hub import hf_hub_download
 import os
 import sys
 import tempfile
+from pathlib import Path
 from scipy.io.wavfile import write
-
-# Клонуем рэпазіторый, калі ён яшчэ не загружаны
-if not os.path.exists("XTTSv2-Finetuning-for-New-Languages"):
-    os.system("git clone https://github.com/hellcatmon/XTTSv2-Finetuning-for-New-Languages.git")
-
-# Перамяшчаем тэчку TTS у асноўную дырэкторыю
-if os.path.exists("XTTSv2-Finetuning-for-New-Languages/TTS"):
-    os.system("mv XTTSv2-Finetuning-for-New-Languages/TTS ./")
-
-# Дадаем тэчку TTS у PYTHONPATH
-sys.path.append("./TTS")
+import numpy as np
 from tqdm import tqdm
 from underthesea import sent_tokenize
-from TTS.tts.configs.xtts_config import XttsConfig
-from TTS.tts.models.xtts import Xtts
 
-# Шлях да мадэлі ў Hugging Face
-repo_id = "archivartaunik/BE_XTTS_V2_60epoch3Dataset"
-checkpoint_file = hf_hub_download(repo_id, filename="model.pth")
-config_file = hf_hub_download(repo_id, filename="config.json")
-vocab_file = hf_hub_download(repo_id, filename="vocab.json")
-default_voice_file = hf_hub_download(repo_id, filename="voice.wav")
+try:
+    from TTS.tts.configs.xtts_config import XttsConfig
+    from TTS.tts.models.xtts import Xtts
+except ImportError:
+    os.system("git clone https://github.com/hellcatmon/XTTSv2-Finetuning-for-New-Languages.git")
+    if os.path.exists("XTTSv2-Finetuning-for-New-Languages/TTS"):
+        os.system("mv XTTSv2-Finetuning-for-New-Languages/TTS ./")
+    sys.path.append("./TTS")
+    from TTS.tts.configs.xtts_config import XttsConfig
+    from TTS.tts.models.xtts import Xtts
 
-# Загрузка канфігурацыі мадэлі
+# Шляхі да файлаў
+repo_id = "archivartaunik/BE_XTTS_V2_60epoch3Dataset"
+model_path = Path("./model")
+model_path.mkdir(exist_ok=True)
+checkpoint_file = model_path / "model.pth"
+config_file = model_path / "config.json"
+vocab_file = model_path / "vocab.json"
+default_voice_file = model_path / "voice.wav"
+
+if not checkpoint_file.exists():
+    hf_hub_download(repo_id, filename="model.pth", local_dir=model_path)
+if not config_file.exists():
+    hf_hub_download(repo_id, filename="config.json", local_dir=model_path)
+if not vocab_file.exists():
+    hf_hub_download(repo_id, filename="vocab.json", local_dir=model_path)
+if not default_voice_file.exists():
+    hf_hub_download(repo_id, filename="voice.wav", local_dir=model_path)
+
+# Загрузка канфігурацыі і мадэлі адзін раз
 config = XttsConfig()
 config.load_json(config_file)
-
-# Ініцыялізацыя і загрузка мадэлі
 XTTS_MODEL = Xtts.init_from_config(config)
 XTTS_MODEL.load_checkpoint(config, checkpoint_path=checkpoint_file, vocab_path=vocab_file, use_deepspeed=False)
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+XTTS_MODEL.to(device)
+sampling_rate = XTTS_MODEL.config.audio["sample_rate"]
 
 @spaces.GPU(duration=60)
-def text_to_speech(belarusian_story, lang="be", speaker_audio_file=None):
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    XTTS_MODEL.to(device)
-    # Калі файл не пададзены, выкарыстоўваем голас па змаўчанні
+def text_to_speech(belarusian_story, speaker_audio_file=None): # Прыбралі аргумент lang
     if not speaker_audio_file or (not isinstance(speaker_audio_file, str) and speaker_audio_file.name == ""):
-        speaker_audio_file = default_voice_file
-
-    # Атрыманне латэнтных умоў і эмацый
-    gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(
-        audio_path=speaker_audio_file,
-        gpt_cond_len=XTTS_MODEL.config.gpt_cond_len,
-        max_ref_length=XTTS_MODEL.config.max_ref_len,
-        sound_norm_refs=XTTS_MODEL.config.sound_norm_refs,
-    )
+        speaker_audio_file = str(default_voice_file)
+
+    try:
+        gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(
+            audio_path=speaker_audio_file,
+            gpt_cond_len=XTTS_MODEL.config.gpt_cond_len,
+            max_ref_length=XTTS_MODEL.config.max_ref_len,
+            sound_norm_refs=XTTS_MODEL.config.sound_norm_refs,
+        )
+    except Exception as e:
+        return f"Error getting conditioning latents: {e}"
 
-    # Токенізацыя тэксту на асобныя сказы
-    tts_texts = sent_tokenize(belarusian_story)
+    try:
+        tts_texts = sent_tokenize(belarusian_story)
+    except Exception as e:
+        return f"Error tokenizing text: {e}"
 
-    # Генерацыя аўдыё для кожнага сказы
-    wav_chunks = []
+    all_wavs = []
     for text in tqdm(tts_texts):
-        wav_chunk = XTTS_MODEL.inference(
-            text=text,
-            language=lang,
-            gpt_cond_latent=gpt_cond_latent,
-            speaker_embedding=speaker_embedding,
-            temperature=0.1,
-            length_penalty=1.0,
-            repetition_penalty=10.0,
-            top_k=10,
-            top_p=0.3,
-        )
-        wav_chunks.append(torch.tensor(wav_chunk["wav"]))
-
-    # Аб'ядноўваем усе часткі аўдыё ў адзін масіў
-    out_wav = torch.cat(wav_chunks, dim=0).squeeze().cpu().numpy()
+        try:
+            with torch.no_grad():
+                wav_chunk = XTTS_MODEL.inference(
+                    text=text,
+                    language="be", # Зафіксавалі мову як "be"
+                    gpt_cond_latent=gpt_cond_latent,
+                    speaker_embedding=speaker_embedding,
+                    temperature=0.1,
+                    length_penalty=1.0,
+                    repetition_penalty=10.0,
+                    top_k=10,
+                    top_p=0.3,
+                )
+            all_wavs.append(wav_chunk["wav"])
+        except Exception as e:
+            return f"Error generating audio: {e}"
+
+    try:
+        out_wav = np.concatenate(all_wavs)
+    except ValueError:
+        return "Немагчыма згенерыраваць аўдыё. Праверце ўваходны тэкст і аўдыёфайл."
+    except Exception as e:
+        return f"Error concatenating audio: {e}"
 
-    # Захоўваем аўдыё ў часовы файл
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
-    write(temp_file.name, 24000, out_wav)
+    write(temp_file.name, sampling_rate, out_wav)
 
     return temp_file.name
 
-
 demo = gr.Interface(
     fn=text_to_speech,
     inputs=[
         gr.Textbox(lines=5, label="Тэкст на беларускай мове"),
-        gr.Textbox(value="be", label="Мова (па змаўчанні BE)", visible=False),
         gr.Audio(type="filepath", label="Запішыце або загрузіце файл голасу (без іншых гукаў) не карацей 7 секунд", interactive=True),
     ],
     outputs="audio",