xinference 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

xinference/thirdparty/deepseek_vl2/utils/io.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2023-2024 DeepSeek.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to
+# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+# the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+import json
+from typing import Dict, List
+
+import PIL.Image
+import torch
+from transformers import AutoModelForCausalLM
+
+
+def load_pretrained_model(model_path: str):
+
+    from deepseek_vl2.models.processing_deepseek_vl_v2 import DeepseekVLV2Processor
+    from deepseek_vl2.models.modeling_deepseek_vl_v2 import DeepseekVLV2ForCausalLM
+
+    vl_chat_processor = DeepseekVLV2Processor.from_pretrained(model_path)
+    tokenizer = vl_chat_processor.tokenizer
+
+    vl_gpt: DeepseekVLV2ForCausalLM = AutoModelForCausalLM.from_pretrained(
+        model_path, trust_remote_code=True
+    )
+    vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
+
+    return tokenizer, vl_chat_processor, vl_gpt
+
+
+def load_pil_images(conversations: List[Dict[str, str]]) -> List[PIL.Image.Image]:
+    """
+
+    Args:
+        conversations (List[Dict[str, str]]): the conversations with a list of messages. An example is :
+            [
+                {
+                    "role": "User",
+                    "content": "<image>\nExtract all information from this image and convert them into markdown format.",
+                    "images": ["./examples/table_datasets.png"]
+                },
+                {"role": "Assistant", "content": ""},
+            ]
+
+    Returns:
+        pil_images (List[PIL.Image.Image]): the list of PIL images.
+
+    """
+
+    pil_images = []
+
+    for message in conversations:
+        if "images" not in message:
+            continue
+
+        for image_path in message["images"]:
+            pil_img = PIL.Image.open(image_path)
+            pil_img = pil_img.convert("RGB")
+            pil_images.append(pil_img)
+
+    return pil_images
+
+
+def load_json(filepath):
+    with open(filepath, "r") as f:
+        data = json.load(f)
+        return data

xinference/thirdparty/megatts3/tts/frontend_function.py

