Adding deepspeed to xtts-streaming

xtts-streaming/config.yaml

@@ -1,14 +1,14 @@
+base_image:
+  image: htrivedi05/xtts-streaming
+  python_executable_path: /opt/conda/bin/python
 environment_variables:
   COQUI_TOS_AGREED: '1'
 external_package_dirs: []
 model_metadata: {}
 model_name: XTTS Streaming
-python_version: py310
-requirements_file: ./requirements.txt
 resources:
-  accelerator: T4
+  accelerator: H100
   cpu: '3'
   memory: 10Gi
   use_gpu: true
 secrets: {}
-system_packages: []

xtts-streaming/model/model.py

@@ -2,6 +2,7 @@
 import io
 import logging
 import os
+import time
 import wave
 
 import numpy as np
@@ -32,8 +33,12 @@ def load(self):
         config = XttsConfig()
         config.load_json(os.path.join(model_path, "config.json"))
         self.model = Xtts.init_from_config(config)
-        self.model.load_checkpoint(config, checkpoint_dir=model_path, eval=True)
+        # self.model.load_checkpoint(config, checkpoint_dir=model_path, eval=True)
+        self.model.load_checkpoint(
+            config, checkpoint_dir=model_path, eval=True, use_deepspeed=True
+        )
         self.model.to(device)
+        # self.compiled_model = torch.compile(self.model.inference_stream)
 
         self.speaker = {
             "speaker_embedding": self.model.speaker_manager.speakers[SPEAKER_NAME][
@@ -51,6 +56,17 @@ def load(self):
             .half()
             .tolist(),
         }
+
+        self.speaker_embedding = (
+            torch.tensor(self.speaker.get("speaker_embedding"))
+            .unsqueeze(0)
+            .unsqueeze(-1)
+        )
+        self.gpt_cond_latent = (
+            torch.tensor(self.speaker.get("gpt_cond_latent"))
+            .reshape((-1, 1024))
+            .unsqueeze(0)
+        )
         logging.info("🔥Model Loaded")
 
     def wav_postprocess(self, wav):
@@ -66,32 +82,21 @@ def predict(self, model_input):
         text = model_input.get("text")
         language = model_input.get("language", "en")
         chunk_size = int(
-            model_input.get("chunk_size", 150)
+            model_input.get("chunk_size", 20)
         )  # Ensure chunk_size is an integer
         add_wav_header = False
 
-        speaker_embedding = (
-            torch.tensor(self.speaker.get("speaker_embedding"))
-            .unsqueeze(0)
-            .unsqueeze(-1)
-        )
-        gpt_cond_latent = (
-            torch.tensor(self.speaker.get("gpt_cond_latent"))
-            .reshape((-1, 1024))
-            .unsqueeze(0)
-        )
-
         streamer = self.model.inference_stream(
             text,
             language,
-            gpt_cond_latent,
-            speaker_embedding,
+            self.gpt_cond_latent,
+            self.speaker_embedding,
             stream_chunk_size=chunk_size,
             enable_text_splitting=True,
+            temperature=0.2,
         )
 
         for chunk in streamer:
-            print(type(chunk))
             processed_chunk = self.wav_postprocess(chunk)
             processed_bytes = processed_chunk.tobytes()
             yield processed_bytes