support voice creation

cli/SparkTTS.py

@@ -15,12 +15,13 @@
 
 import re
 import torch
+from typing import Tuple
 from pathlib import Path
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 from sparktts.utils.file import load_config
 from sparktts.models.audio_tokenizer import BiCodecTokenizer
-from sparktts.utils.token_parser import TASK_TOKEN_MAP
+from sparktts.utils.token_parser import LEVELS_MAP, GENDER_MAP, TASK_TOKEN_MAP
 
 
 class SparkTTS:
@@ -49,36 +50,36 @@ class SparkTTS:
         self.audio_tokenizer = BiCodecTokenizer(self.model_dir, device=self.device)
         self.model.to(self.device)
 
-    @torch.no_grad()
-    def inference(
+    def process_prompt(
         self,
         text: str,
         prompt_speech_path: Path,
         prompt_text: str = None,
-        temperature: float = 0.8,
-        top_k: float = 50,
-        top_p: float = 0.95,
-    ) -> torch.Tensor:
+    ) -> Tuple[str, torch.Tensor]:
         """
-        Performs inference to generate speech from text, incorporating prompt audio and/or text.
+        Process input for voice cloning.
 
         Args:
             text (str): The text input to be converted to speech.
             prompt_speech_path (Path): Path to the audio file used as a prompt.
             prompt_text (str, optional): Transcript of the prompt audio.
-            temperature (float, optional): Sampling temperature for controlling randomness. Default is 0.8.
-            top_k (float, optional): Top-k sampling parameter. Default is 50.
-            top_p (float, optional): Top-p (nucleus) sampling parameter. Default is 0.95.
 
-        Returns:
-            torch.Tensor: Generated waveform as a tensor.
+        Return:
+            Tuple[str, torch.Tensor]: Input prompt; global tokens
         """
-        global_token_ids, semantic_token_ids = self.audio_tokenizer.tokenize(prompt_speech_path)
-        global_tokens = "".join([f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze()])
+
+        global_token_ids, semantic_token_ids = self.audio_tokenizer.tokenize(
+            prompt_speech_path
+        )
+        global_tokens = "".join(
+            [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze()]
+        )
 
         # Prepare the input tokens for the model
         if prompt_text is not None:
-            semantic_tokens = "".join([f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze()])
+            semantic_tokens = "".join(
+                [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze()]
+            )
             inputs = [
                 TASK_TOKEN_MAP["tts"],
                 "<|start_content|>",
@@ -103,7 +104,94 @@ class SparkTTS:
             ]
 
         inputs = "".join(inputs)
