New verbalizer, and better text processing, also upgraded styletts2-inference library with fix for short sentances generation.

app.py

@@ -17,16 +17,19 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 prompts_dir = 'voices'
 
-verbalizer = Verbalizer(device=device)
+verbalizer = Verbalizer()
 
 
 def split_to_parts(text):
+    text = re.sub(r'([^.,!:?\-])\n', r'\1. ', text)
+    text = text.replace('\n', ' ')
     split_symbols = '.?!:'
     parts = ['']
     index = 0
-    for s in text:
+    last = len(text)-1
+    for i, s in enumerate(text):
         parts[index] += s
-        if s in split_symbols and len(parts[index]) > 20:
+        if s in split_symbols and i < last and text[i+1] == ' ' and len(parts[index]) > 20:
             index += 1
             parts.append('')
     return parts
@@ -60,13 +63,12 @@ models = {
 }
 
 
-@spaces.GPU
 def verbalize(text):
     parts = split_to_parts(text)
     verbalized = ''
     for part in parts:
         if part.strip():
-            verbalized += verbalizer.generate_text(part) + ' '
+            verbalized += verbalizer.process_text(part)[0] + ' '
     return verbalized
 
 description = f'''
@@ -103,6 +105,7 @@ def synthesize(model_name, text, speed, voice_name = None, progress=gr.Progress(
 
     result_wav = []
     for t in progress.tqdm(split_to_parts(text)):
+        
 
         t = t.strip()
         t = t.replace('"', '')
@@ -111,8 +114,11 @@ def synthesize(model_name, text, speed, voice_name = None, progress=gr.Progress(
             t = normalize('NFKC', t)
 
             t = re.sub(r'[᠆‐‑‒–—―⁻₋−⸺⸻]', '-', t)
+            if t[-1] not in '.?!:-':
+                t += '.'
             t = re.sub(r' - ', ': ', t)
-            ps = ipa(stressify(t))
+            t = stressify(t)
+            ps = ipa(t)
 
             if ps:
                 tokens = models[model_name]['model'].tokenizer.encode(ps)

requirements.txt

@@ -10,7 +10,9 @@ torch==2.2.0
 transformers==5.0.0
 git+https://github.com/patriotyk/ukrainian-word-stress.git
 git+https://github.com/patriotyk/ipa-uk.git
-git+https://github.com/patriotyk/styletts2-inference@105aed29fa1a7698d08d920986890e9bbd03447c
+git+https://github.com/patriotyk/styletts2-inference@209352700cf703a05ff94501290095dd9072fd5f
 spaces
 numpy<2
-accelerate>=0.26.0
+accelerate>=0.26.0
+ctranslate2
+sentencepiece

verbalizer.py

@@ -1,39 +1,55 @@
-
-from transformers import MBartForConditionalGeneration, AutoTokenizer
-
-verbalizer_model_name = "skypro1111/mbart-large-50-verbalization"
+import os
+import time
+import ctranslate2
+from transformers import M2M100Tokenizer
+from huggingface_hub import snapshot_download
 
 
 class Verbalizer():
-    def __init__(self, device):
-        self.device = device
+    MODEL_PATH = os.getenv("MODEL_PATH", "skypro1111/m2m100-ukr-verbalization-ct2")
+    TOKENIZER_PATH = os.getenv("TOKENIZER_PATH", "skypro1111/m2m100-ukr-verbalization")
+    def __init__(self):
+        
+        print("\nInitializing CTranslate2 model and tokenizer...")
 
-        self.model = MBartForConditionalGeneration.from_pretrained(verbalizer_model_name,
-            low_cpu_mem_usage=True,
-            device_map=device,
-        )
-        self.model.eval()
-    
-        self.tokenizer = AutoTokenizer.from_pretrained(verbalizer_model_name)
-        self.tokenizer.src_lang = "uk_XX"
-        self.tokenizer.tgt_lang = "uk_XX"
     
-    def generate_text(self, text):
-        """Generate text for a single input."""
-        # Prepare input
-        input_text = "<verbalization>:" + text
-
-        encoded_input = self.tokenizer(
-            input_text,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=1024,
-        ).to(self.device)
-        output_ids = self.model.generate(
-            **encoded_input, max_length=1024, num_beams=5, early_stopping=True
+        # Download the model from HuggingFace Hub
+        local_model_path = snapshot_download(
+            repo_id=self.MODEL_PATH,
+            allow_patterns=["*.bin", "*.json", "tokenizer.json", "vocab.json"],
         )
-        normalized_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
+        
+        self.translator = ctranslate2.Translator(
+            local_model_path,
+            device='cpu',
+            compute_type="int8",  # Use INT8 quantization for CPU
+        )
+        
+        # Load tokenizer
+        self.tokenizer = M2M100Tokenizer.from_pretrained(self.TOKENIZER_PATH)
+        self.tokenizer.src_lang = "uk"
+        
+    def process_text(self, text: str):
+        """Process a single text input using the CTranslate2 model."""
+        start_time = time.time()
+        
+        # Tokenize input
+        source = self.tokenizer.convert_ids_to_tokens(self.tokenizer.encode(text))
+        target_prefix = [self.tokenizer.lang_code_to_token["uk"]]
+        
+        # Run inference
+        results = self.translator.translate_batch(
+            [source],
+            target_prefix=[target_prefix],
+            beam_size=1,
+            num_hypotheses=1,
+            use_vmap=True,
+        )
+        
+        # Get target tokens and decode
+        target = results[0].hypotheses[0][1:]  # Remove language token
+        output = self.tokenizer.decode(self.tokenizer.convert_tokens_to_ids(target))
+        
+        inference_time = time.time() - start_time
+        return output, inference_time
 
-        return normalized_text.strip()