diff --git a/.gitignore b/.gitignore
index 146a43c..4007272 100644
--- a/.gitignore
+++ b/.gitignore
@@ -172,3 +172,6 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+# runtime
+triton_python_backend_utils.py
\ No newline at end of file
diff --git a/cli/SparkTTS.py b/cli/SparkTTS.py
index bc86ce3..c927502 100644
--- a/cli/SparkTTS.py
+++ b/cli/SparkTTS.py
@@ -13,15 +13,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+import math
 import re
+from threading import Thread
+import uuid
 import torch
-from typing import Tuple
+from typing import Generator, Tuple
 from pathlib import Path
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+from sparktts.utils import ThreadSafeDict
 from sparktts.utils.file import load_config
 from sparktts.models.audio_tokenizer import BiCodecTokenizer
 from sparktts.utils.token_parser import LEVELS_MAP, GENDER_MAP, TASK_TOKEN_MAP
+from cli.streamer import TokenStreamer
 
 
 class SparkTTS:
@@ -29,7 +35,18 @@ class SparkTTS:
     Spark-TTS for text-to-speech generation.
     """
 
-    def __init__(self, model_dir: Path, device: torch.device = torch.device("cuda:0")):
+    def __init__(
+        self,
+        model_dir: Path,
+        device: torch.device = torch.device("cuda:0"),
+        stream: bool = False,
+        stream_factor: int = 2,
+        stream_scale_factor: float = 1.0,
+        max_stream_factor: int = 2,
+        token_overlap_len: int = 0,
+        input_frame_rate: int = 25,
+        **kwargs,
+    ):
         """
         Initializes the SparkTTS model with the provided configurations and device.
 
@@ -37,11 +54,32 @@ def __init__(self, model_dir: Path, device: torch.device = torch.device("cuda:0"
             model_dir (Path): Directory containing the model and config files.
             device (torch.device): The device (CPU/GPU) to run the model on.
         """
+        if stream is True:
+            # fast path to check params
+            # rtf and decoding related
+            assert (
+                stream_factor >= 2
+            ), f"stream_factor must >=2 increase for better speech quality, but rtf slow (speech quality vs rtf)"
+            self.stream_factor = stream_factor
+            self.max_stream_factor = max_stream_factor
+            assert (
+                stream_scale_factor >= 1.0
+            ), "stream_scale_factor should be greater than 1, change it according to your actual rtf"
+            self.stream_scale_factor = stream_scale_factor  # scale speed
+            assert (
+                token_overlap_len >= 0
+            ), "token_overlap_len should be greater than 0, change it according to your actual rtf"
+            self.token_overlap_len = token_overlap_len
+            self.input_frame_rate = input_frame_rate
+
         self.device = device
         self.model_dir = model_dir
         self.configs = load_config(f"{model_dir}/config.yaml")
         self.sample_rate = self.configs["sample_rate"]
         self._initialize_inference()
+        self.start_global_token_id = self.tokenizer.encode("<|start_global_token|>")[0]
+        self.start_semantic_token_id = self.tokenizer.encode("<|start_semantic_token|>")[0]
+        logging.debug(f"start_global_token_id:{self.start_global_token_id} start_semantic_token_id:{self.start_semantic_token_id}")
 
     def _initialize_inference(self):
         """Initializes the tokenizer, model, and audio tokenizer for inference."""
@@ -68,12 +106,8 @@ def process_prompt(
             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:
@@ -138,9 +172,7 @@ def process_prompt_control(
         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]
-        )
+        attribte_tokens = "".join([gender_tokens, pitch_label_tokens, speed_label_tokens])
 
         control_tts_inputs = [
             TASK_TOKEN_MAP["controllable_tts"],
@@ -154,6 +186,41 @@ def process_prompt_control(
 
         return "".join(control_tts_inputs)
 
+    def token2wav(self, generated_ids: torch.Tensor, gender: str, global_token_ids: torch.Tensor):
+        """
+        generated_ids -- tokenizer.decode --> sematic tokens + global tokens  -- audio_tokenizer.detokenize --> waveform
+        """
+        #print("generated_ids", generated_ids)
+        # Decode the generated tokens into text (just a mapping, so quick,don't worry)
+        predicts = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        #print("predicts", predicts)
+
+        # Extract semantic token IDs from the generated text
+        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:
+            # Tips: generated_id - global_vq_index = 151665
+            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(
+            global_token_ids.to(self.device).squeeze(0),
+            pred_semantic_ids.to(self.device),
+        )
+
+        return wav
+
     @torch.no_grad()
     def inference(
         self,
@@ -184,13 +251,11 @@ def inference(
         Returns:
             torch.Tensor: Generated waveform as a tensor.
         """
+        global_token_ids = None
         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
-            )
+            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
@@ -209,28 +274,123 @@ def inference(
             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]
+        wav = self.token2wav(generated_ids, gender, global_token_ids)
 
-        # Extract semantic token IDs from the generated text
-        pred_semantic_ids = (
-            torch.tensor([int(token) for token in re.findall(r"bicodec_semantic_(\d+)", predicts)])
-            .long()
-            .unsqueeze(0)
-        )
+        return wav
 
+    @torch.no_grad()
+    def inference_stream(
+        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,
+    ) -> Generator[torch.Tensor, None, None]:
+        """
+        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 generator.
+        """
+        global_token_ids = None
         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)