-        model_inputs = self.tokenizer([inputs], return_tensors="pt").to(self.device)
+
+        return inputs, global_token_ids
+
+    def process_prompt_control(
+        self,
+        gender: str,
+        pitch: str,
+        speed: str,
+        text: str,
+    ):
+        """
+        Process input for voice creation.
+
+        Args:
+            gender (str): female | male.
+            pitch (str): very_low | low | moderate | high | very_high
+            speed (str): very_low | low | moderate | high | very_high
+            text (str): The text input to be converted to speech.
+
+        Return:
+            str: Input prompt
+        """
+        assert gender in GENDER_MAP.keys()
+        assert pitch in LEVELS_MAP.keys()
+        assert speed in LEVELS_MAP.keys()
+
+        gender_id = GENDER_MAP[gender]
+        pitch_level_id = LEVELS_MAP[pitch]
+        speed_level_id = LEVELS_MAP[speed]
+
+        pitch_label_tokens = f"<|pitch_label_{pitch_level_id}|>"
+        speed_label_tokens = f"<|speed_label_{speed_level_id}|>"
+        gender_tokens = f"<|gender_{gender_id}|>"
+
+        attribte_tokens = "".join(
+            [gender_tokens, pitch_label_tokens, speed_label_tokens]
+        )
+
+        control_tts_inputs = [
+            TASK_TOKEN_MAP["controllable_tts"],
+            "<|start_content|>",
+            text,
+            "<|end_content|>",
+            "<|start_style_label|>",
+            attribte_tokens,
+            "<|end_style_label|>",
+        ]
+
+        return "".join(control_tts_inputs)
+
+    @torch.no_grad()
+    def inference(
+        self,
+        text: str,
+        prompt_speech_path: Path = None,
+        prompt_text: str = None,
+        gender: str = None,
+        pitch: str = None,
+        speed: str = None,
+        temperature: float = 0.8,
+        top_k: float = 50,
+        top_p: float = 0.95,
+    ) -> torch.Tensor:
+        """
+        Performs inference to generate speech from text, incorporating prompt audio and/or text.
+
+        Args:
+            text (str): The text input to be converted to speech.
+            prompt_speech_path (Path): Path to the audio file used as a prompt.
+            prompt_text (str, optional): Transcript of the prompt audio.
+            gender (str): female | male.
+            pitch (str): very_low | low | moderate | high | very_high
+            speed (str): very_low | low | moderate | high | very_high
+            temperature (float, optional): Sampling temperature for controlling randomness. Default is 0.8.
+            top_k (float, optional): Top-k sampling parameter. Default is 50.
+            top_p (float, optional): Top-p (nucleus) sampling parameter. Default is 0.95.
+
+        Returns:
+            torch.Tensor: Generated waveform as a tensor.
+        """
+        if gender is not None:
+            prompt = self.process_prompt_control(gender, pitch, speed, text)
+
+        else:
+            prompt, global_token_ids = self.process_prompt(
+                text, prompt_speech_path, prompt_text
+            )
+        model_inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device)
 
         # Generate speech using the model
         generated_ids = self.model.generate(
@@ -117,14 +205,27 @@ class SparkTTS:
 
         # Trim the output tokens to remove the input tokens
         generated_ids = [
-            output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+            output_ids[len(input_ids) :]
+            for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
         ]
 
         # Decode the generated tokens into text
         predicts = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
         # Extract semantic token IDs from the generated text
-        pred_semantic_ids = torch.tensor([int(token) for token in re.findall(r"\d+", predicts)]).long().unsqueeze(0)
+        pred_semantic_ids = (
+            torch.tensor([int(token) for token in re.findall(r"bicodec_semantic_(\d+)", predicts)])
+            .long()
+            .unsqueeze(0)
+        )
+
+        if gender is not None:
+            global_token_ids = (
+                torch.tensor([int(token) for token in re.findall(r"bicodec_global_(\d+)", predicts)])
+                .long()
+                .unsqueeze(0)
+                .unsqueeze(0)
+            )
 
         # Convert semantic tokens back to waveform
         wav = self.audio_tokenizer.detokenize(
@@ -132,4 +233,4 @@ class SparkTTS:
             pred_semantic_ids.to(self.device),
         )
 
-        return wav
+        return wav

cli/inference.py

@@ -12,16 +12,35 @@ def parse_args():
     """Parse command-line arguments."""
     parser = argparse.ArgumentParser(description="Run TTS inference.")
 
-    parser.add_argument("--model_dir", type=str, default="pretrained_models/Spark-TTS-0.5B",
-                        help="Path to the model directory")
-    parser.add_argument("--save_dir", type=str, default="example/results",
-                        help="Directory to save generated audio files")
+    parser.add_argument(
+        "--model_dir",
+        type=str,
+        default="pretrained_models/Spark-TTS-0.5B",
+        help="Path to the model directory",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        default="example/results",
+        help="Directory to save generated audio files",
+    )
     parser.add_argument("--device", type=int, default=0, help="CUDA device number")
-    parser.add_argument("--text", type=str, required=True, help="Text for TTS generation")
+    parser.add_argument(
+        "--text", type=str, required=True, help="Text for TTS generation"
+    )
     parser.add_argument("--prompt_text", type=str, help="Transcript of prompt audio")
-    parser.add_argument("--prompt_speech_path", type=str, required=True,
-                        help="Path to the prompt audio file")
-
+    parser.add_argument(
+        "--prompt_speech_path",
+        type=str,
+        help="Path to the prompt audio file",
+    )
+    parser.add_argument("--gender", choices=["male", "pitch"])
+    parser.add_argument(
+        "--pitch", choices=["very_low", "low", "moderate", "high", "very_high"]
+    )
+    parser.add_argument(
+        "--speed", choices=["very_low", "low", "moderate", "high", "very_high"]
+    )
     return parser.parse_args()
 
 
@@ -47,14 +66,23 @@ def run_tts(args):
 
     # Perform inference and save the output audio
     with torch.no_grad():
-        wav = model.inference(args.text, args.prompt_speech_path, prompt_text=args.prompt_text)
+        wav = model.inference(
+            args.text,
+            args.prompt_speech_path,
+            prompt_text=args.prompt_text,
+            gender=args.gender,
+            pitch=args.pitch,
+            speed=args.speed,
+        )
         sf.write(save_path, wav, samplerate=16000)
 
     logging.info(f"Audio saved at: {save_path}")
 
 
 if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    logging.basicConfig(
+        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+    )
 
     args = parse_args()
     run_tts(args)