Update app.py

app.py

@@ -7,12 +7,16 @@ import re
 
 from parler_tts import ParlerTTSForConditionalGeneration
 from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
+from functools import lru_cache
+from torch.cuda.amp import autocast
+import time
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 repo_id = "parler-tts/parler-tts-mini-expresso"
 
-model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id).to(device)
+model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id, torch_dtype=torch.float16).to(device)
+model = torch.compile(model)  # Adiciona otimização com torch.compile
 tokenizer = AutoTokenizer.from_pretrained(repo_id)
 feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)
 
@@ -66,15 +70,30 @@ def preprocess(text):
     return text
 
 
-@spaces.GPU
-def gen_tts(text, description):
-    inputs = tokenizer(description, return_tensors="pt").to(device)
-    prompt = tokenizer(preprocess(text), return_tensors="pt").to(device)
+@lru_cache(maxsize=128)
+def cached_tokenizer(text):
+    return tokenizer(text, return_tensors="pt").to(device)
 
-    set_seed(SEED)
-    generation = model.generate(input_ids=inputs.input_ids, prompt_input_ids=prompt.input_ids)
-    audio_arr = generation.cpu().numpy().squeeze()
 
+@spaces.GPU
+def gen_tts(text, description):
+    start_time = time.time()
+    with torch.no_grad():  # Desativa gradientes
+        inputs = cached_tokenizer(description)
+        prompt = cached_tokenizer(preprocess(text))
+
+        set_seed(SEED)
+        with autocast():  # Habilita precisão mista
+            generation = model.generate(
+                input_ids=inputs.input_ids,
+                prompt_input_ids=prompt.input_ids,
+                max_length=200,  # Limita o comprimento máximo da saída
+                num_beams=3      # Usa beam search com 3 feixes
+            )
+        audio_arr = generation.cpu().numpy().squeeze()
+
+    end_time = time.time()
+    print(f"Generation completed in {end_time - start_time:.2f} seconds")
     return SAMPLE_RATE, audio_arr