gpu based

app.py

@@ -54,15 +54,17 @@ import torchvision.transforms as transforms
 import torch.nn.functional as F
 from OmniAvatar.utils.audio_preprocess import add_silence_to_audio_ffmpeg
 
-from supertonic import generate_speech
+from supertonic import generate_speech, load_text_to_speech
 
 os.environ["PROCESSED_RESULTS"] = f"{os.getcwd()}/proprocess_results"
 
 def tts_from_text(text, tts_dir, voice_choice):
 
-    output = generate_speech([text], tts_dir, voice_choice)[0]
+    text_to_speech = load_text_to_speech(True)
+    output = generate_speech(text_to_speech , [text], tts_dir, voice_choice)[0]
     return output
 
+@spaces.GPU()
 def speak_to_me(session_id, evt: gr.EventData):
     detail = getattr(evt, "data", None) or getattr(evt, "_data", {}) or {}
 

requirements.txt

@@ -17,4 +17,4 @@ gradio_extendedimage @ https://github.com/OutofAi/gradio-extendedimage/releases/
 gradio_extendedaudio @ https://github.com/OutofAi/gradio-extendedaudio/releases/download/0.0.5/gradio_extendedaudio-0.0.5-py3-none-any.whl
 
 flash-attn-3 @ https://huggingface.co/alexnasa/flash-attn-3/resolve/main/128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl
-onnxruntime
+onnxruntime-gpu

supertonic.py

@@ -1,14 +1,15 @@
 import json
 import os
 import time
-from contextlib import contextmanager
+from time_util import timer
 from typing import Optional
 from unicodedata import normalize
-
+import uuid
 import numpy as np
 import onnxruntime as ort
 import soundfile as sf
 from huggingface_hub import snapshot_download
+from typing import Optional, Union
 
 
 class UnicodeProcessor:
@@ -87,24 +88,65 @@ class TextToSpeech:
         noisy_latent = noisy_latent * latent_mask
         return noisy_latent, latent_mask
 
+
     def _infer(
-        self, text_list: list[str], style: Style, total_step: int, speed: float = 1.05
+        self,
+        text_list: list[str],
+        style: Style,
+        total_step: int,
+        speed: float = 1.05,
+        suggested_duration: Optional[Union[float, list[float], np.ndarray]] = None,
+        speed_min_factor: float = 0.75,
+        speed_max_factor: float = 1.2,
     ) -> tuple[np.ndarray, np.ndarray]:
         assert (
             len(text_list) == style.ttl.shape[0]
         ), "Number of texts must match number of style vectors"
         bsz = len(text_list)
+
         text_ids, text_mask = self.text_processor(text_list)
-        dur_onnx, *_ = self.dp_ort.run(
+
+        # 1) Predict base duration
+        dur_pred, *_ = self.dp_ort.run(
             None, {"text_ids": text_ids, "style_dp": style.dp, "text_mask": text_mask}
         )
-        dur_onnx = dur_onnx / speed
+        dur_pred = np.array(dur_pred, dtype=np.float32).reshape(bsz)  # (bsz,)
+
+        # 2) Adjust duration based on suggested_duration (if given)
+        if suggested_duration is not None:
+            sugg = np.array(suggested_duration, dtype=np.float32)
+            if sugg.ndim == 0:
+                # same suggestion for all
+                sugg = np.full((bsz,), float(sugg), dtype=np.float32)
+            else:
+                sugg = sugg.reshape(bsz)
+
+            eps = 1e-3
+            sugg = np.clip(sugg, eps, None)
+
+            # we want dur_used ≈ sugg
+            # dur_used = dur_pred / speed_used  => speed_target = dur_pred / sugg
+            speed_target = dur_pred / sugg
+
+            speed_min = speed * speed_min_factor
+            speed_max = speed * speed_max_factor
+            speed_used = np.clip(speed_target, speed_min, speed_max)
+
+            dur_used = dur_pred / speed_used
+        else:
+            # default behaviour
+            speed_used = np.full((bsz,), speed, dtype=np.float32)
+            dur_used = dur_pred / speed_used
+
+        # 3) Continue as before, using dur_used
         text_emb_onnx, *_ = self.text_enc_ort.run(
             None,
             {"text_ids": text_ids, "style_ttl": style.ttl, "text_mask": text_mask},
-        )  # dur_onnx: [bsz]
-        xt, latent_mask = self.sample_noisy_latent(dur_onnx)
+        )
+
+        xt, latent_mask = self.sample_noisy_latent(dur_used)
         total_step_np = np.array([total_step] * bsz, dtype=np.float32)
+
         for step in range(total_step):
             current_step = np.array([step] * bsz, dtype=np.float32)
             xt, *_ = self.vector_est_ort.run(
@@ -119,8 +161,30 @@ class TextToSpeech:
                     "total_step": total_step_np,
                 },
             )
+
         wav, *_ = self.vocoder_ort.run(None, {"latent": xt})