-            )
+            prompt = self.process_prompt_control(gender, pitch, speed, text)
 
-        # Convert semantic tokens back to waveform
-        wav = self.audio_tokenizer.detokenize(
-            global_token_ids.to(self.device).squeeze(0),
-            pred_semantic_ids.to(self.device),
-        )
+        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)
 
-        return wav
\ No newline at end of file
+        # session streamer, skip input prompt
+        streamer = TokenStreamer(skip_prompt=True)
+
+        generation_kwargs = dict(
+            **model_inputs,
+            streamer=streamer,
+            max_new_tokens=3000,
+            do_sample=True,
+            top_k=top_k,
+            top_p=top_p,
+            temperature=temperature,
+        )
+        # print("generation_kwargs", generation_kwargs)
+
+        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        is_meet_start_global_token = False
+        is_meet_start_semantic_token = False
+        controll_gen_global_token_ids = []
+        semantic_token_ids = []
+
+        max_batch_size = math.ceil(self.max_stream_factor * self.input_frame_rate)
+        batch_size = math.ceil(self.stream_factor * self.input_frame_rate)
+        logging.info(f"init batch_size: {batch_size} max_batch_size: {max_batch_size}")
+
+        for token_id in streamer:
+            if gender is not None:  # Inference Overview of Controlled Generation
+                if is_meet_start_global_token is False and token_id != self.start_global_token_id:
+                    continue
+                if is_meet_start_global_token is False and token_id == self.start_global_token_id:
+                    is_meet_start_global_token = True
+                    controll_gen_global_token_ids.append(token_id)
+                    continue
+                # append global token until meet start_global_token
+                if (
+                    is_meet_start_global_token is True
+                    and is_meet_start_semantic_token is False
+                    and token_id != self.start_global_token_id
+                ):
+                    controll_gen_global_token_ids.append(token_id)
+
+                if is_meet_start_semantic_token is False and token_id != self.start_semantic_token_id:
+                    continue
+                if is_meet_start_semantic_token is False and token_id == self.start_semantic_token_id:
+                    is_meet_start_semantic_token = True
+                    continue
+                # do batch stream until meet start_semantic_token
+                if is_meet_start_semantic_token is True and token_id != self.start_semantic_token_id:
+                    # print(controll_gen_global_token_ids)
+                    pass
+
+            semantic_token_ids.append(token_id)
+            # if len(semantic_token_ids) % batch_size == 0:
+            if len(semantic_token_ids) >= batch_size + self.token_overlap_len:
+                batch = semantic_token_ids[: batch_size + self.token_overlap_len]
+                # Process each batch
+                sub_tts_speech = self.token2wav(
+                    [controll_gen_global_token_ids + batch], gender, global_token_ids
+                )  # one batch
+                yield {"tts_speech": sub_tts_speech, "sample_rate": self.sample_rate}
+                semantic_token_ids = semantic_token_ids[batch_size:]
+                # increase token_hop_len for better speech quality
+                batch_size = min(max_batch_size, int(batch_size * self.stream_scale_factor))
+                logging.info(
+                    f"increase batch_size: {batch_size} token_overlap_len:{self.token_overlap_len}"
+                )
+
+        if len(semantic_token_ids) > 0:  # end to finalize
+            # Process each batch
+            sub_tts_speech = self.token2wav(
+                [controll_gen_global_token_ids + semantic_token_ids], gender, global_token_ids
+            )  # one batch
+            yield {"tts_speech": sub_tts_speech, "sample_rate": self.sample_rate}
+            logging.info(f"last batch len: {len(semantic_token_ids)}")
+
+        torch.cuda.empty_cache()
diff --git a/cli/__init__.py b/cli/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/cli/inference.py b/cli/inference.py
index 349f7ea..781f39a 100644
--- a/cli/inference.py
+++ b/cli/inference.py
@@ -42,9 +42,7 @@ def parse_args():
         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",