@@ -0,0 +1,175 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# 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 torch
+import torch.nn.functional as F
+import whisper
+import librosa
+from copy import deepcopy
+from tts.utils.text_utils.ph_tone_convert import split_ph_timestamp, split_ph
+from tts.utils.audio_utils.align import mel2token_to_dur
+
+''' Graphme to phoneme function '''
+def g2p(self, text_inp):
+    # prepare inputs
+    txt_token = self.g2p_tokenizer('<BOT>' + text_inp + '<BOS>')['input_ids']
+    input_ids = torch.LongTensor([txt_token+[145+self.speech_start_idx]]).to(self.device)
+
+    # model forward
+    with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):
+        outputs = self.g2p_model.generate(input_ids, max_new_tokens=256, do_sample=True, top_k=1, eos_token_id=800+1+self.speech_start_idx)
+    
+    # process outputs
+    ph_tokens = outputs[:, len(txt_token):-1]-self.speech_start_idx
+    ph_pred, tone_pred = split_ph(ph_tokens[0])
+    ph_pred, tone_pred = ph_pred[None, :].to(self.device), tone_pred[None, :].to(self.device)
+    return ph_pred, tone_pred
+
+''' Get phoneme2mel align of prompt speech '''
+def align(self, wav):
+    with torch.inference_mode():
+        whisper_wav = librosa.resample(wav, orig_sr=self.sr, target_sr=16000)
+        mel = torch.FloatTensor(whisper.log_mel_spectrogram(whisper_wav).T).to(self.device)[None].transpose(1,2)
+        prompt_max_frame = mel.size(2) // self.fm * self.fm
+        mel = mel[:, :, :prompt_max_frame]
+        token = torch.LongTensor([[798]]).to(self.device)
+        audio_features = self.aligner_lm.embed_audio(mel)
+        for i in range(768):
+            with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):
+                logits = self.aligner_lm.logits(token, audio_features, None)
+                token_pred = torch.argmax(F.softmax(logits[:, -1], dim=-1), 1)[None]
+                token = torch.cat([token, token_pred], dim=1)
+                if token_pred[0] == 799:
+                    break
+        alignment_tokens = token
+    
+    ph_ref, tone_ref, dur_ref, _ = split_ph_timestamp(deepcopy(alignment_tokens)[0, 1:-1])
+    ph_ref = torch.Tensor(ph_ref)[None].to(self.device)
+    tone_ref = torch.Tensor(tone_ref)[None].to(self.device)
+    if dur_ref.sum() < prompt_max_frame:
+        dur_ref[-1] += prompt_max_frame - dur_ref.sum()
+    elif dur_ref.sum() > prompt_max_frame:
+        len_diff = dur_ref.sum() - prompt_max_frame
+        while True:
+            for i in range(len(dur_ref)):
+                dur_ref[i] -= 1
+                len_diff -= 1
+                if len_diff == 0:
+                    break
+            if len_diff == 0:
+                break
+    mel2ph_ref = self.length_regulator(dur_ref[None]).to(self.device)
+    mel2ph_ref = mel2ph_ref[:, :mel2ph_ref.size(1)//self.fm*self.fm]
+    return ph_ref, tone_ref, mel2ph_ref
+
+''' Duration Prompting '''
+def make_dur_prompt(self, mel2ph_ref, ph_ref, tone_ref):
+    dur_tokens_2d_ = mel2token_to_dur(mel2ph_ref, ph_ref.shape[1]).clamp(
+                    max=self.hp_dur_model['dur_code_size'] - 1) + 1
+
+    ctx_dur_tokens = dur_tokens_2d_.clone().flatten(0, 1).to(self.device)
+    txt_tokens_flat_ = ph_ref.flatten(0, 1)
+    ctx_dur_tokens = ctx_dur_tokens[txt_tokens_flat_ > 0][None]
+
+    last_dur_pos_prompt = ctx_dur_tokens.shape[1]
+    dur_spk_pos_ids_flat = range(0, last_dur_pos_prompt)
+    dur_spk_pos_ids_flat = torch.LongTensor([dur_spk_pos_ids_flat]).to(self.device)
+    with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):
+        _, incremental_state_dur_prompt = self.dur_model.infer(
+            ph_ref, {'tone': tone_ref}, None, None, None,
+            ctx_vqcodes=ctx_dur_tokens, spk_pos_ids_flat=dur_spk_pos_ids_flat, return_state=True)
+    return incremental_state_dur_prompt, ctx_dur_tokens
+
+''' Duration Prediction '''
+def dur_pred(self, ctx_dur_tokens, incremental_state_dur_prompt, ph_pred, tone_pred, seg_i, dur_disturb, dur_alpha, is_first, is_final):
+    last_dur_token = ctx_dur_tokens[:, -1:]
+    last_dur_pos_prompt = ctx_dur_tokens.shape[1]
+    incremental_state_dur = deepcopy(incremental_state_dur_prompt)
+    txt_len = ph_pred.shape[1]
+    dur_spk_pos_ids_flat = range(last_dur_pos_prompt, last_dur_pos_prompt + txt_len)
+    dur_spk_pos_ids_flat = torch.LongTensor([dur_spk_pos_ids_flat]).to(self.device)
+    last_dur_pos_prompt = last_dur_pos_prompt + txt_len
+
+    with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):
+        dur_pred = self.dur_model.infer(
+            ph_pred, {'tone': tone_pred}, None, None, None,
+            incremental_state=incremental_state_dur,
+            first_decoder_inp=last_dur_token,
+            spk_pos_ids_flat=dur_spk_pos_ids_flat,
+        )
+
+    dur_pred = dur_pred - 1
+    dur_pred = dur_pred.clamp(0, self.hp_dur_model['dur_code_size'] - 1)
+    # if is_final:
+    #     dur_pred[:, -1] = dur_pred[:, -1].clamp(64, 128)
+    # else:
+    #     dur_pred[:, -1] = dur_pred[:, -1].clamp(48, 128)
+    # if seg_i > 0:
+        # dur_pred[:, 0] = 0
+    # ['。', '！', '？', 'sil']
+    for sil_token in [148, 153, 166, 145]:
+        dur_pred[ph_pred==sil_token].clamp_min(32)
+    # ['，', '；'] 
+    for sil_token in [163, 165]:
+        dur_pred[ph_pred==sil_token].clamp_min(16)
+    if not is_final:
+        # add 0.32ms for crossfade
+        dur_pred[:, -1] =  dur_pred[:, -1] + 32
+    else:
+        dur_pred[:, -1] = dur_pred[:, -1].clamp(64, 128)
+
+    ''' DiT target speech generation '''
+    dur_disturb_choice = (torch.rand_like(dur_pred.float()) > 0.5).float()
+    dur_disturb_r = 1 + torch.rand_like(dur_pred.float()) * dur_disturb
+    dur_pred = dur_pred * dur_disturb_r * dur_disturb_choice + \
+            dur_pred / dur_disturb_r * (1 - dur_disturb_choice)
+    dur_pred = torch.round(dur_pred * dur_alpha).clamp(0, 127)
+    if is_first:
+        dur_pred[:, 0] = 8
+    
+    dur_sum = dur_pred.sum()
+    npad = self.fm - dur_sum % self.fm
+    if npad < self.fm:
+        dur_pred[:, -1] += npad
+    mel2ph_pred = self.length_regulator(dur_pred).to(self.device)
+    return mel2ph_pred
+
+def prepare_inputs_for_dit(self, mel2ph_ref, mel2ph_pred, ph_ref, tone_ref, ph_pred, tone_pred, vae_latent):
+    # Prepare duration token 
+    mel2ph_pred = torch.cat((mel2ph_ref, mel2ph_pred+ph_ref.size(1)), dim=1)
+    mel2ph_pred = mel2ph_pred[:, :mel2ph_pred.size(1)//self.fm*self.fm].repeat(3, 1)
+    # Prepare phone and tone token
+    ph_pred = torch.cat((ph_ref, ph_pred), dim=1)
+    tone_pred = torch.cat((tone_ref, tone_pred), dim=1)
+    # Disable the English tone (set them to 3)"""
+    en_tone_idx = ~((tone_pred == 4) | ( (11 <= tone_pred) & (tone_pred <= 15)) | (tone_pred == 0))
+    tone_pred[en_tone_idx] = 3
+    
+    # Prepare cfg inputs
+    ph_seq = torch.cat([ph_pred, ph_pred, torch.full(ph_pred.size(), self.cfg_mask_token_phone, device=self.device)], 0)
+    tone_seq = torch.cat([tone_pred, tone_pred, torch.full(tone_pred.size(), self.cfg_mask_token_tone, device=self.device)], 0)
+    target_size = mel2ph_pred.size(1)//self.vae_stride
+    vae_latent_ = vae_latent.repeat(3, 1, 1)
+    ctx_mask = torch.ones_like(vae_latent_[:, :, 0:1])
+    vae_latent_ = F.pad(vae_latent_, (0, 0, 0, target_size - vae_latent.size(1)), mode='constant', value=0)
+    vae_latent_[1:] = 0.0
+    ctx_mask = F.pad(ctx_mask, (0, 0, 0, target_size - vae_latent.size(1)), mode='constant', value=0)
+
+    return {
+        'phone': ph_seq,
+        'tone': tone_seq,
+        "lat_ctx": vae_latent_ * ctx_mask,
+        "ctx_mask": ctx_mask,
+        "dur": mel2ph_pred,
+    }

xinference/thirdparty/megatts3/tts/gradio_api.py

@@ -0,0 +1,93 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# 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 multiprocessing as mp
+import torch
+import os
+from functools import partial
+import gradio as gr
+import traceback
+from tts.infer_cli import MegaTTS3DiTInfer, convert_to_wav, cut_wav
+
+
+def model_worker(input_queue, output_queue, device_id):
+    device = None
+    if device_id is not None:
+        device = torch.device(f'cuda:{device_id}')
+    infer_pipe = MegaTTS3DiTInfer(device=device)
+
+    while True:
+        task = input_queue.get()
+        inp_audio_path, inp_npy_path, inp_text, infer_timestep, p_w, t_w = task
+        try:
+            convert_to_wav(inp_audio_path)
+            wav_path = os.path.splitext(inp_audio_path)[0] + '.wav'
+            cut_wav(wav_path, max_len=28)
+            with open(wav_path, 'rb') as file:
+                file_content = file.read()
+            resource_context = infer_pipe.preprocess(file_content, latent_file=inp_npy_path)
+            wav_bytes = infer_pipe.forward(resource_context, inp_text, time_step=infer_timestep, p_w=p_w, t_w=t_w)
+            output_queue.put(wav_bytes)
+        except Exception as e:
+            traceback.print_exc()
+            print(task, str(e))
+            output_queue.put(None)
+
+
+def main(inp_audio, inp_npy, inp_text, infer_timestep, p_w, t_w, processes, input_queue, output_queue):
+    print("Push task to the inp queue |", inp_audio, inp_npy, inp_text, infer_timestep, p_w, t_w)
+    input_queue.put((inp_audio, inp_npy, inp_text, infer_timestep, p_w, t_w))
+    res = output_queue.get()
+    if res is not None:
+        return res
+    else:
+        print("")
+        return None
+
+
+if __name__ == '__main__':
+    mp.set_start_method('spawn', force=True)
+    mp_manager = mp.Manager()
+
+    devices = os.environ.get('CUDA_VISIBLE_DEVICES', '')
+    if devices != '':
+        devices = os.environ.get('CUDA_VISIBLE_DEVICES', '').split(",")
+    else:
+        devices = None
+    
+    num_workers = 1
+    input_queue = mp_manager.Queue()
+    output_queue = mp_manager.Queue()
+    processes = []
+
+    print("Start open workers")
+    for i in range(num_workers):
+        p = mp.Process(target=model_worker, args=(input_queue, output_queue, i % len(devices) if devices is not None else None))
+        p.start()
+        processes.append(p)
+
+    api_interface = gr.Interface(fn=
+                                partial(main, processes=processes, input_queue=input_queue, 
+                                        output_queue=output_queue), 
+                                inputs=[gr.Audio(type="filepath", label="Upload .wav"), gr.File(type="filepath", label="Upload .npy"), "text", 
+                                        gr.Number(label="infer timestep", value=32),
+                                        gr.Number(label="Intelligibility Weight", value=1.4),
+                                        gr.Number(label="Similarity Weight", value=3.0)], outputs=[gr.Audio(label="Synthesized Audio")],
+                                title="MegaTTS3",  
+                                description="Upload a speech clip as a reference for timbre, " +
+                                "upload the pre-extracted latent file, "+
+                                "input the target text, and receive the cloned voice.", concurrency_limit=1)
+    api_interface.launch(server_name='0.0.0.0', server_port=7929, debug=True)
+    for p in processes:
+        p.join()

xinference/thirdparty/megatts3/tts/infer_cli.py

@@ -0,0 +1,277 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# 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 json
+import os
+import argparse
+import librosa
+import numpy as np
+import torch
+
+from tn.chinese.normalizer import Normalizer as ZhNormalizer
+from tn.english.normalizer import Normalizer as EnNormalizer
+from langdetect import detect as classify_language
+from pydub import AudioSegment
+import pyloudnorm as pyln
+
+from tts.modules.ar_dur.commons.nar_tts_modules import LengthRegulator
+from tts.frontend_function import g2p, align, make_dur_prompt, dur_pred, prepare_inputs_for_dit
+from tts.utils.audio_utils.io import save_wav, to_wav_bytes, convert_to_wav_bytes, combine_audio_segments
+from tts.utils.commons.ckpt_utils import load_ckpt
+from tts.utils.commons.hparams import set_hparams, hparams
+from tts.utils.text_utils.text_encoder import TokenTextEncoder
+from tts.utils.text_utils.split_text import chunk_text_chinese, chunk_text_english
+from tts.utils.commons.hparams import hparams, set_hparams
+
+
+if "TOKENIZERS_PARALLELISM" not in os.environ:
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+def convert_to_wav(wav_path):
+    # Check if the file exists
+    if not os.path.exists(wav_path):
+        print(f"The file '{wav_path}' does not exist.")
+        return
+
+    # Check if the file already has a .wav extension
+    if not wav_path.endswith(".wav"):
+        # Define the output path with a .wav extension
+        out_path = os.path.splitext(wav_path)[0] + ".wav"
+
+        # Load the audio file using pydub and convert it to WAV
+        audio = AudioSegment.from_file(wav_path)
+        audio.export(out_path, format="wav")
+
+        print(f"Converted '{wav_path}' to '{out_path}'")
+
+
+def cut_wav(wav_path, max_len=28):
+    audio = AudioSegment.from_file(wav_path)
+    audio = audio[:int(max_len * 1000)]
+    audio.export(wav_path, format="wav")
+
+class MegaTTS3DiTInfer():
+    def __init__(
+            self, 
+            device=None,
+            ckpt_root='./checkpoints',
+            dit_exp_name='diffusion_transformer',
+            frontend_exp_name='aligner_lm',
+            wavvae_exp_name='wavvae',
+            dur_ckpt_path='duration_lm',
+            g2p_exp_name='g2p',
+            precision=torch.float16,
+            **kwargs
+        ):
+        self.sr = 24000
+        self.fm = 8
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        self.precision = precision
+
+        # build models
+        self.dit_exp_name = os.path.join(ckpt_root, dit_exp_name)
+        self.frontend_exp_name = os.path.join(ckpt_root, frontend_exp_name)
+        self.wavvae_exp_name = os.path.join(ckpt_root, wavvae_exp_name)
+        self.dur_exp_name = os.path.join(ckpt_root, dur_ckpt_path)
+        self.g2p_exp_name = os.path.join(ckpt_root, g2p_exp_name)
+        self.build_model(self.device)
+
+        # init text normalizer
+        self.zh_normalizer = ZhNormalizer(overwrite_cache=False, remove_erhua=False, remove_interjections=False)
+        self.en_normalizer = EnNormalizer(overwrite_cache=False)
+        # loudness meter
+        self.loudness_meter = pyln.Meter(self.sr)
+
+    def build_model(self, device):
+        set_hparams(exp_name=self.dit_exp_name, print_hparams=False)
+
+        ''' Load Dict '''
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        ling_dict = json.load(open(f"{current_dir}/utils/text_utils/dict.json", encoding='utf-8-sig'))
+        self.ling_dict = {k: TokenTextEncoder(None, vocab_list=ling_dict[k], replace_oov='<UNK>') for k in ['phone', 'tone']}
+        self.token_encoder = token_encoder = self.ling_dict['phone']
+        ph_dict_size = len(token_encoder)
+
+        ''' Load Duration LM '''
+        from tts.modules.ar_dur.ar_dur_predictor import ARDurPredictor
+        hp_dur_model = self.hp_dur_model = set_hparams(f'{self.dur_exp_name}/config.yaml', global_hparams=False)
+        hp_dur_model['frames_multiple'] = hparams['frames_multiple']
+        self.dur_model = ARDurPredictor(
+            hp_dur_model, hp_dur_model['dur_txt_hs'], hp_dur_model['dur_model_hidden_size'],
+            hp_dur_model['dur_model_layers'], ph_dict_size,
+            hp_dur_model['dur_code_size'],
+            use_rot_embed=hp_dur_model.get('use_rot_embed', False))
+        self.length_regulator = LengthRegulator()
+        load_ckpt(self.dur_model, f'{self.dur_exp_name}', 'dur_model')
+        self.dur_model.eval()
+        self.dur_model.to(device)
+
+        ''' Load Diffusion Transformer '''
+        from tts.modules.llm_dit.dit import Diffusion
+        self.dit = Diffusion()
+        load_ckpt(self.dit, f'{self.dit_exp_name}', 'dit', strict=False)
+        self.dit.eval()
+        self.dit.to(device)
+        self.cfg_mask_token_phone = 302 - 1
+        self.cfg_mask_token_tone = 32 - 1
+
+        ''' Load Frontend LM '''
+        from tts.modules.aligner.whisper_small import Whisper
+        self.aligner_lm = Whisper()
+        load_ckpt(self.aligner_lm, f'{self.frontend_exp_name}', 'model')
+        self.aligner_lm.eval()
+        self.aligner_lm.to(device)
+        self.kv_cache = None
+        self.hooks = None
+
+        ''' Load G2P LM'''
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        g2p_tokenizer = AutoTokenizer.from_pretrained(self.g2p_exp_name, padding_side="right")
+        g2p_tokenizer.padding_side = "right"
+        self.g2p_model = AutoModelForCausalLM.from_pretrained(self.g2p_exp_name).eval().to(device)
+        self.g2p_tokenizer = g2p_tokenizer
+        self.speech_start_idx = g2p_tokenizer.encode('<Reserved_TTS_0>')[0]
+
+        ''' Wav VAE '''
+        self.hp_wavvae = hp_wavvae = set_hparams(f'{self.wavvae_exp_name}/config.yaml', global_hparams=False)
+        from tts.modules.wavvae.decoder.wavvae_v3 import WavVAE_V3
+        self.wavvae = WavVAE_V3(hparams=hp_wavvae)
+        if os.path.exists(f'{self.wavvae_exp_name}/model_only_last.ckpt'):
+            load_ckpt(self.wavvae, f'{self.wavvae_exp_name}/model_only_last.ckpt', 'model_gen', strict=True)
+            self.has_vae_encoder = True
+        else:
+            load_ckpt(self.wavvae, f'{self.wavvae_exp_name}/decoder.ckpt', 'model_gen', strict=False)
+            self.has_vae_encoder = False
+        self.wavvae.eval()
+        self.wavvae.to(device)
+        self.vae_stride = hp_wavvae.get('vae_stride', 4)
+        self.hop_size = hp_wavvae.get('hop_size', 4)
+    
+    def preprocess(self, audio_bytes, latent_file=None, topk_dur=1, **kwargs):
+        wav_bytes = convert_to_wav_bytes(audio_bytes)
+
+        ''' Load wav '''
+        wav, _ = librosa.core.load(wav_bytes, sr=self.sr)
+        # Pad wav if necessary
+        ws = hparams['win_size']
+        if len(wav) % ws < ws - 1:
+            wav = np.pad(wav, (0, ws - 1 - (len(wav) % ws)), mode='constant', constant_values=0.0).astype(np.float32)
+        wav = np.pad(wav, (0, 12000), mode='constant', constant_values=0.0).astype(np.float32)
+        self.loudness_prompt = self.loudness_meter.integrated_loudness(wav.astype(float))
+
+        ''' obtain alignments with aligner_lm '''
+        ph_ref, tone_ref, mel2ph_ref = align(self, wav)
+
+        with torch.inference_mode():
+            ''' Forward WaveVAE to obtain: prompt latent '''
+            if self.has_vae_encoder:
+                wav = torch.FloatTensor(wav)[None].to(self.device)
+                vae_latent = self.wavvae.encode_latent(wav)
+                vae_latent = vae_latent[:, :mel2ph_ref.size(1)//4]
+            else:
+                assert latent_file is not None, "Please provide latent_file in WaveVAE decoder-only mode"
+                vae_latent = torch.from_numpy(np.load(latent_file)).to(self.device)
+                vae_latent = vae_latent[:, :mel2ph_ref.size(1)//4]
+        
+            ''' Duration Prompting '''
+            self.dur_model.hparams["infer_top_k"] = topk_dur if topk_dur > 1 else None
+            incremental_state_dur_prompt, ctx_dur_tokens = make_dur_prompt(self, mel2ph_ref, ph_ref, tone_ref)
+            
+        return {
+            'ph_ref': ph_ref,
+            'tone_ref': tone_ref,
+            'mel2ph_ref': mel2ph_ref,
+            'vae_latent': vae_latent,
+            'incremental_state_dur_prompt': incremental_state_dur_prompt,
+            'ctx_dur_tokens': ctx_dur_tokens,
+        }
+
+    def forward(self, resource_context, input_text, time_step, p_w, t_w, dur_disturb=0.1, dur_alpha=1.0, **kwargs):
+        device = self.device
+
+        ph_ref = resource_context['ph_ref'].to(device)
+        tone_ref = resource_context['tone_ref'].to(device)
+        mel2ph_ref = resource_context['mel2ph_ref'].to(device)
+        vae_latent = resource_context['vae_latent'].to(device)
+        ctx_dur_tokens = resource_context['ctx_dur_tokens'].to(device)
+        incremental_state_dur_prompt = resource_context['incremental_state_dur_prompt']
+
+        with torch.inference_mode():
+            ''' Generating '''
+            wav_pred_ = []
+            language_type = classify_language(input_text)
+            if language_type == 'en':
+                input_text = self.en_normalizer.normalize(input_text)
+                text_segs = chunk_text_english(input_text, max_chars=130)
+            else:
+                input_text = self.zh_normalizer.normalize(input_text)
+                text_segs = chunk_text_chinese(input_text, limit=60)
+
+            for seg_i, text in enumerate(text_segs):
+                ''' G2P '''
+                ph_pred, tone_pred = g2p(self, text)
+
+                ''' Duration Prediction '''
+                mel2ph_pred = dur_pred(self, ctx_dur_tokens, incremental_state_dur_prompt, ph_pred, tone_pred, seg_i, dur_disturb, dur_alpha, is_first=seg_i==0, is_final=seg_i==len(text_segs)-1)
+                
+                inputs = prepare_inputs_for_dit(self, mel2ph_ref, mel2ph_pred, ph_ref, tone_ref, ph_pred, tone_pred, vae_latent)
+                # Speech dit inference
+                with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):
+                    x = self.dit.inference(inputs, timesteps=time_step, seq_cfg_w=[p_w, t_w]).float()
+                
+                # WavVAE decode
+                x[:, :vae_latent.size(1)] = vae_latent
+                wav_pred = self.wavvae.decode(x)[0,0].to(torch.float32)
+                
+                ''' Post-processing '''
+                # Trim prompt wav
+                wav_pred = wav_pred[vae_latent.size(1)*self.vae_stride*self.hop_size:].cpu().numpy()
+                # Norm generated wav to prompt wav's level
+                meter = pyln.Meter(self.sr)  # create BS.1770 meter
+                loudness_pred = self.loudness_meter.integrated_loudness(wav_pred.astype(float))
+                wav_pred = pyln.normalize.loudness(wav_pred, loudness_pred, self.loudness_prompt)
+                if np.abs(wav_pred).max() >= 1:
+                    wav_pred = wav_pred / np.abs(wav_pred).max() * 0.95
+
+                # Apply hamming window
+                wav_pred_.append(wav_pred)
+
+            return combine_audio_segments(wav_pred_, sr=self.sr).astype(float)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_wav', type=str)
+    parser.add_argument('--input_text', type=str)
+    parser.add_argument('--output_dir', type=str)
+    parser.add_argument('--time_step', type=int, default=32, help='Inference steps of Diffusion Transformer')
+    parser.add_argument('--p_w', type=float, default=1.6, help='Intelligibility Weight')
+    parser.add_argument('--t_w', type=float, default=2.5, help='Similarity Weight')
+    args = parser.parse_args()
+    wav_path, input_text, out_path, time_step, p_w, t_w = args.input_wav, args.input_text, args.output_dir, args.time_step, args.p_w, args.t_w
+
+    infer_ins = MegaTTS3DiTInfer()
+
+    with open(wav_path, 'rb') as file:
+        file_content = file.read()
+
+    print(f"| Start processing {wav_path}+{input_text}")
+    resource_context = infer_ins.preprocess(file_content, latent_file=wav_path.replace('.wav', '.npy'))
+    wav_bytes = infer_ins.forward(resource_context, input_text, time_step=time_step, p_w=p_w, t_w=t_w)
+
+    print(f"| Saving results to {out_path}/[P]{input_text[:20]}.wav")
+    os.makedirs(out_path, exist_ok=True)
+    save_wav(wav_bytes, f'{out_path}/[P]{input_text[:20]}.wav')