-        return wav, dur_onnx
+        return wav, dur_used
+
+    def batch(
+        self,
+        text_list: list[str],
+        style: Style,
+        total_step: int,
+        speed: float = 1.05,
+        suggested_duration: Optional[Union[float, list[float], np.ndarray]] = None,
+        speed_min_factor: float = 0.75,
+        speed_max_factor: float = 1.2,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        return self._infer(
+            text_list,
+            style,
+            total_step,
+            speed=speed,
+            suggested_duration=suggested_duration,
+            speed_min_factor=speed_min_factor,
+            speed_max_factor=speed_max_factor,
+        )
+
 
     def __call__(
         self,
@@ -149,10 +213,6 @@ class TextToSpeech:
                 dur_cat += dur_onnx + silence_duration
         return wav_cat, dur_cat
 
-    def batch(
-        self, text_list: list[str], style: Style, total_step: int, speed: float = 1.05
-    ) -> tuple[np.ndarray, np.ndarray]:
-        return self._infer(text_list, style, total_step, speed)
 
 
 def length_to_mask(lengths: np.ndarray, max_len: Optional[int] = None) -> np.ndarray:
@@ -219,11 +279,17 @@ def load_text_processor(onnx_dir: str) -> UnicodeProcessor:
     text_processor = UnicodeProcessor(unicode_indexer_path)
     return text_processor
 
+# text_to_speech = load_text_to_speech(False)
+
+
+model_dir = snapshot_download("Supertone/supertonic")
+onnx_dir = f"{model_dir}/onnx"
+
+def load_text_to_speech(use_gpu: bool = False) -> TextToSpeech:
 
-def load_text_to_speech(onnx_dir: str, use_gpu: bool = False) -> TextToSpeech:
     opts = ort.SessionOptions()
     if use_gpu:
-        raise NotImplementedError("GPU mode is not fully tested")
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
     else:
         providers = ["CPUExecutionProvider"]
         print("Using CPU for inference")
@@ -268,12 +334,6 @@ def load_voice_style(voice_style_paths: list[str], verbose: bool = False) -> Sty
     return Style(ttl_style, dp_style)
 
 
-@contextmanager
-def timer(name: str):
-    start = time.time()
-    print(f"{name}...")
-    yield
-    print(f"  -> {name} completed in {time.time() - start:.2f} sec")
 
 
 def sanitize_filename(text: str, max_len: int) -> str:
@@ -327,11 +387,19 @@ def chunk_text(text: str, max_len: int = 300) -> list[str]:
 
     return chunks
 
-model_dir = snapshot_download("Supertone/supertonic")
-onnx_dir = f"{model_dir}/onnx"
-text_to_speech = load_text_to_speech(onnx_dir, False)
-
-def generate_speech(text_list, save_dir, voice_style="M1", total_step=5, speed=1.05, n_test=1, batch=None):
+def generate_speech(
+    text_to_speech,
+    text_list,
+    save_dir,
+    voice_style="M1",
+    total_step=5,
+    speed=1.05,
+    n_test=1,
+    batch=None,
+    suggested_durations=None,      # NEW: list/np.ndarray of seconds, len == len(text_list)
+    speed_min_factor=0.75,
+    speed_max_factor=1.2,
+):
 
     saved_files_list = []
 
@@ -345,20 +413,30 @@ def generate_speech(text_list, save_dir, voice_style="M1", total_step=5, speed=1
     style = load_voice_style(voice_style_paths, verbose=True)
 
     for n in range(n_test):
-        print(f"\n[{n+1}/{n_test}] Starting synthesis...")
-        with timer("Generating speech from text"):
-            if batch:
-                wav, duration = text_to_speech.batch(text_list, style, total_step, speed)
-            else:
-                wav, duration = text_to_speech(text_list[0], style, total_step, speed)
+        if batch:
+            wav, duration = text_to_speech.batch(
+                text_list,
+                style,
+                total_step,
+                speed=speed,
+                suggested_duration=suggested_durations,
+                speed_min_factor=speed_min_factor,
+                speed_max_factor=speed_max_factor,
+            )
+        else:
+            # optional: could support suggested_durations[0] here too
+            wav, duration = text_to_speech(
+                text_list[0], style, total_step, speed
+            )
+
         if not os.path.exists(save_dir):
             os.makedirs(save_dir)
+
         for b in range(bsz):
-            fname = f"{sanitize_filename(text_list[b], 20)}_{n+1}.wav"
-            w = wav[b, : int(text_to_speech.sample_rate * duration[b].item())]  # [T_trim]
+            unique = uuid.uuid4().hex[:8]
+            fname = f"{sanitize_filename(text_list[b], 20)}_{unique}_{n+1}.wav"
+            w = wav[b, : int(text_to_speech.sample_rate * duration[b].item())]
             sf.write(os.path.join(save_dir, fname), w, text_to_speech.sample_rate)
             saved_files_list.append(f"{save_dir}/{fname}")
-            # print(f"Saved: {save_dir}/{fname}")
-    print("\n=== Synthesis completed successfully! ===")
 
     return saved_files_list

time_util.py

@@ -0,0 +1,9 @@
+import time
+from contextlib import contextmanager
+
+@contextmanager
+def timer(name: str):
+    start = time.time()
+    print(f"{name}...")
+    yield
+    print(f"  -> {name} completed in {time.time() - start:.2f} sec")