@@ -52,12 +50,8 @@ def parse_args():
         help="Path to the prompt audio file",
     )
     parser.add_argument("--gender", choices=["male", "female"])
-    parser.add_argument(
-        "--pitch", choices=["very_low", "low", "moderate", "high", "very_high"]
-    )
-    parser.add_argument(
-        "--speed", choices=["very_low", "low", "moderate", "high", "very_high"]
-    )
+    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()
 
 
@@ -107,10 +101,24 @@ def run_tts(args):
     logging.info(f"Audio saved at: {save_path}")
 
 
+"""
+# Inference Overview of Controlled Generation
+PYTHONPATH=./ python cli/inference.py \
+    --text "身临其境，换新体验。塑造开源语音合成新范式，让智能语音更自然。" \
+    --save_dir "example/results" \
+    --model_dir ../../models/SparkAudio/Spark-TTS-0.5B \
+    --gender female --pitch  moderate --speed high
+
+# Inference Overview of Voice Cloning
+PYTHONPATH=./ python cli/inference.py \
+    --text "身临其境，换新体验。塑造开源语音合成新范式，让智能语音更自然。" \
+    --save_dir "example/results" \
+    --model_dir ../../models/SparkAudio/Spark-TTS-0.5B \
+    --prompt_text "吃燕窝就选燕之屋，本节目由26年专注高品质燕窝的燕之屋冠名播出。豆奶牛奶换着喝，营养更均衡，本节目由豆本豆豆奶特约播出。" \
+    --prompt_speech_path "example/prompt_audio.wav"
+"""
 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)
