Update app.py

app.py

@@ -1,6 +1,7 @@
 # Imports
 import gradio as gr
 import spaces
+import torch
 
 from kokoro import KModel, KPipeline
 
@@ -10,6 +11,8 @@ if DEVICE == "auto":
     DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"[SYSTEM] | Using {DEVICE} type compute device.")
 
+torch.set_num_threads(4)
+
 # Variables
 CHAR_LIMIT = 2000
 
@@ -51,10 +54,12 @@ PIPELINES = {lang: KPipeline(lang_code=lang, model=False) for lang in "ab"}
 PIPELINES["a"].g2p.lexicon.golds["kokoro"] = "kˈOkəɹO"
 PIPELINES["b"].g2p.lexicon.golds["kokoro"] = "kˈQkəɹQ"
 
+VOICE_PACKS = {}
 for v in CHOICES.values():
-    PIPELINES[v[0]].load_voice(v)
-
-MODEL = KModel().eval()
+    VOICE_PACKS[v] = PIPELINES[v[0]].load_voice(v)
+    
+model_instance = KModel().to(DEVICE).eval()
+MODEL = torch.jit.script(model_instance)
 
 css = '''
 .gradio-container{max-width: 560px !important}
@@ -66,10 +71,10 @@ footer {
 
 # Functions
 def generate(text=DEFAULT_INPUT, voice=DEFAULT_VOICE, speed=1):
-    text = text.strip()[:CHAR_LIMIT] + "."
+text = text.strip()[:CHAR_LIMIT] + "."
     pipeline = PIPELINES[voice[0]]
-    pack = pipeline.load_voice(voice)
-    
+    pack = VOICE_PACKS[voice]
+
     for _, ps, _ in pipeline(text, voice, speed):
         ref_s = pack[len(ps) - 1]
         audio = MODEL(ps, ref_s, speed)