diff --git a/cli/inference_stream.py b/cli/inference_stream.py
new file mode 100644
index 0000000..abec662
--- /dev/null
+++ b/cli/inference_stream.py
@@ -0,0 +1,173 @@
+# Copyright (c) 2025 SparkAudio
+#               2025 Xinsheng Wang (w.xinshawn@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import argparse
+import torch
+import soundfile as sf
+import logging
+from datetime import datetime
+import platform
+
+from cli.SparkTTS import SparkTTS
+from sparktts.utils.audio import merge_numpy_darray
+
+
+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("--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("--prompt_text", type=str, help="Transcript of prompt audio")
+    parser.add_argument(
+        "--prompt_speech_path",
+        type=str,
+        help="Path to the prompt audio file",
+    )
+    parser.add_argument("--gender", choices=["male", "female"])
+    parser.add_argument("--pitch", choices=["very_low", "low", "moderate", "high", "very_high"])
+    parser.add_argument("--speed", choices=["very_low", "low", "moderate", "high", "very_high"])
+    parser.add_argument(
+        "--stream-factor", type=int, default=2, help="Synthesis audios stream factor"
+    )
+    parser.add_argument(
+        "--stream-scale-factor",
+        type=float,
+        default=1.0,
+        help="Synthesis audios stream scale factor",
+    )
+    parser.add_argument(
+        "--max-stream-factor", type=int, default=2, help="Synthesis audios max stream factor"
+    )
+    parser.add_argument(
+        "--token-overlap-len", type=int, default=0, help="Synthesis audios token overlap len"
+    )
+    return parser.parse_args()
+
+
+def run_tts(args):
+    """Perform TTS inference and save the generated audio."""
+    logging.info(f"Using model from: {args.model_dir}")
+    logging.info(f"Saving audio to: {args.save_dir}")
+
+    # Ensure the save directory exists
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    # Convert device argument to torch.device
+    if platform.system() == "Darwin" and torch.backends.mps.is_available():
+        # macOS with MPS support (Apple Silicon)
+        device = torch.device(f"mps:{args.device}")
+        logging.info(f"Using MPS device: {device}")
+    elif torch.cuda.is_available():
+        # System with CUDA support
+        device = torch.device(f"cuda:{args.device}")
+        logging.info(f"Using CUDA device: {device}")
+    else:
+        # Fall back to CPU
+        device = torch.device("cpu")
+        logging.info("GPU acceleration not available, using CPU")
+
+    # Initialize the model
+    model = SparkTTS(
+        args.model_dir,
+        device,
+        stream=True,
+        stream_factor=args.stream_factor,
+        stream_scale_factor=args.stream_scale_factor,
+        max_stream_factor=args.max_stream_factor,
+        token_overlap_len=args.token_overlap_len,
+    )
+
+    # Generate unique filename using timestamp
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    save_path = os.path.join(args.save_dir, f"{timestamp}.wav")
+
+    logging.info("Starting stream inference...")
+
+    sub_tts_speechs = []
+    # Perform inference and save the output audio
+    with torch.no_grad():
+        batch_stream = model.inference_stream(
+            args.text,
+            args.prompt_speech_path,
+            prompt_text=args.prompt_text,
+            gender=args.gender,
+            pitch=args.pitch,
+            speed=args.speed,
+        )
+        for item in batch_stream:
+            sub_tts_speechs.append(item["tts_speech"])
+
+    output_audio = merge_numpy_darray(sub_tts_speechs)  # [[T],...] -> [T]
+    sf.write(save_path, output_audio, samplerate=16000)
+    logging.info(f"Audio saved at: {save_path}")
+
+
+"""
+# Inference Overview of Controlled Generation
+PYTHONPATH=./ python cli/inference_stream.py \
+    --text "身临其境，换新体验。塑造开源语音合成新范式，让智能语音更自然。" \
+    --save_dir "example/results" \
+    --model_dir ../../models/SparkAudio/Spark-TTS-0.5B \
+    --gender female --pitch  moderate --speed high
+
+PYTHONPATH=./ python cli/inference_stream.py \
+    --text "万物之始,大道至简,衍化至繁。君不见黄河之水天上来，奔流到海不复回。君不见高堂明镜悲白发，朝如青丝暮成雪。人生得意须尽欢，莫使金樽空对月。天生我材必有用，千金散尽还复来。" \
+    --save_dir "example/results" \
+    --model_dir ../../models/SparkAudio/Spark-TTS-0.5B \
+    --gender female --pitch  moderate --speed high
+
+# Inference Overview of Voice Cloning
+# default use static batch is ok
+PYTHONPATH=./ python cli/inference_stream.py \
+    --text "身临其境，换新体验。塑造开源语音合成新范式，让智能语音更自然。" \
+    --save_dir "example/results" \
+    --model_dir ../../models/SparkAudio/Spark-TTS-0.5B \
+    --prompt_text "吃燕窝就选燕之屋，本节目由26年专注高品质燕窝的燕之屋冠名播出。豆奶牛奶换着喝，营养更均衡，本节目由豆本豆豆奶特约播出。" \
+    --prompt_speech_path "example/prompt_audio.wav"
+
+PYTHONPATH=./ python cli/inference_stream.py \
+    --text "万物之始,大道至简,衍化至繁。君不见黄河之水天上来，奔流到海不复回。君不见高堂明镜悲白发，朝如青丝暮成雪。人生得意须尽欢，莫使金樽空对月。天生我材必有用，千金散尽还复来。" \
+    --save_dir "example/results" \
+    --model_dir ../../models/SparkAudio/Spark-TTS-0.5B \
+    --prompt_text "吃燕窝就选燕之屋，本节目由26年专注高品质燕窝的燕之屋冠名播出。豆奶牛奶换着喝，营养更均衡，本节目由豆本豆豆奶特约播出。" \
+    --prompt_speech_path "example/prompt_audio.wav"
+
+PYTHONPATH=./ python cli/inference_stream.py \
+    --text "万物之始,大道至简,衍化至繁。" \
+    --save_dir "example/results" \
+    --model_dir ../../models/SparkAudio/Spark-TTS-0.5B \
+    --prompt_text "欢迎大家来体验达摩院推出的语音识别模型" \
+    --prompt_speech_path "../../test/audio_files/asr_example_zh.wav"
+"""
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+    args = parse_args()
+    run_tts(args)
diff --git a/cli/streamer.py b/cli/streamer.py
new file mode 100644
index 0000000..57bce40
--- /dev/null
+++ b/cli/streamer.py
@@ -0,0 +1,41 @@
+from queue import Queue
+
+from transformers.generation.streamers import BaseStreamer
+
+
+class TokenStreamer(BaseStreamer):
+    def __init__(self, skip_prompt: bool = False, timeout=None):
+        self.skip_prompt = skip_prompt
+
+        # variables used in the streaming process
+        self.token_queue = Queue()
+        self.stop_signal = None
+        self.next_tokens_are_prompt = True
+        self.timeout = timeout
+
+    def put(self, value):
+        if len(value.shape) > 1 and value.shape[0] > 1:
+            raise ValueError("TextStreamer only supports batch size 1")
+        elif len(value.shape) > 1:
+            value = value[0]
+        # print(value)
+
+        if self.skip_prompt and self.next_tokens_are_prompt:
+            self.next_tokens_are_prompt = False
+            return
+
+        for token in value.tolist():
+            self.token_queue.put(token)
+
+    def end(self):
+        self.token_queue.put(self.stop_signal)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value = self.token_queue.get(timeout=self.timeout)
+        if value == self.stop_signal:
+            raise StopIteration()
+        else:
+            return value
diff --git a/example/prompt_recon.wav b/example/prompt_recon.wav
new file mode 100644
index 0000000..47799d4
Binary files /dev/null and b/example/prompt_recon.wav differ
diff --git a/sparktts/models/audio_tokenizer.py b/sparktts/models/audio_tokenizer.py
index d7065eb..587c96f 100644
--- a/sparktts/models/audio_tokenizer.py
+++ b/sparktts/models/audio_tokenizer.py
@@ -43,9 +43,7 @@ def __init__(self, model_dir: Path, device: torch.device = None, **kwargs):
 
     def _initialize_model(self):
         """Load and initialize the BiCodec model and Wav2Vec2 feature extractor."""
-        self.model = BiCodec.load_from_checkpoint(f"{self.model_dir}/BiCodec").to(
-            self.device
-        )
+        self.model = BiCodec.load_from_checkpoint(f"{self.model_dir}/BiCodec").to(self.device)
         self.processor = Wav2Vec2FeatureExtractor.from_pretrained(
             f"{self.model_dir}/wav2vec2-large-xlsr-53"
         )
@@ -82,7 +80,7 @@ def process_audio(self, wav_path: Path) -> Tuple[np.ndarray, torch.Tensor]:
         wav_ref = torch.from_numpy(wav_ref).unsqueeze(0).float()
         return wav, wav_ref
 
-    def extract_wav2vec2_features(self, wavs: torch.Tensor) -> torch.Tensor:
+    def extract_wav2vec2_features(self, wavs: np.ndarray) -> torch.Tensor:
         """extract wav2vec2 features"""
         inputs = self.processor(
             wavs,
@@ -92,9 +90,7 @@ def extract_wav2vec2_features(self, wavs: torch.Tensor) -> torch.Tensor:
             output_hidden_states=True,
         ).input_values
         feat = self.feature_extractor(inputs.to(self.feature_extractor.device))
-        feats_mix = (
-            feat.hidden_states[11] + feat.hidden_states[14] + feat.hidden_states[16]
-        ) / 3
+        feats_mix = (feat.hidden_states[11] + feat.hidden_states[14] + feat.hidden_states[16]) / 3
 
         return feats_mix
 
@@ -129,9 +125,7 @@ def tokenize(self, audio_path: str) -> Tuple[torch.Tensor, torch.Tensor]:
 
         return global_tokens, semantic_tokens
 
-    def detokenize(
-        self, global_tokens: torch.Tensor, semantic_tokens: torch.Tensor
-    ) -> np.array:
+    def detokenize(self, global_tokens: torch.Tensor, semantic_tokens: torch.Tensor) -> np.array:
         """detokenize the tokens to waveform
 
         Args:
@@ -149,10 +143,11 @@ def detokenize(
 # test
 if __name__ == "__main__":
     import soundfile as sf
+    import os
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     tokenizer = BiCodecTokenizer(
-        model_dir="pretrained_models/Spark-TTS-0.5B",
+        model_dir=os.getenv("MODEL_DIR", "pretrained_models/Spark-TTS-0.5B"),
         device=device,
     )
     wav_path = "example/prompt_audio.wav"
diff --git a/sparktts/models/bicodec.py b/sparktts/models/bicodec.py
index 8cab2f0..82fe1df 100644
--- a/sparktts/models/bicodec.py
+++ b/sparktts/models/bicodec.py
@@ -17,7 +17,6 @@
 import torch.nn as nn
 from pathlib import Path
 from typing import Dict, Any
-from omegaconf import DictConfig
 from safetensors.torch import load_file
 
 from sparktts.utils.file import load_config
@@ -43,7 +42,7 @@ def __init__(
         speaker_encoder: nn.Module,
         prenet: nn.Module,
         postnet: nn.Module,
-        **kwargs
+        **kwargs,
     ) -> None:
         """
         Initializes the BiCodec model with the required components.
@@ -73,12 +72,12 @@ def load_from_checkpoint(cls, model_dir: Path, **kwargs) -> "BiCodec":
 
         Args:
             model_dir (Path): Path to the model directory containing checkpoint and config.
-        
+
         Returns:
             BiCodec: The initialized BiCodec model.
         """
-        ckpt_path = f'{model_dir}/model.safetensors'
-        config = load_config(f'{model_dir}/config.yaml')['audio_tokenizer']
+        ckpt_path = f"{model_dir}/model.safetensors"
+        config = load_config(f"{model_dir}/config.yaml")["audio_tokenizer"]
         mel_params = config["mel_params"]
         encoder = Encoder(**config["encoder"])
         quantizer = FactorizedVectorQuantize(**config["quantizer"])
@@ -116,7 +115,7 @@ def forward(self, batch: Dict[str, Any]) -> Dict[str, Any]:
 
         Args:
             batch (dict): A dictionary containing features, reference waveform, and target waveform.
-        
+
         Returns:
             dict: A dictionary containing the reconstruction, features, and other metrics.
         """
@@ -212,6 +211,7 @@ def init_mel_transformer(self, config: Dict[str, Any]):
 
     def remove_weight_norm(self):
         """Removes weight normalization from all layers."""
+
         def _remove_weight_norm(m):
             try:
                 torch.nn.utils.remove_weight_norm(m)
@@ -223,16 +223,20 @@ def _remove_weight_norm(m):
 
 # Test the model
 if __name__ == "__main__":
-
     config = load_config("pretrained_models/SparkTTS-0.5B/BiCodec/config.yaml")
     model = BiCodec.load_from_checkpoint(
         model_dir="pretrained_models/SparkTTS-0.5B/BiCodec",
     )
+    device = "cpu" if not torch.cuda.is_available() else "cuda"
+    print(model)
+    model_million_params = sum(p.numel() for p in model.parameters()) / 1e6
+    print(f"{model_million_params}M parameters")
+    model.to(device)
 
     # Generate random inputs for testing
     duration = 0.96
-    x = torch.randn(20, 1, int(duration * 16000))
-    feat = torch.randn(20, int(duration * 50), 1024)
+    x = torch.randn(20, 1, int(duration * 16000)).to(device)
+    feat = torch.randn(20, int(duration * 50), 1024).to(device)
     inputs = {"feat": feat, "wav": x, "ref_wav": x}
 
     # Forward pass
@@ -241,7 +245,8 @@ def _remove_weight_norm(m):
     wav_recon = model.detokenize(semantic_tokens, global_tokens)
 
     # Verify if the reconstruction matches
-    if torch.allclose(outputs["recons"].detach(), wav_recon):
+    if torch.allclose(outputs["recons"].detach(), wav_recon, rtol=1e-3, atol=1e-5):
+        # if torch.allclose(outputs["recons"].detach(), wav_recon):
         print("Test successful")
     else:
         print("Test failed")
diff --git a/sparktts/utils/__init__.py b/sparktts/utils/__init__.py
index e69de29..1e8f788 100644
--- a/sparktts/utils/__init__.py
+++ b/sparktts/utils/__init__.py
@@ -0,0 +1,20 @@
+import threading
+
+
+class ThreadSafeDict:
+    def __init__(self):
+        self._dict = {}
+        # 使用 RLock 可重入锁，避免死锁
+        self._lock = threading.RLock()
+
+    def get(self, key, default=None):
+        with self._lock:
+            return self._dict.get(key, default)
+
+    def set(self, key, value):
+        with self._lock:
+            self._dict[key] = value
+
+    def pop(self, key):
+        with self._lock:
+            return self._dict.pop(key, None)
diff --git a/sparktts/utils/audio.py b/sparktts/utils/audio.py
index 105cd9c..1bb028e 100644
--- a/sparktts/utils/audio.py
+++ b/sparktts/utils/audio.py
@@ -46,9 +46,7 @@ def audio_volume_normalize(audio: np.ndarray, coeff: float = 0.2) -> np.ndarray:
 
     # If the maximum value is less than 0.1, scale the array to have a maximum of 0.1
     if temp[-1] < 0.1:
-        scaling_factor = max(
-            temp[-1], 1e-3
-        )  # Prevent division by zero with a small constant
+        scaling_factor = max(temp[-1], 1e-3)  # Prevent division by zero with a small constant
         audio = audio / scaling_factor * 0.1
 
     # Filter out values less than 0.01 from temp
@@ -168,9 +166,7 @@ def stft(
         Tensor: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
     """
 
-    x_stft = torch.stft(
-        x, fft_size, hop_size, win_length, window.to(x.device), return_complex=True
-    )
+    x_stft = torch.stft(x, fft_size, hop_size, win_length, window.to(x.device), return_complex=True)
 
     # clamp is needed to avoid nan or inf
     if not use_complex:
@@ -188,73 +184,64 @@ def detect_speech_boundaries(
     sample_rate: int,
     window_duration: float = 0.1,
     energy_threshold: float = 0.01,
-    margin_factor: int = 2
+    margin_factor: int = 2,
 ) -> Tuple[int, int]:
     """Detect the start and end points of speech in an audio signal using RMS energy.
-    
+
     Args:
         wav: Input audio signal array with values in [-1, 1]
         sample_rate: Audio sample rate in Hz
         window_duration: Duration of detection window in seconds
         energy_threshold: RMS energy threshold for speech detection
         margin_factor: Factor to determine extra margin around detected boundaries
-        
+
     Returns:
         tuple: (start_index, end_index) of speech segment
-        
+
     Raises:
         ValueError: If the audio contains only silence
     """
     window_size = int(window_duration * sample_rate)
     margin = margin_factor * window_size
     step_size = window_size // 10
-    
+
     # Create sliding windows using stride tricks to avoid loops
     windows = sliding_window_view(wav, window_size)[::step_size]
-    
+
     # Calculate RMS energy for each window
-    energy = np.sqrt(np.mean(windows ** 2, axis=1))
+    energy = np.sqrt(np.mean(windows**2, axis=1))
     speech_mask = energy >= energy_threshold
-    
+
     if not np.any(speech_mask):
         raise ValueError("No speech detected in audio (only silence)")
-    
+
     start = max(0, np.argmax(speech_mask) * step_size - margin)
     end = min(len(wav), (len(speech_mask) - 1 - np.argmax(speech_mask[::-1])) * step_size + margin)
-    
+
     return start, end
 
 
 def remove_silence_on_both_ends(
-    wav: np.ndarray,
-    sample_rate: int,
-    window_duration: float = 0.1,
-    volume_threshold: float = 0.01
+    wav: np.ndarray, sample_rate: int, window_duration: float = 0.1, volume_threshold: float = 0.01
 ) -> np.ndarray:
     """Remove silence from both ends of an audio signal.
-    
+
     Args:
         wav: Input audio signal array
         sample_rate: Audio sample rate in Hz
         window_duration: Duration of detection window in seconds
         volume_threshold: Amplitude threshold for silence detection
-        
+
     Returns:
         np.ndarray: Audio signal with silence removed from both ends
-        
+
     Raises:
         ValueError: If the audio contains only silence
     """
-    start, end = detect_speech_boundaries(
-        wav,
-        sample_rate,
-        window_duration,
-        volume_threshold
-    )
+    start, end = detect_speech_boundaries(wav, sample_rate, window_duration, volume_threshold)
     return wav[start:end]
 
 
-
 def hertz_to_mel(pitch: float) -> float:
     """
     Converts a frequency from the Hertz scale to the Mel scale.
@@ -268,4 +255,35 @@ def hertz_to_mel(pitch: float) -> float:
         Frequency in Mel scale.
     """
     mel = 2595 * np.log10(1 + pitch / 700)
-    return mel
\ No newline at end of file
+    return mel
+
+
+def merge_numpy_darray(sub_arrays: list[np.ndarray]) -> np.ndarray | None:
+    """
+    Merges a list of NumPy arrays into a single NumPy array.
+    This function is designed to handle arrays that represent audio
+    waveforms, where each array has shape (num_channels, sequence_length)
+    but may have different lengths along the sequence_length dimension.
+
+    Args:
+        sub_arrays: A list of NumPy arrays. [1d array, 1d array, ...]
+
+    Returns:
+        A single NumPy array with all the sub-arrays concatenated along
+        the sequence_length (time) dimension.
+        Returns None if the input list is empty or if the arrays have
+        inconsistent shapes (different number of channels).
+    """
+    if not sub_arrays:
+        return None
+
+    total_length = sum(arr.shape[0] for arr in sub_arrays)
+    dtype = sub_arrays[0].dtype
+    merged_array = np.empty((total_length), dtype=dtype)
+    current_position = 0
+
+    for arr in sub_arrays:
+        merged_array[current_position : current_position + arr.shape[0]] = arr
+        current_position += arr.shape[0]
+
+    return merged_array
diff --git a/sparktts/utils/token_parser.py b/sparktts/utils/token_parser.py
index cc43782..9dbd13c 100644
--- a/sparktts/utils/token_parser.py
+++ b/sparktts/utils/token_parser.py
@@ -156,10 +156,11 @@ def emotion(emotion: str):
 
 # test
 if __name__ == "__main__":
+    import os
     from transformers import AutoTokenizer
 
     tokenizer = AutoTokenizer.from_pretrained(
-        "/aifs4su/xinshengwang/code/StyleCraft/tokenizer/stylecraft-bicodec-pitch-loudness-speed-emotion-tokenizer"
+        os.getenv("TOKENIZER_PATH", "/aifs4su/xinshengwang/code/StyleCraft/tokenizer/stylecraft-bicodec-pitch-loudness-speed-emotion-tokenizer")
     )
 
     tasks = ["tts", "tts", "understand", "controllable_tts", "prompt_tts"]
@@ -183,5 +184,7 @@ def emotion(emotion: str):
         inputs = [task, age, gender, mel, mel_level, loudness, loudness_level, emotion]
         inputs = "".join(inputs)
         ids = tokenizer.encode(inputs, add_special_tokens=False)
-        print(ids)
-        print("decode", tokenizer.decode(ids))
+        print("tokenized ids",ids)
+        tokens = tokenizer.decode(ids)
+        print("decoded tokens", tokens)
+        assert tokens == inputs