init

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,37 @@
+# general things to ignore
+.DS_Store
+build/
+build_contrib/
+dist/
+.cache/
+*.egg-info/
+*.egg
+*.py[cod]
+__pycache__/
+*.so
+*~
+
+# IDE
+.vscode/
+
+# misc
+checkpoints/
+test_waves/
+reconstructed/
+.python-version
+ruff.log
+/configs/inuse/
+*.wav
+*.ogg
+*.mp3
+demo_dir/*
+*.pt 
+*.json
+*.txt
+*.ipynb
+asset/
+*.csv
+*.xlsx
+*.jpg
+*.log
+*.pth

Dockerfile

@@ -0,0 +1,30 @@
+# 第一阶段：构建 OpenVoice 环境（锁定 numpy==1.22.0）
+FROM python:3.9-slim AS openvoice-builder
+
+# 安装系统依赖
+RUN apt-get update && apt-get install -y \
+    ffmpeg libsndfile1 git \
+ && rm -rf /var/lib/apt/lists/*
+
+# 强制安装 numpy==1.22.0 并构建 OpenVoice
+WORKDIR /app
+COPY OpenVoice/ ./OpenVoice
+RUN pip install numpy==1.22.0 && \
+    pip install -e ./OpenVoice
+
+
+# 第二阶段：构建主程序环境（允许 numpy>=2.0.0）
+FROM python:3.9-slim AS main-app
+
+# 安装主程序依赖（隔离于 OpenVoice 的环境）
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# 从第一阶段复制已编译的 OpenVoice 和兼容库
+COPY --from=openvoice-builder /usr/local/lib/python3.9/site-packages /usr/local/lib/python3.9/site-packages
+COPY --from=openvoice-builder /app/OpenVoice /app/OpenVoice
+
+# 暴露端口并启动
+EXPOSE 7860
+CMD ["python", "app.py"]

OpenVoice/.gitignore

@@ -0,0 +1,13 @@
+__pycache__/
+.ipynb_checkpoints/
+processed
+outputs
+outputs_v2
+checkpoints
+checkpoints_v2
+trash
+examples*
+.env
+build
+*.egg-info/
+*.zip

OpenVoice/LICENSE

@@ -0,0 +1,7 @@
+Copyright 2024 MyShell.ai
+
+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.

OpenVoice/README.md

@@ -0,0 +1,70 @@
+<div align="center">
+  <div>&nbsp;</div>
+  <img src="resources/openvoicelogo.jpg" width="400"/> 
+
+[Paper](https://arxiv.org/abs/2312.01479) |
+[Website](https://research.myshell.ai/open-voice) <br> <br>
+<a href="https://trendshift.io/repositories/6161" target="_blank"><img src="https://trendshift.io/api/badge/repositories/6161" alt="myshell-ai%2FOpenVoice | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
+</div>
+
+## Introduction
+
+### OpenVoice V1
+
+As we detailed in our [paper](https://arxiv.org/abs/2312.01479) and [website](https://research.myshell.ai/open-voice), the advantages of OpenVoice are three-fold:
+
+**1. Accurate Tone Color Cloning.**
+OpenVoice can accurately clone the reference tone color and generate speech in multiple languages and accents.
+
+**2. Flexible Voice Style Control.**
+OpenVoice enables granular control over voice styles, such as emotion and accent, as well as other style parameters including rhythm, pauses, and intonation. 
+
+**3. Zero-shot Cross-lingual Voice Cloning.**
+Neither of the language of the generated speech nor the language of the reference speech needs to be presented in the massive-speaker multi-lingual training dataset.
+
+### OpenVoice V2
+
+In April 2024, we released OpenVoice V2, which includes all features in V1 and has:
+
+**1. Better Audio Quality.**
+OpenVoice V2 adopts a different training strategy that delivers better audio quality.
+
+**2. Native Multi-lingual Support.**
+English, Spanish, French, Chinese, Japanese and Korean are natively supported in OpenVoice V2.
+
+**3. Free Commercial Use.**
+Starting from April 2024, both V2 and V1 are released under MIT License. Free for commercial use.
+
+[Video](https://github.com/myshell-ai/OpenVoice/assets/40556743/3cba936f-82bf-476c-9e52-09f0f417bb2f)
+
+OpenVoice has been powering the instant voice cloning capability of [myshell.ai](https://app.myshell.ai/explore) since May 2023. Until Nov 2023, the voice cloning model has been used tens of millions of times by users worldwide, and witnessed the explosive user growth on the platform.
+
+## Main Contributors
+
+- [Zengyi Qin](https://www.qinzy.tech) at MIT
+- [Wenliang Zhao](https://wl-zhao.github.io) at Tsinghua University
+- [Xumin Yu](https://yuxumin.github.io) at Tsinghua University
+- [Ethan Sun](https://twitter.com/ethan_myshell) at MyShell
+
+## How to Use
+Please see [usage](docs/USAGE.md) for detailed instructions.
+
+## Common Issues
+
+Please see [QA](docs/QA.md) for common questions and answers. We will regularly update the question and answer list.
+
+## Citation
+```
+@article{qin2023openvoice,
+  title={OpenVoice: Versatile Instant Voice Cloning},
+  author={Qin, Zengyi and Zhao, Wenliang and Yu, Xumin and Sun, Xin},
+  journal={arXiv preprint arXiv:2312.01479},
+  year={2023}
+}
+```
+
+## License
+OpenVoice V1 and V2 are MIT Licensed. Free for both commercial and research use.
+
+## Acknowledgements
+This implementation is based on several excellent projects, [TTS](https://github.com/coqui-ai/TTS), [VITS](https://github.com/jaywalnut310/vits), and [VITS2](https://github.com/daniilrobnikov/vits2). Thanks for their awesome work!

OpenVoice/docs/QA.md

@@ -0,0 +1,39 @@
+# Common Questions and Answers
+
+## General Comments
+
+**OpenVoice is a Technology, not a Product**
+
+Although it works on a majority of voices if used correctly, please do not expect it to work perfectly on every case, as it takes a lot of engineering effort to translate a technology to a stable product. The targeted users of this technology are developers and researchers, not end users. End users expects a perfect product. However, we are confident to say that OpenVoice is the state-of-the-art among the source-available voice cloning technologies.
+
+The contribution of OpenVoice is a versatile instant voice cloning technical approach, not a ready-to-use perfect voice cloning product. However, we firmly believe that by releasing OpenVoice, we can accelerate the open research community's progress on instant voice cloning, and someday in the future the free voice cloning methods will be as good as commercial ones.
+
+## Issues with Voice Quality
+
+**Accent and Emotion of the Generated Voice is not Similar to the Reference Voice**
+
+First of all, OpenVoice only clones the tone color of the reference speaker. It does NOT clone the accent or emotion. The accent and emotion is controlled by the base speaker TTS model, not cloned by the tone color converter (please refer to our [paper](https://arxiv.org/pdf/2312.01479.pdf) for technical details). If the user wants to change the accent or emotion of the output, they need to have a base speaker model with that accent. OpenVoice provides sufficient flexibility for users to integrate their own base speaker model into the framework by simply replacing the current base speaker we provided.
+
+**Bad Audio Quality of the Generated Speech** 
+
+Please check the followings:
+- Is your reference audio is clean enough without any background noise? You can find some high-quality reference speech [here](https://aiartes.com/voiceai)
+- Is your audio too short?
+- Does your audio contain speech from more than one person?
+- Does the reference audio contain long blank sections?
+- Did you name the reference audio the same name you used before but forgot to delete the `processed` folder?
+
+## Issues with Languages
+
+**Support of Other Languages**
+
+For multi-lingual and cross-lingual usage, please refer to [`demo_part2.ipynb`](https://github.com/myshell-ai/OpenVoice/blob/main/demo_part2.ipynb). OpenVoice supports any language as long as you have a base speaker in that language. The OpenVoice team already did the most difficult part (tone color converter training) for you. Base speaker TTS model is relatively easy to train, and multiple existing open-source repositories support it. If you don't want to train by yourself, simply use the OpenAI TTS model as the base speaker.
+
+## Issues with Installation
+**Error Related to Silero**
+
+When calling `get_vad_segments` from `se_extractor.py`, there should be a message like this:
+```
+Downloading: "https://github.com/snakers4/silero-vad/zipball/master" to /home/user/.cache/torch/hub/master.zip
+```
+The download would fail if your machine can not access github. Please download the zip from "https://github.com/snakers4/silero-vad/zipball/master" manually and unzip it to `/home/user/.cache/torch/hub/snakers4_silero-vad_master`. You can also see [this issue](https://github.com/myshell-ai/OpenVoice/issues/57) for solutions for other versions of silero.

OpenVoice/docs/USAGE.md

@@ -0,0 +1,83 @@
+# Usage
+
+## Table of Content
+
+- [Quick Use](#quick-use): directly use OpenVoice without installation.
+- [Linux Install](#linux-install): for researchers and developers only.
+    - [V1](#openvoice-v1)
+    - [V2](#openvoice-v2)
+- [Install on Other Platforms](#install-on-other-platforms): unofficial installation guide contributed by the community
+
+## Quick Use
+
+The input speech audio of OpenVoice can be in **Any Language**. OpenVoice can clone the voice in that speech audio, and use the voice to speak in multiple languages. For quick use, we recommend you to try the already deployed services:
+
+- [British English](https://app.myshell.ai/widget/vYjqae)
+- [American English](https://app.myshell.ai/widget/nEFFJf)
+- [Indian English](https://app.myshell.ai/widget/V3iYze)
+- [Australian English](https://app.myshell.ai/widget/fM7JVf)
+- [Spanish](https://app.myshell.ai/widget/NNFFVz)
+- [French](https://app.myshell.ai/widget/z2uyUz)
+- [Chinese](https://app.myshell.ai/widget/fU7nUz)
+- [Japanese](https://app.myshell.ai/widget/IfIB3u)
+- [Korean](https://app.myshell.ai/widget/q6ZjIn)
+
+## Minimal Demo
+
+For users who want to quickly try OpenVoice and do not require high quality or stability, click any of the following links:
+
+<div align="center">
+    <a href="https://app.myshell.ai/bot/z6Bvua/1702636181"><img src="../resources/myshell-hd.png" height="28"></a>
+    &nbsp;&nbsp;&nbsp;&nbsp;
+    <a href="https://huggingface.co/spaces/myshell-ai/OpenVoice"><img src="../resources/huggingface.png" height="32"></a>
+</div>
+
+## Linux Install
+
+This section is only for developers and researchers who are familiar with Linux, Python and PyTorch. Clone this repo, and run
+
+```
+conda create -n openvoice python=3.9
+conda activate openvoice
+git clone git@github.com:myshell-ai/OpenVoice.git
+cd OpenVoice
+pip install -e .
+```
+
+No matter if you are using V1 or V2, the above installation is the same.
+
+### OpenVoice V1
+
+Download the checkpoint from [here](https://myshell-public-repo-host.s3.amazonaws.com/openvoice/checkpoints_1226.zip) and extract it to the `checkpoints` folder.
+
+**1. Flexible Voice Style Control.**
+Please see [`demo_part1.ipynb`](../demo_part1.ipynb) for an example usage of how OpenVoice enables flexible style control over the cloned voice.
+
+**2. Cross-Lingual Voice Cloning.**
+Please see [`demo_part2.ipynb`](../demo_part2.ipynb) for an example for languages seen or unseen in the MSML training set.
+
+**3. Gradio Demo.**. We provide a minimalist local gradio demo here. We strongly suggest the users to look into `demo_part1.ipynb`, `demo_part2.ipynb` and the [QnA](QA.md) if they run into issues with the gradio demo. Launch a local gradio demo with `python -m openvoice_app --share`.
+
+### OpenVoice V2
+
+Download the checkpoint from [here](https://myshell-public-repo-host.s3.amazonaws.com/openvoice/checkpoints_v2_0417.zip) and extract it to the `checkpoints_v2` folder.
+
+Install [MeloTTS](https://github.com/myshell-ai/MeloTTS):
+```
+pip install git+https://github.com/myshell-ai/MeloTTS.git
+python -m unidic download
+```
+
+**Demo Usage.** Please see [`demo_part3.ipynb`](../demo_part3.ipynb) for example usage of OpenVoice V2. Now it natively supports English, Spanish, French, Chinese, Japanese and Korean.
+
+
+## Install on Other Platforms
+
+This section provides the unofficial installation guides by open-source contributors in the community:
+
+- Windows
+  - [Guide](https://github.com/Alienpups/OpenVoice/blob/main/docs/USAGE_WINDOWS.md) by [@Alienpups](https://github.com/Alienpups)
+  - You are welcome to contribute if you have a better installation guide. We will list you here.
+- Docker
+  - [Guide](https://github.com/StevenJSCF/OpenVoice/blob/update-docs/docs/DF_USAGE.md) by [@StevenJSCF](https://github.com/StevenJSCF)
+  - You are welcome to contribute if you have a better installation guide. We will list you here.

OpenVoice/openvoice/api.py

@@ -0,0 +1,202 @@
+import torch
+import numpy as np
+import re
+import soundfile
+from openvoice import utils
+from openvoice import commons
+import os
+import librosa
+from openvoice.text import text_to_sequence
+from openvoice.mel_processing import spectrogram_torch
+from openvoice.models import SynthesizerTrn
+
+
+class OpenVoiceBaseClass(object):
+    def __init__(self, 
+                config_path, 
+                device='cuda:0'):
+        if 'cuda' in device:
+            assert torch.cuda.is_available()
+
+        hps = utils.get_hparams_from_file(config_path)
+
+        model = SynthesizerTrn(
+            len(getattr(hps, 'symbols', [])),
+            hps.data.filter_length // 2 + 1,
+            n_speakers=hps.data.n_speakers,
+            **hps.model,
+        ).to(device)
+
+        model.eval()
+        self.model = model
+        self.hps = hps
+        self.device = device
+
+    def load_ckpt(self, ckpt_path):
+        checkpoint_dict = torch.load(ckpt_path, map_location=torch.device(self.device))
+        a, b = self.model.load_state_dict(checkpoint_dict['model'], strict=False)
+        print("Loaded checkpoint '{}'".format(ckpt_path))
+        print('missing/unexpected keys:', a, b)
+
+
+class BaseSpeakerTTS(OpenVoiceBaseClass):
+    language_marks = {
+        "english": "EN",
+        "chinese": "ZH",
+    }
+
+    @staticmethod
+    def get_text(text, hps, is_symbol):
+        text_norm = text_to_sequence(text, hps.symbols, [] if is_symbol else hps.data.text_cleaners)
+        if hps.data.add_blank:
+            text_norm = commons.intersperse(text_norm, 0)
+        text_norm = torch.LongTensor(text_norm)
+        return text_norm
+
+    @staticmethod
+    def audio_numpy_concat(segment_data_list, sr, speed=1.):
+        audio_segments = []
+        for segment_data in segment_data_list:
+            audio_segments += segment_data.reshape(-1).tolist()
+            audio_segments += [0] * int((sr * 0.05)/speed)
+        audio_segments = np.array(audio_segments).astype(np.float32)
+        return audio_segments
+
+    @staticmethod
+    def split_sentences_into_pieces(text, language_str):
+        texts = utils.split_sentence(text, language_str=language_str)
+        print(" > Text splitted to sentences.")
+        print('\n'.join(texts))
+        print(" > ===========================")
+        return texts
+
+    def tts(self, text, output_path, speaker, language='English', speed=1.0):
+        mark = self.language_marks.get(language.lower(), None)
+        assert mark is not None, f"language {language} is not supported"
+
+        texts = self.split_sentences_into_pieces(text, mark)
+
+        audio_list = []
+        for t in texts:
+            t = re.sub(r'([a-z])([A-Z])', r'\1 \2', t)
+            t = f'[{mark}]{t}[{mark}]'
+            stn_tst = self.get_text(t, self.hps, False)
+            device = self.device
+            speaker_id = self.hps.speakers[speaker]
+            with torch.no_grad():
+                x_tst = stn_tst.unsqueeze(0).to(device)
+                x_tst_lengths = torch.LongTensor([stn_tst.size(0)]).to(device)
+                sid = torch.LongTensor([speaker_id]).to(device)
+                audio = self.model.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=0.667, noise_scale_w=0.6,
+                                    length_scale=1.0 / speed)[0][0, 0].data.cpu().float().numpy()
+            audio_list.append(audio)
+        audio = self.audio_numpy_concat(audio_list, sr=self.hps.data.sampling_rate, speed=speed)
+
+        if output_path is None:
+            return audio
+        else:
+            soundfile.write(output_path, audio, self.hps.data.sampling_rate)
+
+
+class ToneColorConverter(OpenVoiceBaseClass):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if kwargs.get('enable_watermark', True):
+            import wavmark
+            self.watermark_model = wavmark.load_model().to(self.device)
+        else:
+            self.watermark_model = None
+        self.version = getattr(self.hps, '_version_', "v1")
+
+
+
+    def extract_se(self, ref_wav_list, se_save_path=None):
+        if isinstance(ref_wav_list, str):
+            ref_wav_list = [ref_wav_list]
+        
+        device = self.device
+        hps = self.hps
+        gs = []
+        
+        for fname in ref_wav_list:
+            audio_ref, sr = librosa.load(fname, sr=hps.data.sampling_rate)
+            y = torch.FloatTensor(audio_ref)
+            y = y.to(device)
+            y = y.unsqueeze(0)
+            y = spectrogram_torch(y, hps.data.filter_length,
+                                        hps.data.sampling_rate, hps.data.hop_length, hps.data.win_length,
+                                        center=False).to(device)
+            with torch.no_grad():
+                g = self.model.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+                gs.append(g.detach())
+        gs = torch.stack(gs).mean(0)
+
+        if se_save_path is not None:
+            os.makedirs(os.path.dirname(se_save_path), exist_ok=True)
+            torch.save(gs.cpu(), se_save_path)
+
+        return gs
+
+    def convert(self, audio_src_path, src_se, tgt_se, output_path=None, tau=0.3, message="default"):
+        hps = self.hps
+        # load audio
+        audio, sample_rate = librosa.load(audio_src_path, sr=hps.data.sampling_rate)
+        audio = torch.tensor(audio).float()
+        
+        with torch.no_grad():
+            y = torch.FloatTensor(audio).to(self.device)
+            y = y.unsqueeze(0)
+            spec = spectrogram_torch(y, hps.data.filter_length,
+                                    hps.data.sampling_rate, hps.data.hop_length, hps.data.win_length,
+                                    center=False).to(self.device)
+            spec_lengths = torch.LongTensor([spec.size(-1)]).to(self.device)
+            audio = self.model.voice_conversion(spec, spec_lengths, sid_src=src_se, sid_tgt=tgt_se, tau=tau)[0][
+                        0, 0].data.cpu().float().numpy()
+            audio = self.add_watermark(audio, message)
+            if output_path is None:
+                return audio
+            else:
+                soundfile.write(output_path, audio, hps.data.sampling_rate)
+    
+    def add_watermark(self, audio, message):
+        if self.watermark_model is None:
+            return audio
+        device = self.device
+        bits = utils.string_to_bits(message).reshape(-1)
+        n_repeat = len(bits) // 32
+
+        K = 16000
+        coeff = 2
+        for n in range(n_repeat):
+            trunck = audio[(coeff * n) * K: (coeff * n + 1) * K]
+            if len(trunck) != K:
+                print('Audio too short, fail to add watermark')
+                break
+            message_npy = bits[n * 32: (n + 1) * 32]
+            
+            with torch.no_grad():
+                signal = torch.FloatTensor(trunck).to(device)[None]
+                message_tensor = torch.FloatTensor(message_npy).to(device)[None]
+                signal_wmd_tensor = self.watermark_model.encode(signal, message_tensor)
+                signal_wmd_npy = signal_wmd_tensor.detach().cpu().squeeze()
+            audio[(coeff * n) * K: (coeff * n + 1) * K] = signal_wmd_npy
+        return audio
+
+    def detect_watermark(self, audio, n_repeat):
+        bits = []
+        K = 16000
+        coeff = 2
+        for n in range(n_repeat):
+            trunck = audio[(coeff * n) * K: (coeff * n + 1) * K]
+            if len(trunck) != K:
+                print('Audio too short, fail to detect watermark')
+                return 'Fail'
+            with torch.no_grad():
+                signal = torch.FloatTensor(trunck).to(self.device).unsqueeze(0)
+                message_decoded_npy = (self.watermark_model.decode(signal) >= 0.5).int().detach().cpu().numpy().squeeze()
+            bits.append(message_decoded_npy)
+        bits = np.stack(bits).reshape(-1, 8)
+        message = utils.bits_to_string(bits)
+        return message
+    

OpenVoice/openvoice/attentions.py

@@ -0,0 +1,465 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from openvoice import commons
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        window_size=4,
+        isflow=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        # if isflow:
+        #  cond_layer = torch.nn.Conv1d(256, 2*hidden_channels*n_layers, 1)
+        #  self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1)
+        #  self.cond_layer = weight_norm(cond_layer, name='weight')
+        #  self.gin_channels = 256
+        self.cond_layer_idx = self.n_layers
+        if "gin_channels" in kwargs:
+            self.gin_channels = kwargs["gin_channels"]
+            if self.gin_channels != 0:
+                self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
+                # vits2 says 3rd block, so idx is 2 by default
+                self.cond_layer_idx = (
+                    kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2
+                )
+                # logging.debug(self.gin_channels, self.cond_layer_idx)
+                assert (
+                    self.cond_layer_idx < self.n_layers
+                ), "cond_layer_idx should be less than n_layers"
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+
+        for i in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, g=None):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            if i == self.cond_layer_idx and g is not None:
+                g = self.spk_emb_linear(g.transpose(1, 2))
+                g = g.transpose(1, 2)
+                x = x + g
+                x = x * x_mask
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(LayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                    causal=True,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, h, h_mask):
+        """
+        x: decoder input
+        h: encoder output
+        """
+        self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        n_heads,
+        p_dropout=0.0,
+        window_size=None,
+        heads_share=True,
+        block_length=None,
+        proximal_bias=False,
+        proximal_init=False,
+    ):
+        super().__init__()
+        assert channels % n_heads == 0
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
+
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                self.conv_k.bias.copy_(self.conv_q.bias)
+
+    def forward(self, x, c, attn_mask=None):
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask=None):
+        # reshape [b, d, t] -> [b, n_h, t, d_k]
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            )
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)  # [b, n_h, t_t, t_s]
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_v, t_s
+            )
+            output = output + self._matmul_with_relative_values(
+                relative_weights, value_relative_embeddings
+            )
+        output = (
+            output.transpose(2, 3).contiguous().view(b, d, t_t)
+        )  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+        return output, p_attn
+
+    def _matmul_with_relative_values(self, x, y):
+        """
+        x: [b, h, l, m]
+        y: [h or 1, m, d]
+        ret: [b, h, l, d]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
+
+    def _matmul_with_relative_keys(self, x, y):
+        """
+        x: [b, h, l, d]
+        y: [h or 1, m, d]
+        ret: [b, h, l, m]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        return ret
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        2 * self.window_size + 1
+        # Pad first before slice to avoid using cond ops.
+        pad_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
+
+    def _relative_position_to_absolute_position(self, x):
+        """
+        x: [b, h, l, 2*l-1]
+        ret: [b, h, l, l]
+        """
+        batch, heads, length, _ = x.size()
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
+
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(
+            x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [0, length - 1]])
+        )
+
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
+            :, :, :length, length - 1 :
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        """
+        x: [b, h, l, l]
+        ret: [b, h, l, 2*l-1]
+        """
+        batch, heads, length, _ = x.size()
+        # pad along column
+        x = F.pad(
+            x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
+        )
+        x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
+        # add 0's in the beginning that will skew the elements after reshape
+        x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length):
+        """Bias for self-attention to encourage attention to close positions.
+        Args:
+          length: an integer scalar.
+        Returns:
+          a Tensor with shape [1, 1, length, length]
+        """
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        filter_channels,
+        kernel_size,
+        p_dropout=0.0,
+        activation=None,
+        causal=False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.activation = activation
+        self.causal = causal
+
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
+
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding(x * x_mask))
+        if self.activation == "gelu":
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
+
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x

OpenVoice/openvoice/commons.py

@@ -0,0 +1,160 @@
+import math
+import torch
+from torch.nn import functional as F
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def convert_pad_shape(pad_shape):
+    layer = pad_shape[::-1]
+    pad_shape = [item for sublist in layer for item in sublist]
+    return pad_shape
+
+
+def intersperse(lst, item):
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+    """KL(P||Q)"""
+    kl = (logs_q - logs_p) - 0.5
+    kl += (
+        0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
+    )
+    return kl
+
+
+def rand_gumbel(shape):
+    """Sample from the Gumbel distribution, protect from overflows."""
+    uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+    return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+    g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+    return g
+
+
+def slice_segments(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = ids_str[i]
+        idx_end = idx_str + segment_size
+        ret[i] = x[i, :, idx_str:idx_end]
+    return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
+    position = torch.arange(length, dtype=torch.float)
+    num_timescales = channels // 2
+    log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
+        num_timescales - 1
+    )
+    inv_timescales = min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
+    )
+    scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+    signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+    signal = F.pad(signal, [0, 0, 0, channels % 2])
+    signal = signal.view(1, channels, length)
+    return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def convert_pad_shape(pad_shape):
+    layer = pad_shape[::-1]
+    pad_shape = [item for sublist in layer for item in sublist]
+    return pad_shape
+
+
+def shift_1d(x):
+    x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+    return x
+
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def generate_path(duration, mask):
+    """
+    duration: [b, 1, t_x]
+    mask: [b, 1, t_y, t_x]
+    """
+
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm

OpenVoice/openvoice/mel_processing.py

@@ -0,0 +1,183 @@
+import torch
+import torch.utils.data
+from librosa.filters import mel as librosa_mel_fn
+
+MAX_WAV_VALUE = 32768.0
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    """
+    PARAMS
+    ------
+    C: compression factor
+    """
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    """
+    PARAMS
+    ------
+    C: compression factor used to compress
+    """
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+    output = dynamic_range_decompression_torch(magnitudes)
+    return output
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
+    if torch.min(y) < -1.1:
+        print("min value is ", torch.min(y))
+    if torch.max(y) > 1.1:
+        print("max value is ", torch.max(y))
+
+    global hann_window
+    dtype_device = str(y.dtype) + "_" + str(y.device)
+    wnsize_dtype_device = str(win_size) + "_" + dtype_device
+    if wnsize_dtype_device not in hann_window:
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
+            dtype=y.dtype, device=y.device
+        )
+
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
+    y = y.squeeze(1)
+
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window[wnsize_dtype_device],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=False,
+    )
+
+    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
+    return spec
+
+
+def spectrogram_torch_conv(y, n_fft, sampling_rate, hop_size, win_size, center=False):
+    # if torch.min(y) < -1.:
+    #     print('min value is ', torch.min(y))
+    # if torch.max(y) > 1.:
+    #     print('max value is ', torch.max(y))
+
+    global hann_window
+    dtype_device = str(y.dtype) + '_' + str(y.device)
+    wnsize_dtype_device = str(win_size) + '_' + dtype_device
+    if wnsize_dtype_device not in hann_window:
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)
+
+    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
+    
+    # ******************** original ************************#
+    # y = y.squeeze(1)
+    # spec1 = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
+    #                   center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
+
+    # ******************** ConvSTFT ************************#
+    freq_cutoff = n_fft // 2 + 1
+    fourier_basis = torch.view_as_real(torch.fft.fft(torch.eye(n_fft)))
+    forward_basis = fourier_basis[:freq_cutoff].permute(2, 0, 1).reshape(-1, 1, fourier_basis.shape[1])
+    forward_basis = forward_basis * torch.as_tensor(librosa.util.pad_center(torch.hann_window(win_size), size=n_fft)).float()
+
+    import torch.nn.functional as F
+
+    # if center:
+    #     signal = F.pad(y[:, None, None, :], (n_fft // 2, n_fft // 2, 0, 0), mode = 'reflect').squeeze(1)
+    assert center is False
+
+    forward_transform_squared = F.conv1d(y, forward_basis.to(y.device), stride = hop_size)
+    spec2 = torch.stack([forward_transform_squared[:, :freq_cutoff, :], forward_transform_squared[:, freq_cutoff:, :]], dim = -1)
+
+
+    # ******************** Verification ************************#
+    spec1 = torch.stft(y.squeeze(1), n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
+                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
+    assert torch.allclose(spec1, spec2, atol=1e-4)
+
+    spec = torch.sqrt(spec2.pow(2).sum(-1) + 1e-6)
+    return spec
+
+
+def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
+    global mel_basis
+    dtype_device = str(spec.dtype) + "_" + str(spec.device)
+    fmax_dtype_device = str(fmax) + "_" + dtype_device
+    if fmax_dtype_device not in mel_basis:
+        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
+        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
+            dtype=spec.dtype, device=spec.device
+        )
+    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
+    spec = spectral_normalize_torch(spec)
+    return spec
+
+
+def mel_spectrogram_torch(
+    y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
+):
+    if torch.min(y) < -1.0:
+        print("min value is ", torch.min(y))
+    if torch.max(y) > 1.0:
+        print("max value is ", torch.max(y))
+
+    global mel_basis, hann_window
+    dtype_device = str(y.dtype) + "_" + str(y.device)
+    fmax_dtype_device = str(fmax) + "_" + dtype_device
+    wnsize_dtype_device = str(win_size) + "_" + dtype_device
+    if fmax_dtype_device not in mel_basis:
+        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
+        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
+            dtype=y.dtype, device=y.device
+        )
+    if wnsize_dtype_device not in hann_window:
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
+            dtype=y.dtype, device=y.device
+        )
+
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
+    y = y.squeeze(1)
+
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window[wnsize_dtype_device],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=False,
+    )
+
+    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
+
+    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
+    spec = spectral_normalize_torch(spec)
+
+    return spec

OpenVoice/openvoice/models.py

@@ -0,0 +1,499 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from openvoice import commons
+from openvoice import modules
+from openvoice import attentions
+
+from torch.nn import Conv1d, ConvTranspose1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+
+from openvoice.commons import init_weights, get_padding
+
+
+class TextEncoder(nn.Module):
+	def __init__(self,
+			n_vocab,
+			out_channels,
+			hidden_channels,
+			filter_channels,
+			n_heads,
+			n_layers,
+			kernel_size,
+			p_dropout):
+		super().__init__()
+		self.n_vocab = n_vocab
+		self.out_channels = out_channels
+		self.hidden_channels = hidden_channels
+		self.filter_channels = filter_channels
+		self.n_heads = n_heads
+		self.n_layers = n_layers
+		self.kernel_size = kernel_size
+		self.p_dropout = p_dropout
+
+		self.emb = nn.Embedding(n_vocab, hidden_channels)
+		nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+
+		self.encoder = attentions.Encoder(
+			hidden_channels,
+			filter_channels,
+			n_heads,
+			n_layers,
+			kernel_size,
+			p_dropout)
+		self.proj= nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+	def forward(self, x, x_lengths):
+		x = self.emb(x) * math.sqrt(self.hidden_channels) # [b, t, h]
+		x = torch.transpose(x, 1, -1) # [b, h, t]
+		x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+
+		x = self.encoder(x * x_mask, x_mask)
+		stats = self.proj(x) * x_mask
+
+		m, logs = torch.split(stats, self.out_channels, dim=1)
+		return x, m, logs, x_mask
+     
+
+class DurationPredictor(nn.Module):
+    def __init__(
+        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+
+        self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(
+            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_1 = modules.LayerNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_2 = modules.LayerNorm(filter_channels)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+    def forward(self, x, x_mask, g=None):
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask
+               
+class StochasticDurationPredictor(nn.Module):
+	def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):
+		super().__init__()
+		filter_channels = in_channels # it needs to be removed from future version.
+		self.in_channels = in_channels
+		self.filter_channels = filter_channels
+		self.kernel_size = kernel_size
+		self.p_dropout = p_dropout
+		self.n_flows = n_flows
+		self.gin_channels = gin_channels
+
+		self.log_flow = modules.Log()
+		self.flows = nn.ModuleList()
+		self.flows.append(modules.ElementwiseAffine(2))
+		for i in range(n_flows):
+			self.flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+			self.flows.append(modules.Flip())
+
+		self.post_pre = nn.Conv1d(1, filter_channels, 1)
+		self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+		self.post_convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
+		self.post_flows = nn.ModuleList()
+		self.post_flows.append(modules.ElementwiseAffine(2))
+		for i in range(4):
+			self.post_flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+			self.post_flows.append(modules.Flip())
+
+		self.pre = nn.Conv1d(in_channels, filter_channels, 1)
+		self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
+		self.convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
+		if gin_channels != 0:
+			self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+
+	def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
+		x = torch.detach(x)
+		x = self.pre(x)
+		if g is not None:
+			g = torch.detach(g)
+			x = x + self.cond(g)
+		x = self.convs(x, x_mask)
+		x = self.proj(x) * x_mask
+
+		if not reverse:
+			flows = self.flows
+			assert w is not None
+
+			logdet_tot_q = 0
+			h_w = self.post_pre(w)
+			h_w = self.post_convs(h_w, x_mask)
+			h_w = self.post_proj(h_w) * x_mask
+			e_q = torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) * x_mask
+			z_q = e_q
+			for flow in self.post_flows:
+				z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
+				logdet_tot_q += logdet_q
+			z_u, z1 = torch.split(z_q, [1, 1], 1)
+			u = torch.sigmoid(z_u) * x_mask
+			z0 = (w - u) * x_mask
+			logdet_tot_q += torch.sum((F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1,2])
+			logq = torch.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - logdet_tot_q
+
+			logdet_tot = 0
+			z0, logdet = self.log_flow(z0, x_mask)
+			logdet_tot += logdet
+			z = torch.cat([z0, z1], 1)
+			for flow in flows:
+				z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+				logdet_tot = logdet_tot + logdet
+			nll = torch.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - logdet_tot
+			return nll + logq # [b]
+		else:
+			flows = list(reversed(self.flows))
+			flows = flows[:-2] + [flows[-1]] # remove a useless vflow
+			z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale
+			for flow in flows:
+				z = flow(z, x_mask, g=x, reverse=reverse)
+			z0, z1 = torch.split(z, [1, 1], 1)
+			logw = z0
+			return logw
+
+class PosteriorEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, g=None, tau=1.0):
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * tau * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+
+class Generator(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        gin_channels=0,
+    ):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(self, x, g=None):
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        print("Removing weight norm...")
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+
+
+class ReferenceEncoder(nn.Module):
+    """
+    inputs --- [N, Ty/r, n_mels*r]  mels
+    outputs --- [N, ref_enc_gru_size]
+    """
+
+    def __init__(self, spec_channels, gin_channels=0, layernorm=True):
+        super().__init__()
+        self.spec_channels = spec_channels
+        ref_enc_filters = [32, 32, 64, 64, 128, 128]
+        K = len(ref_enc_filters)
+        filters = [1] + ref_enc_filters
+        convs = [
+            weight_norm(
+                nn.Conv2d(
+                    in_channels=filters[i],
+                    out_channels=filters[i + 1],
+                    kernel_size=(3, 3),
+                    stride=(2, 2),
+                    padding=(1, 1),
+                )
+            )
+            for i in range(K)
+        ]
+        self.convs = nn.ModuleList(convs)
+
+        out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
+        self.gru = nn.GRU(
+            input_size=ref_enc_filters[-1] * out_channels,
+            hidden_size=256 // 2,
+            batch_first=True,
+        )
+        self.proj = nn.Linear(128, gin_channels)
+        if layernorm:
+            self.layernorm = nn.LayerNorm(self.spec_channels)
+        else:
+            self.layernorm = None
+
+    def forward(self, inputs, mask=None):
+        N = inputs.size(0)
+
+        out = inputs.view(N, 1, -1, self.spec_channels)  # [N, 1, Ty, n_freqs]
+        if self.layernorm is not None:
+            out = self.layernorm(out)
+
+        for conv in self.convs:
+            out = conv(out)
+            # out = wn(out)
+            out = F.relu(out)  # [N, 128, Ty//2^K, n_mels//2^K]
+
+        out = out.transpose(1, 2)  # [N, Ty//2^K, 128, n_mels//2^K]
+        T = out.size(1)
+        N = out.size(0)
+        out = out.contiguous().view(N, T, -1)  # [N, Ty//2^K, 128*n_mels//2^K]
+
+        self.gru.flatten_parameters()
+        memory, out = self.gru(out)  # out --- [1, N, 128]
+
+        return self.proj(out.squeeze(0))
+
+    def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
+        for i in range(n_convs):
+            L = (L - kernel_size + 2 * pad) // stride + 1
+        return L
+
+
+class ResidualCouplingBlock(nn.Module):
+    def __init__(self,
+            channels,
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            n_flows=4,
+            gin_channels=0):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(modules.ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
+            self.flows.append(modules.Flip())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+class SynthesizerTrn(nn.Module):
+    """
+    Synthesizer for Training
+    """
+
+    def __init__(
+        self,
+        n_vocab,
+        spec_channels,
+        inter_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        n_speakers=256,
+        gin_channels=256,
+        zero_g=False,
+        **kwargs
+    ):
+        super().__init__()
+
+        self.dec = Generator(
+            inter_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        )
+        self.enc_q = PosteriorEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+
+        self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
+
+        self.n_speakers = n_speakers
+        if n_speakers == 0:
+            self.ref_enc = ReferenceEncoder(spec_channels, gin_channels)
+        else:
+            self.enc_p = TextEncoder(n_vocab,
+                inter_channels,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout)
+            self.sdp = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels)
+            self.dp = DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels)
+            self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        self.zero_g = zero_g
+
+    def infer(self, x, x_lengths, sid=None, noise_scale=1, length_scale=1, noise_scale_w=1., sdp_ratio=0.2, max_len=None):
+        x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths)
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]
+        else:
+            g = None
+
+        logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * sdp_ratio \
+            + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
+
+        w = torch.exp(logw) * x_mask * length_scale
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(x_mask.dtype)
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.generate_path(w_ceil, attn_mask)
+
+        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
+        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        o = self.dec((z * y_mask)[:,:,:max_len], g=g)
+        return o, attn, y_mask, (z, z_p, m_p, logs_p)
+
+    def voice_conversion(self, y, y_lengths, sid_src, sid_tgt, tau=1.0):
+        g_src = sid_src
+        g_tgt = sid_tgt
+        z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src if not self.zero_g else torch.zeros_like(g_src), tau=tau)
+        z_p = self.flow(z, y_mask, g=g_src)
+        z_hat = self.flow(z_p, y_mask, g=g_tgt, reverse=True)
+        o_hat = self.dec(z_hat * y_mask, g=g_tgt if not self.zero_g else torch.zeros_like(g_tgt))
+        return o_hat, y_mask, (z, z_p, z_hat)

OpenVoice/openvoice/modules.py

@@ -0,0 +1,598 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from torch.nn import Conv1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from openvoice import commons
+from openvoice.commons import init_weights, get_padding
+from openvoice.transforms import piecewise_rational_quadratic_transform
+from openvoice.attentions import Encoder
+
+LRELU_SLOPE = 0.1
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        out_channels,
+        kernel_size,
+        n_layers,
+        p_dropout,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+        assert n_layers > 1, "Number of layers should be larger than 0."
+
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.conv_layers.append(
+            nn.Conv1d(
+                in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
+            )
+        )
+        self.norm_layers.append(LayerNorm(hidden_channels))
+        self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
+        for _ in range(n_layers - 1):
+            self.conv_layers.append(
+                nn.Conv1d(
+                    hidden_channels,
+                    hidden_channels,
+                    kernel_size,
+                    padding=kernel_size // 2,
+                )
+            )
+            self.norm_layers.append(LayerNorm(hidden_channels))
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask):
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
+
+
+class DDSConv(nn.Module):
+    """
+    Dilated and Depth-Separable Convolution
+    """
+
+    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+
+        self.drop = nn.Dropout(p_dropout)
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(n_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(LayerNorm(channels))
+            self.norms_2.append(LayerNorm(channels))
+
+    def forward(self, x, x_mask, g=None):
+        if g is not None:
+            x = x + g
+        for i in range(self.n_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.drop(y)
+            x = x + y
+        return x * x_mask
+
+
+class WN(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+        p_dropout=0,
+    ):
+        super(WN, self).__init__()
+        assert kernel_size % 2 == 1
+        self.hidden_channels = hidden_channels
+        self.kernel_size = (kernel_size,)
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = p_dropout
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
+
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(
+                gin_channels, 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
+
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, x, x_mask, g=None, **kwargs):
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
+            acts = self.drop(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+        self.convs.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class Log(nn.Module):
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+            logdet = torch.sum(-y, [1, 2])
+            return y, logdet
+        else:
+            x = torch.exp(x) * x_mask
+            return x
+
+
+class Flip(nn.Module):
+    def forward(self, x, *args, reverse=False, **kwargs):
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
+
+
+class ElementwiseAffine(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.m = nn.Parameter(torch.zeros(channels, 1))
+        self.logs = nn.Parameter(torch.zeros(channels, 1))
+
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = self.m + torch.exp(self.logs) * x
+            y = y * x_mask
+            logdet = torch.sum(self.logs * x_mask, [1, 2])
+            return y, logdet
+        else:
+            x = (x - self.m) * torch.exp(-self.logs) * x_mask
+            return x
+
+
+class ResidualCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        p_dropout=0,
+        gin_channels=0,
+        mean_only=False,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=p_dropout,
+            gin_channels=gin_channels,
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+
+class ConvFlow(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        filter_channels,
+        kernel_size,
+        n_layers,
+        num_bins=10,
+        tail_bound=5.0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.half_channels = in_channels // 2
+
+        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+        self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
+        self.proj = nn.Conv1d(
+            filter_channels, self.half_channels * (num_bins * 3 - 1), 1
+        )
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
+
+        b, c, t = x0.shape
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)  # [b, cx?, t] -> [b, c, t, ?]
+
+        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
+            self.filter_channels
+        )
+        unnormalized_derivatives = h[..., 2 * self.num_bins :]
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x
+
+
+class TransformerCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        n_layers,
+        n_heads,
+        p_dropout=0,
+        filter_channels=0,
+        mean_only=False,
+        wn_sharing_parameter=None,
+        gin_channels=0,
+    ):
+        assert n_layers == 3, n_layers
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = (
+            Encoder(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=gin_channels,
+            )
+            if wn_sharing_parameter is None
+            else wn_sharing_parameter
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x

OpenVoice/openvoice/openvoice_app.py

@@ -0,0 +1,275 @@
+import os
+import torch
+import argparse
+import gradio as gr
+from zipfile import ZipFile
+import langid
+from openvoice import se_extractor
+from openvoice.api import BaseSpeakerTTS, ToneColorConverter
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--share", action='store_true', default=False, help="make link public")
+args = parser.parse_args()
+
+en_ckpt_base = 'checkpoints/base_speakers/EN'
+zh_ckpt_base = 'checkpoints/base_speakers/ZH'
+ckpt_converter = 'checkpoints/converter'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+output_dir = 'outputs'
+os.makedirs(output_dir, exist_ok=True)
+
+# load models
+en_base_speaker_tts = BaseSpeakerTTS(f'{en_ckpt_base}/config.json', device=device)
+en_base_speaker_tts.load_ckpt(f'{en_ckpt_base}/checkpoint.pth')
+zh_base_speaker_tts = BaseSpeakerTTS(f'{zh_ckpt_base}/config.json', device=device)
+zh_base_speaker_tts.load_ckpt(f'{zh_ckpt_base}/checkpoint.pth')
+tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json', device=device)
+tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth')
+
+# load speaker embeddings
+en_source_default_se = torch.load(f'{en_ckpt_base}/en_default_se.pth').to(device)
+en_source_style_se = torch.load(f'{en_ckpt_base}/en_style_se.pth').to(device)
+zh_source_se = torch.load(f'{zh_ckpt_base}/zh_default_se.pth').to(device)
+
+# This online demo mainly supports English and Chinese
+supported_languages = ['zh', 'en']
+
+def predict(prompt, style, audio_file_pth, agree):
+    # initialize a empty info
+    text_hint = ''
+    # agree with the terms
+    if agree == False:
+        text_hint += '[ERROR] Please accept the Terms & Condition!\n'
+        gr.Warning("Please accept the Terms & Condition!")
+        return (
+            text_hint,
+            None,
+            None,
+        )
+
+    # first detect the input language
+    language_predicted = langid.classify(prompt)[0].strip()  
+    print(f"Detected language:{language_predicted}")
+
+    if language_predicted not in supported_languages:
+        text_hint += f"[ERROR] The detected language {language_predicted} for your input text is not in our Supported Languages: {supported_languages}\n"
+        gr.Warning(
+            f"The detected language {language_predicted} for your input text is not in our Supported Languages: {supported_languages}"
+        )
+
+        return (
+            text_hint,
+            None,
+            None,
+        )
+    
+    if language_predicted == "zh":
+        tts_model = zh_base_speaker_tts
+        source_se = zh_source_se
+        language = 'Chinese'
+        if style not in ['default']:
+            text_hint += f"[ERROR] The style {style} is not supported for Chinese, which should be in ['default']\n"
+            gr.Warning(f"The style {style} is not supported for Chinese, which should be in ['default']")
+            return (
+                text_hint,
+                None,
+                None,
+            )
+
+    else:
+        tts_model = en_base_speaker_tts
+        if style == 'default':
+            source_se = en_source_default_se
+        else:
+            source_se = en_source_style_se
+        language = 'English'
+        if style not in ['default', 'whispering', 'shouting', 'excited', 'cheerful', 'terrified', 'angry', 'sad', 'friendly']:
+            text_hint += f"[ERROR] The style {style} is not supported for English, which should be in ['default', 'whispering', 'shouting', 'excited', 'cheerful', 'terrified', 'angry', 'sad', 'friendly']\n"
+            gr.Warning(f"The style {style} is not supported for English, which should be in ['default', 'whispering', 'shouting', 'excited', 'cheerful', 'terrified', 'angry', 'sad', 'friendly']")
+            return (
+                text_hint,
+                None,
+                None,
+            )
+
+    speaker_wav = audio_file_pth
+
+    if len(prompt) < 2:
+        text_hint += f"[ERROR] Please give a longer prompt text \n"
+        gr.Warning("Please give a longer prompt text")
+        return (
+            text_hint,
+            None,
+            None,
+        )
+    if len(prompt) > 200:
+        text_hint += f"[ERROR] Text length limited to 200 characters for this demo, please try shorter text. You can clone our open-source repo and try for your usage \n"
+        gr.Warning(
+            "Text length limited to 200 characters for this demo, please try shorter text. You can clone our open-source repo for your usage"
+        )
+        return (
+            text_hint,
+            None,
+            None,
+        )
+    
+    # note diffusion_conditioning not used on hifigan (default mode), it will be empty but need to pass it to model.inference
+    try:
+        target_se, audio_name = se_extractor.get_se(speaker_wav, tone_color_converter, target_dir='processed', vad=True)
+    except Exception as e:
+        text_hint += f"[ERROR] Get target tone color error {str(e)} \n"
+        gr.Warning(
+            "[ERROR] Get target tone color error {str(e)} \n"
+        )
+        return (
+            text_hint,
+            None,
+            None,
+        )
+
+    src_path = f'{output_dir}/tmp.wav'
+    tts_model.tts(prompt, src_path, speaker=style, language=language)
+
+    save_path = f'{output_dir}/output.wav'
+    # Run the tone color converter
+    encode_message = "@MyShell"
+    tone_color_converter.convert(
+        audio_src_path=src_path, 
+        src_se=source_se, 
+        tgt_se=target_se, 
+        output_path=save_path,
+        message=encode_message)
+
+    text_hint += f'''Get response successfully \n'''
+
+    return (
+        text_hint,
+        save_path,
+        speaker_wav,
+    )
+
+
+
+title = "MyShell OpenVoice"
+
+description = """
+We introduce OpenVoice, a versatile instant voice cloning approach that requires only a short audio clip from the reference speaker to replicate their voice and generate speech in multiple languages. OpenVoice enables granular control over voice styles, including emotion, accent, rhythm, pauses, and intonation, in addition to replicating the tone color of the reference speaker. OpenVoice also achieves zero-shot cross-lingual voice cloning for languages not included in the massive-speaker training set.
+"""
+
+markdown_table = """
+<div align="center" style="margin-bottom: 10px;">
+
+|               |               |               |
+| :-----------: | :-----------: | :-----------: | 
+| **OpenSource Repo** | **Project Page** | **Join the Community** |        
+| <div style='text-align: center;'><a style="display:inline-block,align:center" href='https://github.com/myshell-ai/OpenVoice'><img src='https://img.shields.io/github/stars/myshell-ai/OpenVoice?style=social' /></a></div> | [OpenVoice](https://research.myshell.ai/open-voice) | [![Discord](https://img.shields.io/discord/1122227993805336617?color=%239B59B6&label=%20Discord%20)](https://discord.gg/myshell) |
+
+</div>
+"""
+
+markdown_table_v2 = """
+<div align="center" style="margin-bottom: 2px;">
+
+|               |               |               |              |
+| :-----------: | :-----------: | :-----------: | :-----------: | 
+| **OpenSource Repo** | <div style='text-align: center;'><a style="display:inline-block,align:center" href='https://github.com/myshell-ai/OpenVoice'><img src='https://img.shields.io/github/stars/myshell-ai/OpenVoice?style=social' /></a></div> |  **Project Page** |  [OpenVoice](https://research.myshell.ai/open-voice) |     
+
+| | |
+| :-----------: | :-----------: |
+**Join the Community** |   [![Discord](https://img.shields.io/discord/1122227993805336617?color=%239B59B6&label=%20Discord%20)](https://discord.gg/myshell) |
+
+</div>
+"""
+content = """
+<div>
+  <strong>If the generated voice does not sound like the reference voice, please refer to <a href='https://github.com/myshell-ai/OpenVoice/blob/main/docs/QA.md'>this QnA</a>.</strong> <strong>For multi-lingual & cross-lingual examples, please refer to <a href='https://github.com/myshell-ai/OpenVoice/blob/main/demo_part2.ipynb'>this jupyter notebook</a>.</strong>
+  This online demo mainly supports <strong>English</strong>. The <em>default</em> style also supports <strong>Chinese</strong>. But OpenVoice can adapt to any other language as long as a base speaker is provided.
+</div>
+"""
+wrapped_markdown_content = f"<div style='border: 1px solid #000; padding: 10px;'>{content}</div>"
+
+
+examples = [
+    [
+        "今天天气真好，我们一起出去吃饭吧。",
+        'default',
+        "resources/demo_speaker1.mp3",
+        True,
+    ],[
+        "This audio is generated by open voice with a half-performance model.",
+        'whispering',
+        "resources/demo_speaker2.mp3",
+        True,
+    ],
+    [
+        "He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered, flour-fattened sauce.",
+        'sad',
+        "resources/demo_speaker0.mp3",
+        True,
+    ],
+]
+
+with gr.Blocks(analytics_enabled=False) as demo:
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                gr.Markdown(
+                    """
+                    ## <img src="https://huggingface.co/spaces/myshell-ai/OpenVoice/raw/main/logo.jpg" height="40"/>
+                    """
+                )
+            with gr.Row():    
+                gr.Markdown(markdown_table_v2)
+            with gr.Row():
+                gr.Markdown(description)
+        with gr.Column():
+            gr.Video('https://github.com/myshell-ai/OpenVoice/assets/40556743/3cba936f-82bf-476c-9e52-09f0f417bb2f', autoplay=True)
+            
+    with gr.Row():
+        gr.HTML(wrapped_markdown_content)
+
+    with gr.Row():
+        with gr.Column():
+            input_text_gr = gr.Textbox(
+                label="Text Prompt",
+                info="One or two sentences at a time is better. Up to 200 text characters.",
+                value="He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered, flour-fattened sauce.",
+            )
+            style_gr = gr.Dropdown(
+                label="Style",
+                info="Select a style of output audio for the synthesised speech. (Chinese only support 'default' now)",
+                choices=['default', 'whispering', 'cheerful', 'terrified', 'angry', 'sad', 'friendly'],
+                max_choices=1,
+                value="default",
+            )
+            ref_gr = gr.Audio(
+                label="Reference Audio",
+                info="Click on the ✎ button to upload your own target speaker audio",
+                type="filepath",
+                value="resources/demo_speaker2.mp3",
+            )
+            tos_gr = gr.Checkbox(
+                label="Agree",
+                value=False,
+                info="I agree to the terms of the cc-by-nc-4.0 license-: https://github.com/myshell-ai/OpenVoice/blob/main/LICENSE",
+            )
+
+            tts_button = gr.Button("Send", elem_id="send-btn", visible=True)
+
+
+        with gr.Column():
+            out_text_gr = gr.Text(label="Info")
+            audio_gr = gr.Audio(label="Synthesised Audio", autoplay=True)
+            ref_audio_gr = gr.Audio(label="Reference Audio Used")
+
+            gr.Examples(examples,
+                        label="Examples",
+                        inputs=[input_text_gr, style_gr, ref_gr, tos_gr],
+                        outputs=[out_text_gr, audio_gr, ref_audio_gr],
+                        fn=predict,
+                        cache_examples=False,)
+            tts_button.click(predict, [input_text_gr, style_gr, ref_gr, tos_gr], outputs=[out_text_gr, audio_gr, ref_audio_gr])
+
+demo.queue()  
+demo.launch(debug=True, show_api=True, share=args.share)

OpenVoice/openvoice/se_extractor.py

@@ -0,0 +1,153 @@
+import os
+import glob
+import torch
+import hashlib
+import librosa
+import base64
+from glob import glob
+import numpy as np
+from pydub import AudioSegment
+from faster_whisper import WhisperModel
+import hashlib
+import base64
+import librosa
+from whisper_timestamped.transcribe import get_audio_tensor, get_vad_segments
+
+model_size = "medium"
+# Run on GPU with FP16
+model = None
+def split_audio_whisper(audio_path, audio_name, target_dir='processed'):
+    global model
+    if model is None:
+        model = WhisperModel(model_size, device="cuda", compute_type="float16")
+    audio = AudioSegment.from_file(audio_path)
+    max_len = len(audio)
+
+    target_folder = os.path.join(target_dir, audio_name)
+    
+    segments, info = model.transcribe(audio_path, beam_size=5, word_timestamps=True)
+    segments = list(segments)    
+
+    # create directory
+    os.makedirs(target_folder, exist_ok=True)
+    wavs_folder = os.path.join(target_folder, 'wavs')
+    os.makedirs(wavs_folder, exist_ok=True)
+
+    # segments
+    s_ind = 0
+    start_time = None
+    
+    for k, w in enumerate(segments):
+        # process with the time
+        if k == 0:
+            start_time = max(0, w.start)
+
+        end_time = w.end
+
+        # calculate confidence
+        if len(w.words) > 0:
+            confidence = sum([s.probability for s in w.words]) / len(w.words)
+        else:
+            confidence = 0.
+        # clean text
+        text = w.text.replace('...', '')
+
+        # left 0.08s for each audios
+        audio_seg = audio[int( start_time * 1000) : min(max_len, int(end_time * 1000) + 80)]
+
+        # segment file name
+        fname = f"{audio_name}_seg{s_ind}.wav"
+
+        # filter out the segment shorter than 1.5s and longer than 20s
+        save = audio_seg.duration_seconds > 1.5 and \
+                audio_seg.duration_seconds < 20. and \
+                len(text) >= 2 and len(text) < 200 
+
+        if save:
+            output_file = os.path.join(wavs_folder, fname)
+            audio_seg.export(output_file, format='wav')
+
+        if k < len(segments) - 1:
+            start_time = max(0, segments[k+1].start - 0.08)
+
+        s_ind = s_ind + 1
+    return wavs_folder
+
+
+def split_audio_vad(audio_path, audio_name, target_dir, split_seconds=10.0):
+    SAMPLE_RATE = 16000
+    audio_vad = get_audio_tensor(audio_path)
+    segments = get_vad_segments(
+        audio_vad,
+        output_sample=True,
+        min_speech_duration=0.1,
+        min_silence_duration=1,
+        method="silero",
+    )
+    segments = [(seg["start"], seg["end"]) for seg in segments]
+    segments = [(float(s) / SAMPLE_RATE, float(e) / SAMPLE_RATE) for s,e in segments]
+    print(segments)
+    audio_active = AudioSegment.silent(duration=0)
+    audio = AudioSegment.from_file(audio_path)
+
+    for start_time, end_time in segments:
+        audio_active += audio[int( start_time * 1000) : int(end_time * 1000)]
+    
+    audio_dur = audio_active.duration_seconds
+    print(f'after vad: dur = {audio_dur}')
+    target_folder = os.path.join(target_dir, audio_name)
+    wavs_folder = os.path.join(target_folder, 'wavs')
+    os.makedirs(wavs_folder, exist_ok=True)
+    start_time = 0.
+    count = 0
+    num_splits = int(np.round(audio_dur / split_seconds))
+    assert num_splits > 0, 'input audio is too short'
+    interval = audio_dur / num_splits
+
+    for i in range(num_splits):
+        end_time = min(start_time + interval, audio_dur)
+        if i == num_splits - 1:
+            end_time = audio_dur
+        output_file = f"{wavs_folder}/{audio_name}_seg{count}.wav"
+        audio_seg = audio_active[int(start_time * 1000): int(end_time * 1000)]
+        audio_seg.export(output_file, format='wav')
+        start_time = end_time
+        count += 1
+    return wavs_folder
+
+def hash_numpy_array(audio_path):
+    array, _ = librosa.load(audio_path, sr=None, mono=True)
+    # Convert the array to bytes
+    array_bytes = array.tobytes()
+    # Calculate the hash of the array bytes
+    hash_object = hashlib.sha256(array_bytes)
+    hash_value = hash_object.digest()
+    # Convert the hash value to base64
+    base64_value = base64.b64encode(hash_value)
+    return base64_value.decode('utf-8')[:16].replace('/', '_^')
+
+def get_se(audio_path, vc_model, target_dir='processed', vad=True):
+    device = vc_model.device
+    version = vc_model.version
+    print("OpenVoice version:", version)
+
+    audio_name = f"{os.path.basename(audio_path).rsplit('.', 1)[0]}_{version}_{hash_numpy_array(audio_path)}"
+    se_path = os.path.join(target_dir, audio_name, 'se.pth')
+
+    # if os.path.isfile(se_path):
+    #     se = torch.load(se_path).to(device)
+    #     return se, audio_name
+    # if os.path.isdir(audio_path):
+    #     wavs_folder = audio_path
+    
+    if vad:
+        wavs_folder = split_audio_vad(audio_path, target_dir=target_dir, audio_name=audio_name)
+    else:
+        wavs_folder = split_audio_whisper(audio_path, target_dir=target_dir, audio_name=audio_name)
+    
+    audio_segs = glob(f'{wavs_folder}/*.wav')
+    if len(audio_segs) == 0:
+        raise NotImplementedError('No audio segments found!')
+    
+    return vc_model.extract_se(audio_segs, se_save_path=se_path), audio_name
+

OpenVoice/openvoice/text/__init__.py

@@ -0,0 +1,79 @@
+""" from https://github.com/keithito/tacotron """
+from openvoice.text import cleaners
+from openvoice.text.symbols import symbols
+
+
+# Mappings from symbol to numeric ID and vice versa:
+_symbol_to_id = {s: i for i, s in enumerate(symbols)}
+_id_to_symbol = {i: s for i, s in enumerate(symbols)}
+
+
+def text_to_sequence(text, symbols, cleaner_names):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+      cleaner_names: names of the cleaner functions to run the text through
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  sequence = []
+  symbol_to_id = {s: i for i, s in enumerate(symbols)}
+  clean_text = _clean_text(text, cleaner_names)
+  print(clean_text)
+  print(f" length:{len(clean_text)}")
+  for symbol in clean_text:
+    if symbol not in symbol_to_id.keys():
+      continue
+    symbol_id = symbol_to_id[symbol]
+    sequence += [symbol_id]
+  print(f" length:{len(sequence)}")
+  return sequence
+
+
+def cleaned_text_to_sequence(cleaned_text, symbols):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  symbol_to_id = {s: i for i, s in enumerate(symbols)}
+  sequence = [symbol_to_id[symbol] for symbol in cleaned_text if symbol in symbol_to_id.keys()]
+  return sequence
+
+
+
+from openvoice.text.symbols import language_tone_start_map
+def cleaned_text_to_sequence_vits2(cleaned_text, tones, language, symbols, languages):
+    """Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+    Returns:
+      List of integers corresponding to the symbols in the text
+    """
+    symbol_to_id = {s: i for i, s in enumerate(symbols)}
+    language_id_map = {s: i for i, s in enumerate(languages)}
+    phones = [symbol_to_id[symbol] for symbol in cleaned_text]
+    tone_start = language_tone_start_map[language]
+    tones = [i + tone_start for i in tones]
+    lang_id = language_id_map[language]
+    lang_ids = [lang_id for i in phones]
+    return phones, tones, lang_ids
+
+
+def sequence_to_text(sequence):
+  '''Converts a sequence of IDs back to a string'''
+  result = ''
+  for symbol_id in sequence:
+    s = _id_to_symbol[symbol_id]
+    result += s
+  return result
+
+
+def _clean_text(text, cleaner_names):
+  for name in cleaner_names:
+    cleaner = getattr(cleaners, name)
+    if not cleaner:
+      raise Exception('Unknown cleaner: %s' % name)
+    text = cleaner(text)
+  return text

OpenVoice/openvoice/text/cleaners.py

@@ -0,0 +1,16 @@
+import re
+from openvoice.text.english import english_to_lazy_ipa, english_to_ipa2, english_to_lazy_ipa2
+from openvoice.text.mandarin import number_to_chinese, chinese_to_bopomofo, latin_to_bopomofo, chinese_to_romaji, chinese_to_lazy_ipa, chinese_to_ipa, chinese_to_ipa2
+
+def cjke_cleaners2(text):
+    text = re.sub(r'\[ZH\](.*?)\[ZH\]',
+                  lambda x: chinese_to_ipa(x.group(1))+' ', text)
+    text = re.sub(r'\[JA\](.*?)\[JA\]',
+                  lambda x: japanese_to_ipa2(x.group(1))+' ', text)
+    text = re.sub(r'\[KO\](.*?)\[KO\]',
+                  lambda x: korean_to_ipa(x.group(1))+' ', text)
+    text = re.sub(r'\[EN\](.*?)\[EN\]',
+                  lambda x: english_to_ipa2(x.group(1))+' ', text)
+    text = re.sub(r'\s+$', '', text)
+    text = re.sub(r'([^\.,!\?\-…~])$', r'\1.', text)
+    return text

OpenVoice/openvoice/text/english.py

@@ -0,0 +1,188 @@
+""" from https://github.com/keithito/tacotron """
+
+'''
+Cleaners are transformations that run over the input text at both training and eval time.
+
+Cleaners can be selected by passing a comma-delimited list of cleaner names as the "cleaners"
+hyperparameter. Some cleaners are English-specific. You'll typically want to use:
+  1. "english_cleaners" for English text
+  2. "transliteration_cleaners" for non-English text that can be transliterated to ASCII using
+     the Unidecode library (https://pypi.python.org/pypi/Unidecode)
+  3. "basic_cleaners" if you do not want to transliterate (in this case, you should also update
+     the symbols in symbols.py to match your data).
+'''
+
+
+# Regular expression matching whitespace:
+
+
+import re
+import inflect
+from unidecode import unidecode
+import eng_to_ipa as ipa
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r'([0-9][0-9\,]+[0-9])')
+_decimal_number_re = re.compile(r'([0-9]+\.[0-9]+)')
+_pounds_re = re.compile(r'£([0-9\,]*[0-9]+)')
+_dollars_re = re.compile(r'\$([0-9\.\,]*[0-9]+)')
+_ordinal_re = re.compile(r'[0-9]+(st|nd|rd|th)')
+_number_re = re.compile(r'[0-9]+')
+
+# List of (regular expression, replacement) pairs for abbreviations:
+_abbreviations = [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [
+    ('mrs', 'misess'),
+    ('mr', 'mister'),
+    ('dr', 'doctor'),
+    ('st', 'saint'),
+    ('co', 'company'),
+    ('jr', 'junior'),
+    ('maj', 'major'),
+    ('gen', 'general'),
+    ('drs', 'doctors'),
+    ('rev', 'reverend'),
+    ('lt', 'lieutenant'),
+    ('hon', 'honorable'),
+    ('sgt', 'sergeant'),
+    ('capt', 'captain'),
+    ('esq', 'esquire'),
+    ('ltd', 'limited'),
+    ('col', 'colonel'),
+    ('ft', 'fort'),
+]]
+
+
+# List of (ipa, lazy ipa) pairs:
+_lazy_ipa = [(re.compile('%s' % x[0]), x[1]) for x in [
+    ('r', 'ɹ'),
+    ('æ', 'e'),
+    ('ɑ', 'a'),
+    ('ɔ', 'o'),
+    ('ð', 'z'),
+    ('θ', 's'),
+    ('ɛ', 'e'),
+    ('ɪ', 'i'),
+    ('ʊ', 'u'),
+    ('ʒ', 'ʥ'),
+    ('ʤ', 'ʥ'),
+    ('ˈ', '↓'),
+]]
+
+# List of (ipa, lazy ipa2) pairs:
+_lazy_ipa2 = [(re.compile('%s' % x[0]), x[1]) for x in [
+    ('r', 'ɹ'),
+    ('ð', 'z'),
+    ('θ', 's'),
+    ('ʒ', 'ʑ'),
+    ('ʤ', 'dʑ'),
+    ('ˈ', '↓'),
+]]
+
+# List of (ipa, ipa2) pairs
+_ipa_to_ipa2 = [(re.compile('%s' % x[0]), x[1]) for x in [
+    ('r', 'ɹ'),
+    ('ʤ', 'dʒ'),
+    ('ʧ', 'tʃ')
+]]
+
+
+def expand_abbreviations(text):
+    for regex, replacement in _abbreviations:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def collapse_whitespace(text):
+    return re.sub(r'\s+', ' ', text)
+
+
+def _remove_commas(m):
+    return m.group(1).replace(',', '')
+
+
+def _expand_decimal_point(m):
+    return m.group(1).replace('.', ' point ')
+
+
+def _expand_dollars(m):
+    match = m.group(1)
+    parts = match.split('.')
+    if len(parts) > 2:
+        return match + ' dollars'  # Unexpected format
+    dollars = int(parts[0]) if parts[0] else 0
+    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if dollars and cents:
+        dollar_unit = 'dollar' if dollars == 1 else 'dollars'
+        cent_unit = 'cent' if cents == 1 else 'cents'
+        return '%s %s, %s %s' % (dollars, dollar_unit, cents, cent_unit)
+    elif dollars:
+        dollar_unit = 'dollar' if dollars == 1 else 'dollars'
+        return '%s %s' % (dollars, dollar_unit)
+    elif cents:
+        cent_unit = 'cent' if cents == 1 else 'cents'
+        return '%s %s' % (cents, cent_unit)
+    else:
+        return 'zero dollars'
+
+
+def _expand_ordinal(m):
+    return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m):
+    num = int(m.group(0))
+    if num > 1000 and num < 3000:
+        if num == 2000:
+            return 'two thousand'
+        elif num > 2000 and num < 2010:
+            return 'two thousand ' + _inflect.number_to_words(num % 100)
+        elif num % 100 == 0:
+            return _inflect.number_to_words(num // 100) + ' hundred'
+        else:
+            return _inflect.number_to_words(num, andword='', zero='oh', group=2).replace(', ', ' ')
+    else:
+        return _inflect.number_to_words(num, andword='')
+
+
+def normalize_numbers(text):
+    text = re.sub(_comma_number_re, _remove_commas, text)
+    text = re.sub(_pounds_re, r'\1 pounds', text)
+    text = re.sub(_dollars_re, _expand_dollars, text)
+    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+    text = re.sub(_ordinal_re, _expand_ordinal, text)
+    text = re.sub(_number_re, _expand_number, text)
+    return text
+
+
+def mark_dark_l(text):
+    return re.sub(r'l([^aeiouæɑɔəɛɪʊ ]*(?: |$))', lambda x: 'ɫ'+x.group(1), text)
+
+
+def english_to_ipa(text):
+    text = unidecode(text).lower()
+    text = expand_abbreviations(text)
+    text = normalize_numbers(text)
+    phonemes = ipa.convert(text)
+    phonemes = collapse_whitespace(phonemes)
+    return phonemes
+
+
+def english_to_lazy_ipa(text):
+    text = english_to_ipa(text)
+    for regex, replacement in _lazy_ipa:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def english_to_ipa2(text):
+    text = english_to_ipa(text)
+    text = mark_dark_l(text)
+    for regex, replacement in _ipa_to_ipa2:
+        text = re.sub(regex, replacement, text)
+    return text.replace('...', '…')
+
+
+def english_to_lazy_ipa2(text):
+    text = english_to_ipa(text)
+    for regex, replacement in _lazy_ipa2:
+        text = re.sub(regex, replacement, text)
+    return text

OpenVoice/openvoice/text/mandarin.py

@@ -0,0 +1,326 @@
+import os
+import sys
+import re
+from pypinyin import lazy_pinyin, BOPOMOFO
+import jieba
+import cn2an
+import logging
+
+
+# List of (Latin alphabet, bopomofo) pairs:
+_latin_to_bopomofo = [(re.compile('%s' % x[0], re.IGNORECASE), x[1]) for x in [
+    ('a', 'ㄟˉ'),
+    ('b', 'ㄅㄧˋ'),
+    ('c', 'ㄙㄧˉ'),
+    ('d', 'ㄉㄧˋ'),
+    ('e', 'ㄧˋ'),
+    ('f', 'ㄝˊㄈㄨˋ'),
+    ('g', 'ㄐㄧˋ'),
+    ('h', 'ㄝˇㄑㄩˋ'),
+    ('i', 'ㄞˋ'),
+    ('j', 'ㄐㄟˋ'),
+    ('k', 'ㄎㄟˋ'),
+    ('l', 'ㄝˊㄛˋ'),
+    ('m', 'ㄝˊㄇㄨˋ'),
+    ('n', 'ㄣˉ'),
+    ('o', 'ㄡˉ'),
+    ('p', 'ㄆㄧˉ'),
+    ('q', 'ㄎㄧㄡˉ'),
+    ('r', 'ㄚˋ'),
+    ('s', 'ㄝˊㄙˋ'),
+    ('t', 'ㄊㄧˋ'),
+    ('u', 'ㄧㄡˉ'),
+    ('v', 'ㄨㄧˉ'),
+    ('w', 'ㄉㄚˋㄅㄨˋㄌㄧㄡˋ'),
+    ('x', 'ㄝˉㄎㄨˋㄙˋ'),
+    ('y', 'ㄨㄞˋ'),
+    ('z', 'ㄗㄟˋ')
+]]
+
+# List of (bopomofo, romaji) pairs:
+_bopomofo_to_romaji = [(re.compile('%s' % x[0]), x[1]) for x in [
+    ('ㄅㄛ', 'p⁼wo'),
+    ('ㄆㄛ', 'pʰwo'),
+    ('ㄇㄛ', 'mwo'),
+    ('ㄈㄛ', 'fwo'),
+    ('ㄅ', 'p⁼'),
+    ('ㄆ', 'pʰ'),
+    ('ㄇ', 'm'),
+    ('ㄈ', 'f'),
+    ('ㄉ', 't⁼'),
+    ('ㄊ', 'tʰ'),
+    ('ㄋ', 'n'),
+    ('ㄌ', 'l'),
+    ('ㄍ', 'k⁼'),
+    ('ㄎ', 'kʰ'),
+    ('ㄏ', 'h'),
+    ('ㄐ', 'ʧ⁼'),
+    ('ㄑ', 'ʧʰ'),
+    ('ㄒ', 'ʃ'),
+    ('ㄓ', 'ʦ`⁼'),
+    ('ㄔ', 'ʦ`ʰ'),
+    ('ㄕ', 's`'),
+    ('ㄖ', 'ɹ`'),
+    ('ㄗ', 'ʦ⁼'),
+    ('ㄘ', 'ʦʰ'),
+    ('ㄙ', 's'),
+    ('ㄚ', 'a'),
+    ('ㄛ', 'o'),
+    ('ㄜ', 'ə'),
+    ('ㄝ', 'e'),
+    ('ㄞ', 'ai'),
+    ('ㄟ', 'ei'),
+    ('ㄠ', 'au'),
+    ('ㄡ', 'ou'),
+    ('ㄧㄢ', 'yeNN'),
+    ('ㄢ', 'aNN'),
+    ('ㄧㄣ', 'iNN'),
+    ('ㄣ', 'əNN'),
+    ('ㄤ', 'aNg'),
+    ('ㄧㄥ', 'iNg'),
+    ('ㄨㄥ', 'uNg'),
+    ('ㄩㄥ', 'yuNg'),
+    ('ㄥ', 'əNg'),
+    ('ㄦ', 'əɻ'),
+    ('ㄧ', 'i'),
+    ('ㄨ', 'u'),
+    ('ㄩ', 'ɥ'),
+    ('ˉ', '→'),
+    ('ˊ', '↑'),
+    ('ˇ', '↓↑'),
+    ('ˋ', '↓'),
+    ('˙', ''),
+    ('，', ','),
+    ('。', '.'),
+    ('！', '!'),
+    ('？', '?'),
+    ('—', '-')
+]]
+
+# List of (romaji, ipa) pairs:
+_romaji_to_ipa = [(re.compile('%s' % x[0], re.IGNORECASE), x[1]) for x in [
+    ('ʃy', 'ʃ'),
+    ('ʧʰy', 'ʧʰ'),
+    ('ʧ⁼y', 'ʧ⁼'),
+    ('NN', 'n'),
+    ('Ng', 'ŋ'),
+    ('y', 'j'),
+    ('h', 'x')
+]]
+
+# List of (bopomofo, ipa) pairs:
+_bopomofo_to_ipa = [(re.compile('%s' % x[0]), x[1]) for x in [
+    ('ㄅㄛ', 'p⁼wo'),
+    ('ㄆㄛ', 'pʰwo'),
+    ('ㄇㄛ', 'mwo'),
+    ('ㄈㄛ', 'fwo'),
+    ('ㄅ', 'p⁼'),
+    ('ㄆ', 'pʰ'),
+    ('ㄇ', 'm'),
+    ('ㄈ', 'f'),
+    ('ㄉ', 't⁼'),
+    ('ㄊ', 'tʰ'),
+    ('ㄋ', 'n'),
+    ('ㄌ', 'l'),
+    ('ㄍ', 'k⁼'),
+    ('ㄎ', 'kʰ'),
+    ('ㄏ', 'x'),
+    ('ㄐ', 'tʃ⁼'),
+    ('ㄑ', 'tʃʰ'),
+    ('ㄒ', 'ʃ'),
+    ('ㄓ', 'ts`⁼'),
+    ('ㄔ', 'ts`ʰ'),
+    ('ㄕ', 's`'),
+    ('ㄖ', 'ɹ`'),
+    ('ㄗ', 'ts⁼'),
+    ('ㄘ', 'tsʰ'),
+    ('ㄙ', 's'),
+    ('ㄚ', 'a'),
+    ('ㄛ', 'o'),
+    ('ㄜ', 'ə'),
+    ('ㄝ', 'ɛ'),
+    ('ㄞ', 'aɪ'),
+    ('ㄟ', 'eɪ'),
+    ('ㄠ', 'ɑʊ'),
+    ('ㄡ', 'oʊ'),
+    ('ㄧㄢ', 'jɛn'),
+    ('ㄩㄢ', 'ɥæn'),
+    ('ㄢ', 'an'),
+    ('ㄧㄣ', 'in'),
+    ('ㄩㄣ', 'ɥn'),
+    ('ㄣ', 'ən'),
+    ('ㄤ', 'ɑŋ'),
+    ('ㄧㄥ', 'iŋ'),
+    ('ㄨㄥ', 'ʊŋ'),
+    ('ㄩㄥ', 'jʊŋ'),
+    ('ㄥ', 'əŋ'),
+    ('ㄦ', 'əɻ'),
+    ('ㄧ', 'i'),
+    ('ㄨ', 'u'),
+    ('ㄩ', 'ɥ'),
+    ('ˉ', '→'),
+    ('ˊ', '↑'),
+    ('ˇ', '↓↑'),
+    ('ˋ', '↓'),
+    ('˙', ''),
+    ('，', ','),
+    ('。', '.'),
+    ('！', '!'),
+    ('？', '?'),
+    ('—', '-')
+]]
+
+# List of (bopomofo, ipa2) pairs:
+_bopomofo_to_ipa2 = [(re.compile('%s' % x[0]), x[1]) for x in [
+    ('ㄅㄛ', 'pwo'),
+    ('ㄆㄛ', 'pʰwo'),
+    ('ㄇㄛ', 'mwo'),
+    ('ㄈㄛ', 'fwo'),
+    ('ㄅ', 'p'),
+    ('ㄆ', 'pʰ'),
+    ('ㄇ', 'm'),
+    ('ㄈ', 'f'),
+    ('ㄉ', 't'),
+    ('ㄊ', 'tʰ'),
+    ('ㄋ', 'n'),
+    ('ㄌ', 'l'),
+    ('ㄍ', 'k'),
+    ('ㄎ', 'kʰ'),
+    ('ㄏ', 'h'),
+    ('ㄐ', 'tɕ'),
+    ('ㄑ', 'tɕʰ'),
+    ('ㄒ', 'ɕ'),
+    ('ㄓ', 'tʂ'),
+    ('ㄔ', 'tʂʰ'),
+    ('ㄕ', 'ʂ'),
+    ('ㄖ', 'ɻ'),
+    ('ㄗ', 'ts'),
+    ('ㄘ', 'tsʰ'),
+    ('ㄙ', 's'),
+    ('ㄚ', 'a'),
+    ('ㄛ', 'o'),
+    ('ㄜ', 'ɤ'),
+    ('ㄝ', 'ɛ'),
+    ('ㄞ', 'aɪ'),
+    ('ㄟ', 'eɪ'),
+    ('ㄠ', 'ɑʊ'),
+    ('ㄡ', 'oʊ'),
+    ('ㄧㄢ', 'jɛn'),
+    ('ㄩㄢ', 'yæn'),
+    ('ㄢ', 'an'),
+    ('ㄧㄣ', 'in'),
+    ('ㄩㄣ', 'yn'),
+    ('ㄣ', 'ən'),
+    ('ㄤ', 'ɑŋ'),
+    ('ㄧㄥ', 'iŋ'),
+    ('ㄨㄥ', 'ʊŋ'),
+    ('ㄩㄥ', 'jʊŋ'),
+    ('ㄥ', 'ɤŋ'),
+    ('ㄦ', 'əɻ'),
+    ('ㄧ', 'i'),
+    ('ㄨ', 'u'),
+    ('ㄩ', 'y'),
+    ('ˉ', '˥'),
+    ('ˊ', '˧˥'),
+    ('ˇ', '˨˩˦'),
+    ('ˋ', '˥˩'),
+    ('˙', ''),
+    ('，', ','),
+    ('。', '.'),
+    ('！', '!'),
+    ('？', '?'),
+    ('—', '-')
+]]
+
+
+def number_to_chinese(text):
+    numbers = re.findall(r'\d+(?:\.?\d+)?', text)
+    for number in numbers:
+        text = text.replace(number, cn2an.an2cn(number), 1)
+    return text
+
+
+def chinese_to_bopomofo(text):
+    text = text.replace('、', '，').replace('；', '，').replace('：', '，')
+    words = jieba.lcut(text, cut_all=False)
+    text = ''
+    for word in words:
+        bopomofos = lazy_pinyin(word, BOPOMOFO)
+        if not re.search('[\u4e00-\u9fff]', word):
+            text += word
+            continue
+        for i in range(len(bopomofos)):
+            bopomofos[i] = re.sub(r'([\u3105-\u3129])$', r'\1ˉ', bopomofos[i])
+        if text != '':
+            text += ' '
+        text += ''.join(bopomofos)
+    return text
+
+
+def latin_to_bopomofo(text):
+    for regex, replacement in _latin_to_bopomofo:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def bopomofo_to_romaji(text):
+    for regex, replacement in _bopomofo_to_romaji:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def bopomofo_to_ipa(text):
+    for regex, replacement in _bopomofo_to_ipa:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def bopomofo_to_ipa2(text):
+    for regex, replacement in _bopomofo_to_ipa2:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def chinese_to_romaji(text):
+    text = number_to_chinese(text)
+    text = chinese_to_bopomofo(text)
+    text = latin_to_bopomofo(text)
+    text = bopomofo_to_romaji(text)
+    text = re.sub('i([aoe])', r'y\1', text)
+    text = re.sub('u([aoəe])', r'w\1', text)
+    text = re.sub('([ʦsɹ]`[⁼ʰ]?)([→↓↑ ]+|$)',
+                  r'\1ɹ`\2', text).replace('ɻ', 'ɹ`')
+    text = re.sub('([ʦs][⁼ʰ]?)([→↓↑ ]+|$)', r'\1ɹ\2', text)
+    return text
+
+
+def chinese_to_lazy_ipa(text):
+    text = chinese_to_romaji(text)
+    for regex, replacement in _romaji_to_ipa:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+def chinese_to_ipa(text):
+    text = number_to_chinese(text)
+    text = chinese_to_bopomofo(text)
+    text = latin_to_bopomofo(text)
+    text = bopomofo_to_ipa(text)
+    text = re.sub('i([aoe])', r'j\1', text)
+    text = re.sub('u([aoəe])', r'w\1', text)
+    text = re.sub('([sɹ]`[⁼ʰ]?)([→↓↑ ]+|$)',
+                  r'\1ɹ`\2', text).replace('ɻ', 'ɹ`')
+    text = re.sub('([s][⁼ʰ]?)([→↓↑ ]+|$)', r'\1ɹ\2', text)
+    return text
+
+
+def chinese_to_ipa2(text):
+    text = number_to_chinese(text)
+    text = chinese_to_bopomofo(text)
+    text = latin_to_bopomofo(text)
+    text = bopomofo_to_ipa2(text)
+    text = re.sub(r'i([aoe])', r'j\1', text)
+    text = re.sub(r'u([aoəe])', r'w\1', text)
+    text = re.sub(r'([ʂɹ]ʰ?)([˩˨˧˦˥ ]+|$)', r'\1ʅ\2', text)
+    text = re.sub(r'(sʰ?)([˩˨˧˦˥ ]+|$)', r'\1ɿ\2', text)
+    return text

OpenVoice/openvoice/text/symbols.py

@@ -0,0 +1,88 @@
+'''
+Defines the set of symbols used in text input to the model.
+'''
+
+# japanese_cleaners
+# _pad        = '_'
+# _punctuation = ',.!?-'
+# _letters = 'AEINOQUabdefghijkmnoprstuvwyzʃʧ↓↑ '
+
+
+'''# japanese_cleaners2
+_pad        = '_'
+_punctuation = ',.!?-~…'
+_letters = 'AEINOQUabdefghijkmnoprstuvwyzʃʧʦ↓↑ '
+'''
+
+
+'''# korean_cleaners
+_pad        = '_'
+_punctuation = ',.!?…~'
+_letters = 'ㄱㄴㄷㄹㅁㅂㅅㅇㅈㅊㅋㅌㅍㅎㄲㄸㅃㅆㅉㅏㅓㅗㅜㅡㅣㅐㅔ '
+'''
+
+'''# chinese_cleaners
+_pad        = '_'
+_punctuation = '，。！？—…'
+_letters = 'ㄅㄆㄇㄈㄉㄊㄋㄌㄍㄎㄏㄐㄑㄒㄓㄔㄕㄖㄗㄘㄙㄚㄛㄜㄝㄞㄟㄠㄡㄢㄣㄤㄥㄦㄧㄨㄩˉˊˇˋ˙ '
+'''
+
+# # zh_ja_mixture_cleaners
+# _pad        = '_'
+# _punctuation = ',.!?-~…'
+# _letters = 'AEINOQUabdefghijklmnoprstuvwyzʃʧʦɯɹəɥ⁼ʰ`→↓↑ '
+
+
+'''# sanskrit_cleaners
+_pad        = '_'
+_punctuation = '।'
+_letters = 'ँंःअआइईउऊऋएऐओऔकखगघङचछजझञटठडढणतथदधनपफबभमयरलळवशषसहऽािीुूृॄेैोौ्ॠॢ '
+'''
+
+'''# cjks_cleaners
+_pad        = '_'
+_punctuation = ',.!?-~…'
+_letters = 'NQabdefghijklmnopstuvwxyzʃʧʥʦɯɹəɥçɸɾβŋɦː⁼ʰ`^#*=→↓↑ '
+'''
+
+'''# thai_cleaners
+_pad        = '_'
+_punctuation = '.!? '
+_letters = 'กขฃคฆงจฉชซฌญฎฏฐฑฒณดตถทธนบปผฝพฟภมยรฤลวศษสหฬอฮฯะัาำิีึืุูเแโใไๅๆ็่้๊๋์'
+'''
+
+# # cjke_cleaners2
+_pad        = '_'
+_punctuation = ',.!?-~…'
+_letters = 'NQabdefghijklmnopstuvwxyzɑæʃʑçɯɪɔɛɹðəɫɥɸʊɾʒθβŋɦ⁼ʰ`^#*=ˈˌ→↓↑ '
+
+
+'''# shanghainese_cleaners
+_pad        = '_'
+_punctuation = ',.!?…'
+_letters = 'abdfghiklmnopstuvyzøŋȵɑɔɕəɤɦɪɿʑʔʰ̩̃ᴀᴇ15678 '
+'''
+
+'''# chinese_dialect_cleaners
+_pad        = '_'
+_punctuation = ',.!?~…─'
+_letters = '#Nabdefghijklmnoprstuvwxyzæçøŋœȵɐɑɒɓɔɕɗɘəɚɛɜɣɤɦɪɭɯɵɷɸɻɾɿʂʅʊʋʌʏʑʔʦʮʰʷˀː˥˦˧˨˩̥̩̃̚ᴀᴇ↑↓∅ⱼ '
+'''
+
+# Export all symbols:
+symbols = [_pad] + list(_punctuation) + list(_letters)
+
+# Special symbol ids
+SPACE_ID = symbols.index(" ")
+
+num_ja_tones = 1
+num_kr_tones = 1
+num_zh_tones = 6
+num_en_tones = 4
+
+language_tone_start_map = {
+    "ZH": 0,
+    "JP": num_zh_tones,
+    "EN": num_zh_tones + num_ja_tones,
+    'KR': num_zh_tones + num_ja_tones + num_en_tones,
+}

OpenVoice/openvoice/transforms.py

@@ -0,0 +1,209 @@
+import torch
+from torch.nn import functional as F
+
+import numpy as np
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails=None,
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if tails is None:
+        spline_fn = rational_quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unconstrained_rational_quadratic_spline
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+    outputs, logabsdet = spline_fn(
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unconstrained_rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails="linear",
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == "linear":
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError("{} tails are not implemented.".format(tails))
+
+    (
+        outputs[inside_interval_mask],
+        logabsdet[inside_interval_mask],
+    ) = rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+
+    return outputs, logabsdet
+
+
+def rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    left=0.0,
+    right=1.0,
+    bottom=0.0,
+    top=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError("Minimal bin width too large for the number of bins")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError("Minimal bin height too large for the number of bins")
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (
+            input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+        )
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet

OpenVoice/openvoice/utils.py

@@ -0,0 +1,194 @@
+import re
+import json
+import numpy as np
+
+
+def get_hparams_from_file(config_path):
+    with open(config_path, "r", encoding="utf-8") as f:
+        data = f.read()
+    config = json.loads(data)
+
+    hparams = HParams(**config)
+    return hparams
+
+class HParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if type(v) == dict:
+                v = HParams(**v)
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+
+def string_to_bits(string, pad_len=8):
+    # Convert each character to its ASCII value
+    ascii_values = [ord(char) for char in string]
+    
+    # Convert ASCII values to binary representation
+    binary_values = [bin(value)[2:].zfill(8) for value in ascii_values]
+    
+    # Convert binary strings to integer arrays
+    bit_arrays = [[int(bit) for bit in binary] for binary in binary_values]
+    
+    # Convert list of arrays to NumPy array
+    numpy_array = np.array(bit_arrays)
+    numpy_array_full = np.zeros((pad_len, 8), dtype=numpy_array.dtype)
+    numpy_array_full[:, 2] = 1
+    max_len = min(pad_len, len(numpy_array))
+    numpy_array_full[:max_len] = numpy_array[:max_len]
+    return numpy_array_full
+
+
+def bits_to_string(bits_array):
+    # Convert each row of the array to a binary string
+    binary_values = [''.join(str(bit) for bit in row) for row in bits_array]
+    
+    # Convert binary strings to ASCII values
+    ascii_values = [int(binary, 2) for binary in binary_values]
+    
+    # Convert ASCII values to characters
+    output_string = ''.join(chr(value) for value in ascii_values)
+    
+    return output_string
+
+
+def split_sentence(text, min_len=10, language_str='[EN]'):
+    if language_str in ['EN']:
+        sentences = split_sentences_latin(text, min_len=min_len)
+    else:
+        sentences = split_sentences_zh(text, min_len=min_len)
+    return sentences
+
+def split_sentences_latin(text, min_len=10):
+    """Split Long sentences into list of short ones
+
+    Args:
+        str: Input sentences.
+
+    Returns:
+        List[str]: list of output sentences.
+    """
+    # deal with dirty sentences
+    text = re.sub('[。！？；]', '.', text)
+    text = re.sub('[，]', ',', text)
+    text = re.sub('[“”]', '"', text)
+    text = re.sub('[‘’]', "'", text)
+    text = re.sub(r"[\<\>\(\)\[\]\"\«\»]+", "", text)
+    text = re.sub('[\n\t ]+', ' ', text)
+    text = re.sub('([,.!?;])', r'\1 $#!', text)
+    # split
+    sentences = [s.strip() for s in text.split('$#!')]
+    if len(sentences[-1]) == 0: del sentences[-1]
+
+    new_sentences = []
+    new_sent = []
+    count_len = 0
+    for ind, sent in enumerate(sentences):
+        # print(sent)
+        new_sent.append(sent)
+        count_len += len(sent.split(" "))
+        if count_len > min_len or ind == len(sentences) - 1:
+            count_len = 0
+            new_sentences.append(' '.join(new_sent))
+            new_sent = []
+    return merge_short_sentences_latin(new_sentences)
+
+
+def merge_short_sentences_latin(sens):
+    """Avoid short sentences by merging them with the following sentence.
+
+    Args:
+        List[str]: list of input sentences.
+
+    Returns:
+        List[str]: list of output sentences.
+    """
+    sens_out = []
+    for s in sens:
+        # If the previous sentence is too short, merge them with
+        # the current sentence.
+        if len(sens_out) > 0 and len(sens_out[-1].split(" ")) <= 2:
+            sens_out[-1] = sens_out[-1] + " " + s
+        else:
+            sens_out.append(s)
+    try:
+        if len(sens_out[-1].split(" ")) <= 2:
+            sens_out[-2] = sens_out[-2] + " " + sens_out[-1]
+            sens_out.pop(-1)
+    except:
+        pass
+    return sens_out
+
+def split_sentences_zh(text, min_len=10):
+    text = re.sub('[。！？；]', '.', text)
+    text = re.sub('[，]', ',', text)
+    # 将文本中的换行符、空格和制表符替换为空格
+    text = re.sub('[\n\t ]+', ' ', text)
+    # 在标点符号后添加一个空格
+    text = re.sub('([,.!?;])', r'\1 $#!', text)
+    # 分隔句子并去除前后空格
+    # sentences = [s.strip() for s in re.split('(。|！|？|；)', text)]
+    sentences = [s.strip() for s in text.split('$#!')]
+    if len(sentences[-1]) == 0: del sentences[-1]
+
+    new_sentences = []
+    new_sent = []
+    count_len = 0
+    for ind, sent in enumerate(sentences):
+        new_sent.append(sent)
+        count_len += len(sent)
+        if count_len > min_len or ind == len(sentences) - 1:
+            count_len = 0
+            new_sentences.append(' '.join(new_sent))
+            new_sent = []
+    return merge_short_sentences_zh(new_sentences)
+
+
+def merge_short_sentences_zh(sens):
+    # return sens
+    """Avoid short sentences by merging them with the following sentence.
+
+    Args:
+        List[str]: list of input sentences.
+
+    Returns:
+        List[str]: list of output sentences.
+    """
+    sens_out = []
+    for s in sens:
+        # If the previous sentense is too short, merge them with
+        # the current sentence.
+        if len(sens_out) > 0 and len(sens_out[-1]) <= 2:
+            sens_out[-1] = sens_out[-1] + " " + s
+        else:
+            sens_out.append(s)
+    try:
+        if len(sens_out[-1]) <= 2:
+            sens_out[-2] = sens_out[-2] + " " + sens_out[-1]
+            sens_out.pop(-1)
+    except:
+        pass
+    return sens_out

OpenVoice/setup.py

@@ -0,0 +1,45 @@
+from setuptools import setup, find_packages
+
+
+setup(name='MyShell-OpenVoice',
+      version='0.0.0',
+      description='Instant voice cloning by MyShell.',
+      long_description=open('README.md').read().strip(),
+      long_description_content_type='text/markdown',
+      keywords=[
+            'text-to-speech',
+            'tts',
+            'voice-clone',
+            'zero-shot-tts'
+      ],
+      url='https://github.com/myshell-ai/OpenVoice',
+      project_urls={
+        'Documentation': 'https://github.com/myshell-ai/OpenVoice/blob/main/docs/USAGE.md',
+        'Changes': 'https://github.com/myshell-ai/OpenVoice/releases',
+        'Code': 'https://github.com/myshell-ai/OpenVoice',
+        'Issue tracker': 'https://github.com/myshell-ai/OpenVoice/issues',
+      },
+      author='MyShell',
+      author_email='ethan@myshell.ai',
+      license='MIT License',
+      packages=find_packages(),
+
+      python_requires='>=3.9',
+      install_requires=[
+            'librosa==0.9.1',
+            'faster-whisper==0.9.0',
+            'pydub==0.25.1',
+            'wavmark==0.0.3',
+            'numpy==1.22.0',
+            'eng_to_ipa==0.0.2',
+            'inflect==7.0.0',
+            'unidecode==1.3.7',
+            'whisper-timestamped==1.14.2',
+            'pypinyin==0.50.0',
+            'cn2an==0.5.22',
+            'jieba==0.42.1',
+            'gradio==3.48.0',
+            'langid==1.1.6'
+      ],
+      zip_safe=False
+      )

app.py

@@ -0,0 +1,548 @@
+import gradio as gr
+import numpy as np
+import json
+from tts_api import TTSapi, DEFAULT_TTS_MODEL_NAME
+from config import *
+from utils import *
+from knowledge_base import LocalRAG, CosPlayer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+def handle_retry(history, thinking_history, config, section_state, retry_data: gr.RetryData):
+    # 获取用户之前的消息
+    previous_message = history[retry_data.index]['content']
+    # 清除后续的回复和思考过程
+    new_history = history[:retry_data.index]
+    section_state['chat_history'] = section_state['chat_history'][:retry_data.index + 1]
+    
+    try:
+        items = thinking_history.split('\n==================\n')
+        if len(items) > 2:
+            new_thinking_history = '\n==================\n'.join(items[:-2])
+        else:
+            new_thinking_history = ''
+            
+        items = section_state['thinking_history'].split('\n==================\n')
+        if len(items) > 2:
+            section_state['thinking_history'] = '\n==================\n'.join(items[:-2])
+        else:
+            section_state['thinking_history'] = ''
+    except Exception as e:
+        print('-----------------------------------')
+        print(e)
+        print('-----------------------------------')
+        print('思考过程发生异常，重置为空')
+        section_state['thinking_history'] = ''
+        new_thinking_history = ''
+    # 重新生成回复
+    return predict(previous_message, new_history, new_thinking_history, config, section_state)
+
+        
+def predict(message, chat_history, thinking_history, config, section_state):
+    global local_rag, TTS_LOADED, LLM_LOADED, synthesiser, core_llm, core_tokenizer
+    print(f"当前模式：{config['mode_selected']}")
+    print(f'角色扮演描述：{config["character_description"]}')
+    print(f"写入角色设定方式：{config['character_setting_mode']}")
+    print(f"选中LLM：{config['llm_model']}")
+    print(f"是否使用RAG本地知识库：{config['kb_on']}")
+    print(f"选中知识库：{config['current_knowledge_base']}")
+    print(f"是否联网搜索：{config['net_on']}")
+    print(f"选中TTS模型：{config['tts_model']}")
+    print(f"是否合成语音：{config['tts_on']}")
+    print(f"参考音频路径：{config['ref_audio']}")
+    print(f"参考音频文本：{config['ref_audio_transcribe']}")
+    
+    context = ''
+    net_search_res = []
+    docs = []
+    if config['kb_on'] and len(config['current_knowledge_base']) > 0:
+        # 检索相似文档  
+        doc_and_scores = local_rag.vector_db.similarity_search(message, k=local_rag.rag_top_k)
+        # doc_and_scores = list(filter(lambda x: x[1] <= 0.4, doc_and_scores))
+        if len(doc_and_scores) > 0:
+            docs, scores = list(zip(*doc_and_scores))
+            docs, scores = list(docs), list(scores)
+            context_local = "【本地知识库】" + "\n".join([concate_metadata(d.metadata) + d.page_content for d in docs])
+            context = context + context_local
+    if config['net_on']:
+        # 检索相似文档
+        ret = web_search(message, max_results=MAX_RESULTS)
+        net_search_res = parse_net_search(ret)
+        context_net = "\n【网络搜索结果】" + ''.join(net_search_res)
+        context = context + context_net
+     
+    if config['character_description']:
+        if config['character_setting_mode'] == 'by system':
+            if len(section_state['chat_history']) == 0 or section_state['chat_history'][0]['role'] != 'system':
+                section_state['chat_history'].insert(0, {"role": "system", "content": config["character_description"]})
+        elif config['character_setting_mode'] == 'by prompt':
+            if len(section_state['chat_history']) > 0 and section_state['chat_history'][0]['role'] == 'system':
+                section_state['chat_history'].pop(0)
+            context = f'【系统核心设定】：{config["character_description"]}\n' if config["character_description"] else '' + context
+        else:
+            raise ValueError(f"未知的角色设定模式：{config['character_setting_mode']}")
+    
+    if len(context) > 0:
+        prompt = f"""请充分理解以下上下文信息，并结合当前及历史对话产生回复'：\n
+                    上下文：{context}
+                    用户当前输入：{message}
+                    回复：
+                """
+        input_message = section_state["chat_history"] + [{"role": "user", "content": prompt}]
+    else:
+        input_message = section_state["chat_history"] + [{"role": "user", "content": message}]
+    
+    #  关闭Qwen3系列默认的思考模式
+    if config['llm_model'].startswith('Qwen3'):
+        input_message[-1]['content'] += '/no_think'
+    # input_message[-1]['content'] += '/no_think'
+    
+    # 添加用户消息到历史
+    section_state["chat_history"].append({"role": "user", "content": message})
+    
+    
+    
+    try:
+        # 调用模型
+        if not LLM_LOADED:
+            core_llm = AutoModelForCausalLM.from_pretrained(
+                    config['llm_model'],
+                    torch_dtype="auto",
+                    device_map="auto"
+                )
+            core_tokenizer = AutoTokenizer.from_pretrained(config['llm_model'])
+            LLM_LOADED = True
+        token_cnt = count_tokens_local(input_message, core_tokenizer)
+        if token_cnt >= MAX_MODEL_CTX:
+            gr.Warning("当前对话已经超出模型上下文长度，请开启新会话...")
+        text = core_tokenizer.apply_chat_template(
+            input_message,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False
+        )
+        model_inputs = core_tokenizer([text], return_tensors="pt").to(core_llm.device)
+        # conduct text completion
+        generated_ids = core_llm.generate(
+            **model_inputs,
+            max_new_tokens=32768
+        )
+        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
+
+        # parsing thinking content
+        # try:
+        #     # rindex finding 151668 (</think>)
+        #     index = len(output_ids) - output_ids[::-1].index(151668)
+        # except ValueError:
+        #     index = 0
+        index = 0
+        # thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
+        thinking = None
+        response_content = core_tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
+        print('回复：', response_content)
+        # 更新对话历史
+        chat_history.append({'role': 'user', 'content': message})
+        if len(context) > 0:
+            # 构建带折叠结构的消息
+            formatted_response = f"""
+            <details class="rag-details">
+                <summary style='cursor: pointer; color: #666;'>
+                    🔍 检索完成✅（共{len(docs)+len(net_search_res)}条）
+                </summary>
+                <div style='margin:10px 0;padding:10px;background:#f5f5f5;border-radius:8px;'>
+                    {
+                        "<br>".join(
+                            ["<br>".join(wash_up_content(content if isinstance(content, str) else (content.page_content, scores[idx]))) 
+                            for idx, content in enumerate(docs + net_search_res)]
+                        ) 
+                    }
+                </div>
+            </details>
+            <div style="margin-top: 10px;">{response_content}</div>  <!-- 增加顶部间距容器 -->
+            """
+            chat_history.append({'role': 'assistant', 'content': formatted_response})
+        else:
+            chat_history.append({'role': 'assistant', 'content': response_content})
+
+        thinking_history += f"User: {message}\nThinking: {thinking}" + '\n==================\n'
+        # 添加助手响应到历史
+        section_state["chat_history"].append({"role": "assistant", "content": response_content})
+        section_state["thinking_history"] += f"User: {message}\nThinking: {thinking}" + '\n==================\n'
+        
+        if (not config['tts_on']) or len(response_content) == 0:
+            audio_output = np.array([0], dtype=np.int16)
+            if len(response_content) == 0:
+                print("LLM 回复为空，无法合成语音") 
+        else:
+            if not TTS_LOADED:
+                print('TTS模型首次加载...')
+                gr.Info("初次加载TTS模型，请稍候..", duration=63)
+                synthesiser = TTSapi(model_name=config['tts_model'])
+                TTS_LOADED = True
+                print('加载完毕...')
+            # 检查当前模型是否是所选
+            if config['tts_model'] != synthesiser.model_name:
+                print(f'当前TTS模型{synthesiser.model_name}非所选，重新加载')
+                synthesiser.reload(model_name=config['tts_model'])
+
+            # 如果提供了参考音频，则需把参考音频的文本加在response_content前面作为前缀
+            if config['ref_audio']:
+                prompt_text = config['ref_audio_transcribe']
+                if prompt_text is None:
+                    # prompt_text = ...
+                    raise NotImplementedError('暂时必须提供文本')  # TODO：考虑后续加入ASR模型
+                response_content = prompt_text + response_content
+                
+            audio_output = synthesiser.forward(response_content, speech_prompt=config['ref_audio'])
+        
+    except Exception as e:
+        print('!!!!!!!!')
+        print(e)
+        print('!!!!!!!!')
+        error_msg = f"Error: {str(e)}"
+        chat_history.append((message, error_msg))
+        thinking_history += f"Error occurred: {str(e)}" + '\n'
+    
+    return "", chat_history, thinking_history, (synthesiser.sr if synthesiser else 16000, audio_output)
+
+
+def init_model(init_llm=True, init_rag=False, init_tts=False):
+    if init_llm:
+        print(f'正在加载LLM:{DEFAULT_MODEL_NAME}...')
+        core_llm = AutoModelForCausalLM.from_pretrained(
+                    DEFAULT_MODEL_NAME,
+                    torch_dtype="auto",
+                    device_map="auto"
+                )
+        print('device:', core_llm.device)
+        core_tokenizer = AutoTokenizer.from_pretrained(DEFAULT_MODEL_NAME)
+        LLM_LOADED = True
+    else:
+        core_llm, core_tokenizer = None, None
+        LLM_LOADED = False
+        
+    if init_rag:
+        gr.Info("正在加载知识库，请稍候...")
+        local_rag = LocalRAG(rag_top_k=RAG_TOP_K)
+    else:
+        local_rag  =None
+    
+    if init_tts:
+        print(f'正在加载TTS模型:{DEFAULT_TTS_MODEL_NAME}...')
+        synthesiser = TTSapi()
+        TTS_LOADED = True
+    else:
+        synthesiser = None
+        TTS_LOADED = False
+    return local_rag, synthesiser, core_llm, core_tokenizer, TTS_LOADED, LLM_LOADED
+        
+
+if __name__ == "__main__":
+    import time
+    st = time.time()
+    print('********************模型加载中************************')
+    local_rag, synthesiser, core_llm, core_tokenizer, TTS_LOADED, LLM_LOADED = init_model()
+    print('********************模型加载完成************************')
+    print('耗时:',time.time() - st)
+    
+    state = {}
+    resp, state = log_in(0, state)
+    cosplayer = CosPlayer(description_file=DEFAULT_COSPLAY_SETTING)
+    print("===== 初始化开始 =====")
+    with gr.Blocks(css=CSS, title="LLM Chat Demo", theme=gr.themes.Soft(font=["sans-serif", "Arial"])) as demo:
+        gr.Markdown("""
+        # LLM Chat Demo
+        ## 用法介绍
+        ### 用户登录
+        * 输入用户名，点击Log In按钮。首次登录会自动创建用户目录，聊天记录会保存在下面，如不登录，默认为公共目录'0'
+        ### 模型选择
+        目前支持Qwen、Deepseek-R1蒸馏系列等部分模型，可下拉菜单选择
+        ### 高级设置
+        * 模式选择：可以选择角色扮演模式/普通模式
+        * 角色设定选择：支持加载不同角色设定文件
+        * 角色配置方式：
+            * by system: 角色设定将作为system prompt存在于输入首部
+            * by prompt: 角色设定每次被添加到当前上下文中
+        * 知识库配置： 支持自由选择、组合知识库
+        """)
+        section_state = gr.State(value=state)  # 创建会话状态对象
+        with gr.Row():
+            uid_input = gr.Textbox(label="Type Your UID:")
+            response = gr.Textbox(label='', value=resp)
+        login_button = gr.Button("Log In")
+        llm_select = gr.Dropdown(label= "模型选择", choices=AVALIABLE_MODELS, value=DEFAULT_MODEL_NAME, visible=True)
+        
+        gr.Markdown("## 高级设置")
+        with gr.Accordion("点击展开折叠", open=False, visible=True):
+            mode_select = gr.Radio(label='模式选择', choices=SUPPORT_MODES, value=DEFAULT_MODE)
+            coser_select = gr.Dropdown(label= "角色设定选择", choices=cosplayer.get_all_characters(), value=DEFAULT_COSPLAY_SETTING, visible=True)
+            coser_setting = gr.Radio(label='角色配置方式', choices=CHARACTER_SETTING_MODES, value=DEFAULT_C_SETTING_MODE, visible=True)
+            kb_select = gr.Dropdown(label= "知识库配置", choices=AVALIABLE_KNOWLEDGE_BASE, value=None, visible=True, multiselect=True)
+            
+        with gr.Row():
+            # 页面左侧
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(label="对话记录", height=500, show_copy_button=True, type='messages')
+                with gr.Row():
+                    msg = gr.Textbox(label="输入消息", placeholder="请输入您的问题...", scale=7)
+                    with gr.Column(scale=1, min_width=15):
+                        with gr.Row():
+                            rag_switch = gr.Checkbox(label="本地RAG", value=False, info="")
+                            net_switch = gr.Checkbox(label="联网搜索", value=False, info="")
+                            
+                        submit_btn = gr.Button("发送", variant="primary", min_width=15)#, , elem_classes=['custom-btn']
+                with gr.Row():
+                    gr.Examples(
+                        examples=[[example] for example in EXAMPLES],
+                        inputs=msg,
+                        outputs=chatbot,
+                        fn=predict,
+                        visible=True,
+                        cache_examples=False
+                    )    
+                with gr.Row():
+                    save_btn = gr.Button("保存对话")
+                    clear_btn = gr.Button("清空对话")
+                    chat_history_select = gr.Dropdown(label='加载历史对话', choices=state['available_history'], visible=True, interactive=True)
+            
+            # 页面右侧
+            with gr.Column(scale=2):
+                thinking_display = gr.TextArea(label="思考过程",interactive=False,
+                                            placeholder="模型思考过程将在此显���..."
+                                            )
+                tts_switch = gr.Checkbox(label="TTS开关", value=False, info="Check me to hear voice")
+                with gr.Tabs() as audio_tabs:
+                    # 选项卡1：音频播放
+                    with gr.Tab("音频输出", id="audio_output"):
+                        audio_player = gr.Audio(
+                            label="听听我声音~",
+                            type="numpy",
+                            interactive=False
+                        )
+                    
+                    # 选项卡2：TTS配置
+                    with gr.Tab("TTS配置", id="tts_config"):
+                        # TTS模型选择
+                        tts_model = gr.Dropdown(
+                            label="选择TTS模型",
+                            choices=["ChatTTS", "HKUSTAudio/LLasa-1B", "HKUSTAudio/LLasa-3B", "HKUSTAudio/LLasa-8B"],
+                            value=DEFAULT_TTS_MODEL_NAME,
+                            interactive=True
+                        )
+                        
+                        # 参考音频上传
+                        ref_audio = gr.Audio(
+                            label="上传参考音频",
+                            type="filepath",
+                            interactive=True
+                        )
+                        ref_audio_transcribe = gr.Textbox(label="参考音频对应文本", visible=True)
+
+        
+        # ================= 状态管理 =================
+        current_config = gr.State({
+            "llm_model": DEFAULT_MODEL_NAME,
+            "tts_model": DEFAULT_TTS_MODEL_NAME,
+            "tts_on": False,
+            "kb_on": False,
+            "net_on": False,
+            "ref_audio": None,
+            "ref_audio_transcribe": None,
+            "mode_selected": DEFAULT_MODE,
+            "character_description": cosplayer.get_core_setting(),
+            "character_setting_mode": DEFAULT_C_SETTING_MODE,
+            "current_knowledge_base": AVALIABLE_KNOWLEDGE_BASE[0]
+        })
+        
+        # 事件处理
+        login_button.click(log_in, inputs=[uid_input, section_state], outputs=[response, section_state])
+        
+        gr.on(triggers=[llm_select.change, tts_model.change, ref_audio.change, 
+                        ref_audio_transcribe.change, tts_switch.select, rag_switch.select, net_switch.select,
+                        mode_select.change], 
+            fn=lambda model1, model2, audio, text, tts_on, kb_on, net_on, mode, character_setting, kb_select: {"llm_model": model1, "tts_model": model2, "ref_audio": audio, 
+                                                            "ref_audio_transcribe": text, "tts_on": tts_on, "kb_on": kb_on, 'net_on': net_on,
+                                                            "mode_selected": mode, "character_description": None if mode == '普通模式' else cosplayer.get_core_setting(),
+                                                            "character_setting_mode": character_setting, "current_knowledge_base": kb_select},
+            inputs=[llm_select, tts_model, ref_audio, ref_audio_transcribe, tts_switch, rag_switch, net_switch, mode_select, coser_setting, kb_select],
+            outputs=current_config
+            )
+        msg.submit(
+            predict,
+            [msg, chatbot, thinking_display, current_config, section_state],
+            [msg, chatbot, thinking_display, audio_player],
+            queue=False
+        )
+        chatbot.retry(fn=handle_retry, 
+                      inputs=[chatbot, thinking_display, current_config, section_state], 
+                      outputs=[msg, chatbot, thinking_display, audio_player])
+        
+        submit_btn.click(
+            predict,
+            [msg, chatbot, thinking_display, current_config, section_state],
+            [msg, chatbot, thinking_display, audio_player],
+            queue=False
+        )
+        
+        def save_chat(state):
+            from datetime import datetime
+            now = datetime.now().strftime('%Y%m%d_%H%M%S')
+            with open(state['user_dir'] / f'chat_history_{now}.json', 'w', encoding='utf-8') as file:
+                json.dump(state["chat_history"], file, ensure_ascii=False, indent=4)
+            with open(state['user_dir'] / f'thinking_history_{now}.txt', 'w') as file:
+                if isinstance(state["thinking_history"], list):
+                    for item in state["thinking_history"]:
+                        file.write(item + '\n')
+                else:
+                    file.write(state["thinking_history"])
+
+            gr.Info("聊天记录已保存！")
+            state['available_history'].append(f'chat_history_{now}')
+            return state
+        
+        def clear_chat(state):
+            state["chat_history"] = []
+            state["thinking_history"] = []
+            prologue = cosplayer.get_prologue()
+            if prologue:
+                state['chat_history'].append({'role': 'assistant', 'content': prologue})
+                chatbot = [{'role': 'assistant', 'content': prologue}]
+            else:
+                chatbot = []
+            return chatbot, [], state
+        
+        def load_chat(state, chat_file):
+            # NOTE: 加载历史聊天记录。一般在对话开始之前加载，如果本次对话已经开始，本操作会覆盖当前会话内容
+            if chat_file:
+                think_file = chat_file.replace("chat_", "thinking_")
+                chat_file_path = state['user_dir'] / (chat_file + '.json')
+                think_file_path = state['user_dir'] / (think_file + '.txt')
+                
+                if not chat_file_path.exists():
+                    gr.Warning(f'聊天记录文件:{chat_file}.json不存在, 加载失败')
+                    return [], '', state
+                
+                with open(chat_file_path, 'r', encoding='utf-8') as f:
+                    content = json.load(f)
+                    state['chat_history'] = content
+                    
+                think = ''
+                if think_file_path.exists():
+                    with open(think_file_path, 'r') as f:
+                        think = f.read()
+                
+                state['thinking_history'] = think
+                
+                # 转换成chatbot可以识别的格式
+                # bot_content = parse_chat_history(content)
+                # 指定chatbot类型为message后，无需解析
+                bot_content = content
+                return bot_content, think, state
+            
+            return [], '', state
+            
+        
+        def update_history(state):
+            return gr.update(choices=state['available_history'])
+        
+        def update_visible(mode):
+            if mode != '普通模式':
+                gr.Warning("当前为角色扮演模式，请确认已配置好该角色的知识库...")
+                
+                return gr.update(visible=True), gr.update(visible=True)
+            return gr.update(visible=False), gr.update(visible=False)
+
+        def update_cosplay(cos_select, config, chatbot, think_display, state):
+            cosplayer.update(cos_select)
+            config['character_description'] = cosplayer.get_core_setting()
+            # 角色设定发生改变后，自动保存当前聊天记录，之后清空历史记录
+            if len(state['chat_history']) > 1:
+                state = save_chat(state)
+            gr.Warning("我的角色已更换，对话已重置。请检查知识库是否需要更新...")
+            chatbot, think_display, state = clear_chat(state)
+            return gr.update(value=cos_select), config, chatbot, think_display, state
+        
+        def update_character_setting_mode(coser_setting, config):
+            config['character_setting_mode'] = coser_setting
+            return gr.update(value=coser_setting), config
+        
+        def update_knowledge_base(knowledge_base, config):
+            global local_rag
+            config['current_knowledge_base'] = knowledge_base
+            if len(knowledge_base) == 0:
+                gr.Warning("当前未选中任何知识库，本地RAG将失效。请确认...")
+            else:
+                if local_rag is None:
+                    gr.Info("初次加载知识库，请稍候...")
+                    local_rag = LocalRAG(rag_top_k=RAG_TOP_K, doc_dir=knowledge_base)
+                    gr.Info("知识库加载完成！")
+                else:
+                    gr.Info("重新加载知识库，请稍候...")
+                    local_rag.reload_knowledge_base(knowledge_base)
+                    gr.Info("知识库加载完成！")
+            return gr.update(value=knowledge_base), config
+                    
+        def init_kb(rag_on, kb_select, config):
+            global local_rag
+            if rag_on:
+                # 初始化本地知识库
+                if config['mode_selected'] == "角色扮演":
+                    gr.Warning("当前为角色扮演模式，请确认已配置好该角色的知识库...")
+                
+                if local_rag is None:
+                    gr.Info("初次加载知识库，请稍候...")
+                    local_rag = LocalRAG(rag_top_k=RAG_TOP_K, doc_dir=kb_select)
+                    gr.Info("知识库加载完成！")
+            return gr.update(value=rag_on)
+                
+        # 选择非普通模式时（角色扮演），会展示可控选择的角色设定列表
+        mode_select.change(update_visible,
+                           inputs=mode_select,
+                           outputs=[coser_select, coser_setting])
+        
+        coser_select.change(update_cosplay,
+                           inputs=[coser_select, current_config, chatbot, thinking_display, section_state],
+                           outputs=[coser_select, current_config, chatbot, thinking_display, section_state])
+        # TODO: 根据角色变化动态展示示例
+        # coser_select.change(update_examples,
+        #                    inputs=[coser_select],
+        #                    outputs=[examples_show])
+        
+        coser_setting.change(update_character_setting_mode,
+                             inputs=[coser_setting, current_config],
+                             outputs=[coser_setting, current_config])
+        
+        kb_select.change(update_knowledge_base, 
+                         inputs=[kb_select, current_config], 
+                         outputs=[kb_select, current_config])
+        
+        # 勾选本地知识库时，若为角色扮演模式，提醒用户设置知识库目录
+        rag_switch.select(init_kb, inputs=[rag_switch, kb_select, current_config], outputs=rag_switch)
+        
+        clear_btn.click(
+            clear_chat,
+            inputs=section_state,
+            outputs=[chatbot, thinking_display, section_state],
+            queue=False
+        )
+        
+        save_btn.click(
+            save_chat,
+            inputs=section_state,
+            outputs=section_state,
+            queue=False
+        ).then(
+            fn=update_history,
+            inputs=section_state,
+            outputs=chat_history_select
+        )
+        
+        chat_history_select.change(load_chat, 
+                                   inputs=[section_state, chat_history_select], 
+                                   outputs=[chatbot, thinking_display, section_state])
+        
+        section_state.change(update_history, 
+                             inputs=section_state,
+                             outputs=chat_history_select)
+    print("===== 初始化完成 =====")
+    demo.launch(share=False)
+    

config.py

@@ -0,0 +1,114 @@
+from pathlib import Path
+import os
+
+DEFAULT_MODEL_NAME = "Qwen/Qwen3-0.6B"
+DEFAULT_MODE = "角色扮演"
+DEFAULT_C_SETTING_MODE = "by system"
+DEFAULT_COSPLAY_SETTING = 'rag/characters/周杰伦.txt'
+AVALIABLE_MODELS = [
+    "Qwen/Qwen3-8B",
+    "Qwen/Qwen3-4B",
+    "Qwen/Qwen3-1.7B",
+    "Qwen/Qwen3-0.6B"
+]
+BASE_MODEL_TABLE = {"qwen7B_jaychou_f16": "qwen2.5:7b-instruct", "qwen0.5B_jaychou13": "qwen2.5:0.5b-instruct",
+                    "qwen14B_jaychou_q8_newdata_add_template": "qwen2.5:14b-instruct",
+                    "qwen2.5_32B_jaychou": "qwen2.5:32b-instruct",
+                    "qwen2.5_0.5B_jaychou_lora": "qwen2.5:0.5b-instruct",
+                    # "qwen2.5_32B_jaychou_tq1": "qwen2.5:32b-instruct"
+                    }
+
+AVALIABLE_KNOWLEDGE_BASE = [
+    "rag/kb/BIGOLIVE及公司介绍",
+    "rag/kb/主播A的直播间对话数据",
+    "rag/kb/周杰伦",
+    "rag/kb/狼人杀"
+]
+SUPPORT_MODES = [
+    "角色扮演",
+    "普通模式",
+]
+CHARACTER_SETTING_MODES = [
+    "by system",
+    "by prompt"
+]
+
+EXAMPLES_changkong = [
+    """
+    [0:00:00]:[0:00:04]	對呀 我肯定沒有回
+    [0:00:04]:[0:00:08]	真的 今天我兒時頭 我感覺頭髮摘
+    [0:00:08]:[0:00:10]	而且我的頭髮越來越爽了
+    [0:00:10]:[0:00:13]	我頭髮越來越爽了
+    [0:00:13]:[0:00:15]	真的越來越少
+    [0:00:15]:[0:00:17]	好煩呀 我經常脫頭髮
+    """,
+    """
+    [0:01:44]:[0:01:46]	我咋覺得這個茶克斯那麼熟呢
+    [0:01:47]:[0:01:49]	哦 我想起來了
+    【观众536644926】
+    好像看过
+
+    [0:01:50]:[0:01:51]	上次跟我一起打BK那個
+    [0:01:52]:[0:01:53]	對面的那個
+    【观众1887407561】
+    嗨😂
+
+    [0:01:54]:[0:01:55]	是不是你
+    [0:01:55]:[0:01:56]	肯定是
+    [0:01:57]:[0:01:58]	我有點想起來了
+    [0:01:59]:[0:02:00]	因為他們老是在叫妳的名字
+    [0:02:01]:[0:02:02]	好像看過
+    [0:02:03]:[0:02:04]	哪一把BK來的
+    [0:02:08]:[0:02:10]	嗨 我記得你了
+    [0:02:11]:[0:02:12]	那個 讓我
+    [0:02:13]:[0:02:14]	讓我 讓我 讓我
+    [0:02:14]:[0:02:15]	做那個懲罰
+    [0:02:15]:[0:02:17]	你知道嗎 我差點進去醫院了
+    """,
+    
+]
+EXAMPLES_zhubo_clone = ["""最近好烦啊"""]
+EXAMPLES_langren = ["""玩家角色分配为一个预言家，一个女巫，一个猎人，三个平民，三个狼人。游戏开始，请开始主持。"""]
+EXAMPLES_jaychou_clone = [
+                """你是谁？""", 
+                """不忙的时候你会做些什么？""",
+                """你创作《素颜》这首歌的时候，背后有什么故事吗；这首歌里面有哪句歌词是你特别喜欢的""",
+                """你的香味一直徘徊,比我知道,秘密躺在我怀抱,还有没有人知道。 这是你的哪首歌？""",
+                ]
+EXAMPLES = EXAMPLES_jaychou_clone
+# 本地RAG
+RAG_TOP_K = 5
+
+# 联网
+MAX_RESULTS= 3
+
+# 目录信息
+BASE_DIR = Path("demo_dir")
+TEMP_DIR = BASE_DIR / "tmp"
+USER_DIR = BASE_DIR / "user"
+TEMP_DIR.mkdir(parents=True, exist_ok=True)
+USER_DIR.mkdir(parents=True, exist_ok=True)
+os.environ['GRADIO_TEMP_DIR'] = str(TEMP_DIR)
+
+CSS = """
+.rag-details:not([open]) > div { 
+    display: none !important;  /* 强制折叠状态 */
+    white-space: pre-wrap;     /* 保留换行符 */
+}
+.rag-details[open] summary::after {
+    content: "▼";
+    float: right;
+}
+"""
+
+# LLM最大上下文长度
+MAX_MODEL_CTX = 32768
+
+# 知识库embedding模型
+AVALIABLE_EMBEDDING_MODELS = [
+    "BAAI/bge-large-zh-v1.5",
+    "sentence-transformers/all-MiniLM-L12-v2",
+    "sentence-transformers/all-MiniLM-L6-v2",
+    "jinaai/jina-embeddings-v2-base-zh",
+]
+DEFAULT_EMBEDDING_MODEL = "BAAI/bge-large-zh-v1.5"

knowledge_base.py

@@ -0,0 +1,245 @@
+from langchain_community.document_loaders import DirectoryLoader, JSONLoader, UnstructuredMarkdownLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter, MarkdownTextSplitter, MarkdownHeaderTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from pathlib import Path
+from transformers import AutoModel, AutoTokenizer
+import torch
+import torch.nn.functional as F
+import numpy as np  
+import config as cfg 
+
+
+
+class LocalRAG:
+    def __init__(self, 
+                 rag_top_k=3,
+                 doc_dir="rag/kb/BIGOLIVE及公司介绍/",    # 默认加载这个，若选择角色扮演模式，可根据角色选择
+                 vector_db_path="rag/vector_db/", 
+                 embed_model=cfg.DEFAULT_EMBEDDING_MODEL
+                 ):
+        self.rag_top_k = rag_top_k
+        self.doc_dir = doc_dir  # 本地知识库的文档目录
+        self.vector_db_path = vector_db_path  # 向量数据库存储路径
+        self.embed_model = embed_model
+        self.build_vector_db()
+        
+    def build_vector_db(self):
+        # 加载文档（支持PDF、TXT、DOCX）
+        if isinstance(self.doc_dir, list):
+            general_docs = []
+            json_docs = []
+            md_docs = []
+            for doc_dir in self.doc_dir:
+                # 处理一般文件，txt等
+                loader = DirectoryLoader(doc_dir, glob="**/*.[!json!md]*")  # "**/[!.]*"
+                tmp_docs = loader.load()
+                general_docs.extend(tmp_docs)
+                # 额外处理json文件
+                for json_file in Path(doc_dir).rglob("*.json"):
+                    loader = JSONLoader(
+                        file_path=str(json_file),
+                        jq_schema=".[] | {spk: .spk, text: .text}",
+                        text_content=False)
+                    
+                    data = loader.load()
+                    for iidx in range(len(data)):
+                        data[iidx].page_content = bytes(data[iidx].page_content, "utf-8").decode("unicode_escape")
+                    json_docs.extend(data)
+                
+                # 额外处理md文件
+                headers_to_split_on = [
+                    ("#", "Header 1"),
+                    ("##", "Header 2"),
+                    ("###", "Header 3"),
+                ]
+                for md_file in Path(doc_dir).rglob("*.md"):
+                    with open(md_file, 'r') as f:
+                        content = f.read()
+                    
+
+                    # 定义拆分器，拆分markdown内容
+                    markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
+                    md_header_splits = markdown_splitter.split_text(content)
+                    md_docs.extend(md_header_splits)
+                    
+                    # loader = UnstructuredMarkdownLoader(md_file, mode="elements")
+                    # data = loader.load()
+                    # docs.extend(data)
+        else:
+            loader = DirectoryLoader(self.doc_dir, glob="**/*.*")
+            docs = loader.load()
+
+        # 文本分块
+        if len(general_docs) > 0:
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=500,
+                chunk_overlap=50
+            )
+            chunks = text_splitter.split_documents(docs)
+        else:
+            chunks = json_docs + md_docs
+
+        # 生成向量并构建FAISS数据库
+        embeddings = HuggingFaceEmbeddings(model_name=self.embed_model)
+        self.vector_db = FAISS.from_documents(chunks, embeddings)
+        self.vector_db.save_local(self.vector_db_path)
+    
+    def reload_knowledge_base(self, target_doc_dir):
+        self.doc_dir = target_doc_dir
+        self.build_vector_db()
+    
+    # def reset(self):
+    #     self.vector_db = None
+
+
+class LocalRAG_new:
+    
+    def __init__(self, 
+                 rag_top_k=3,
+                 doc_dir="rag/kb/BIGOLIVE及公司介绍/",    # 默认加载这个，若选择角色扮演模式，可根据角色选择
+                 vector_db_path="rag/vector_db/", 
+                 embed_model_path="princeton-nlp/sup-simcse-bert-large-uncased",
+                 device=torch.device('cuda:2')):
+        self.rag_top_k = rag_top_k
+        self.doc_dir = doc_dir  # 本地知识库的文档目录
+        self.kb_name = '_'.join([Path(doc_dir[i]).name for i in range(len(doc_dir))])
+        self.embed_model_name = Path(embed_model_path).name
+        self.vector_db_path = vector_db_path  # 向量数据库存储路径
+        self.embed_model = embed_model_path
+        
+        self.device = device
+        # 加载分词器和模型
+        self.tokenizer = AutoTokenizer.from_pretrained(self.embed_model)
+        self.embed_model = AutoModel.from_pretrained(self.embed_model).to(device)
+        self.vector_db = None
+        self._vector_db = None
+        self.build_vector_db()
+    
+    class VectorDB:
+        def __init__(self, rag):
+            self._data = rag._vector_db
+            self.rag = rag
+            
+        def similarity_search(self, query, k):
+            # 可能的输入预处理，暂无
+            # query = input_optimize(query)
+            
+            # 计算query的embedding并与库中比较
+            with torch.inference_mode():
+                query_token = self.rag.tokenizer(query, padding=True, truncation=False, return_tensors="pt").to(self.rag.device)
+                query_embed = self.rag.embed_model(**query_token)['last_hidden_state'].mean(dim=1)
+                sim_query = F.cosine_similarity(query_embed.repeat(len(self._data['embeds']), 1), self._data['embeds'], dim=1, eps=1e-8)
+                max_ids_query = torch.argsort(sim_query, descending=True)[:self.rag.rag_top_k].cpu().detach().numpy()
+                return list(zip(np.array(self._data['chunks'])[max_ids_query], sim_query[max_ids_query]))
+        
+    def build_vector_db(self):
+        # 加载文档（支持PDF、TXT、DOCX）
+        if isinstance(self.doc_dir, list):
+            docs = []
+            for doc_dir in self.doc_dir:
+                loader = DirectoryLoader(doc_dir, glob="**/*.[!json!md]*")  # "**/[!.]*"
+                tmp_docs = loader.load()
+                docs.extend(tmp_docs)
+                # # 额外处理json文件
+                # for json_file in Path(doc_dir).rglob("*.json"):
+                #     loader = JSONLoader(
+                #         file_path=str(json_file),
+                #         jq_schema='.messages[].content',
+                #         text_content=False)
+                    
+                #     data = loader.load()
+                # 额外处理md文件
+                headers_to_split_on = [
+                    ("#", "Header 1"),
+                    ("##", "Header 2"),
+                    ("###", "Header 3"),
+                ]
+                for md_file in Path(doc_dir).rglob("*.md"):
+                    with open(md_file, 'r') as f:
+                        content = f.read()
+                    
+
+                    # 定义拆分器，拆分markdown内容
+                    markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
+                    md_header_splits = markdown_splitter.split_text(content)
+                    docs.extend(md_header_splits)
+                    
+                    # loader = UnstructuredMarkdownLoader(md_file, mode="elements")
+                    # data = loader.load()
+                    # docs.extend(data)
+        else:
+            loader = DirectoryLoader(self.doc_dir, glob="**/*.*")
+            docs = loader.load()
+
+        # 文本分块
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=500,
+            chunk_overlap=50
+        )
+        chunks = text_splitter.split_documents(docs)
+        with torch.inference_mode():
+            chunk_and_embed = []
+            for chunk in chunks:
+                chunk_token = self.tokenizer(chunk.page_content, padding=True, truncation=False, return_tensors="pt").to(self.device)
+                chunk_embed = self.embed_model(**chunk_token)['last_hidden_state'].mean(dim=1)
+                chunk_and_embed.append((chunk, chunk_embed))
+            all_chunks, all_embeds = list(zip(*chunk_and_embed))
+            all_chunks, all_embeds = list(all_chunks), list(all_embeds)
+            all_embeds = torch.cat(all_embeds, dim=0)
+            self._vector_db = {'chunks': all_chunks, 'embeds': all_embeds}
+            self.vector_db = self.VectorDB(self)
+            
+            torch.save(self.vector_db, str(Path(self.vector_db_path) / f'{self.kb_name}_{self.embed_model_name}.pt'))
+    
+    def reload_knowledge_base(self, target_doc_dir):
+        self.doc_dir = target_doc_dir
+        self.build_vector_db()
+    
+    # def reset(self):
+    #     self.vector_db = None
+
+
+class CosPlayer:
+    def __init__(self, description_file):
+        self.update(description_file)
+    
+    def update(self, description_file):
+        self.description_file = description_file
+        with open(description_file, 'r') as f:
+            all_lines = f.readlines()
+        self.core_setting = ''.join(all_lines)
+        self.characters_dir = Path(description_file).parent
+        self.prologue_file = self.description_file.replace('/characters/', '/prologues/')
+        if not Path(self.prologue_file).exists():
+            self.prologue_file = None
+             
+    def get_all_characters(self):
+        return [str(i) for i in list(self.characters_dir.rglob('*.txt'))]
+    
+    def get_core_setting(self):
+        return self.core_setting
+        
+    def get_prologue(self):
+        if self.prologue_file:
+            with open(self.prologue_file, 'r') as f:
+                all_lines = f.readlines()
+            return ''.join(all_lines)
+        else:
+            return None
+    
+        
+if __name__ == "__main__":
+    rag = LocalRAG()
+    # # rag.build_vector_db()
+    # doc_dir = "rag/debug"
+    # loader = DirectoryLoader(doc_dir, glob="**/*.*")
+    # docs = loader.load()
+
+    # # 文本分块
+    # text_splitter = RecursiveCharacterTextSplitter(
+    #     chunk_size=500,
+    #     chunk_overlap=50
+    # )
+    # chunks = text_splitter.split_documents(docs)
+    # pass

rag/kb/周杰伦/周杰伦全部歌曲.md

@@ -0,0 +1,341 @@
+# 周杰伦全部歌曲
+## 十一月的萧邦
+### 浪漫手机
+发行时间:2005.11
+### 漂移
+发行时间:2005.11
+### 一路向北
+发行时间:2005.11
+### 枫
+发行时间:2005.11
+### 黑色毛衣
+发行时间:2005.11
+### 麦芽糖
+发行时间:2005.11
+### 夜曲
+发行时间:2005.11
+### 发如雪
+发行时间:2005.11
+### 蓝色风暴
+发行时间:2005.11
+### 珊瑚海
+发行时间:2005.11
+### 四面楚歌
+发行时间:2005.11
+### 逆鳞
+发行时间:2005.11
+## 叶惠美
+### 东风破
+发行时间:2003.7
+### 三年二班
+发行时间:2003.7
+### 晴天
+发行时间:2003.7
+### 你听得到
+发行时间:2003.7
+### 同一种调调
+发行时间:2003.7
+### 她的睫毛
+发行时间:2003.7
+### 以父之名
+发行时间:2003.7
+### 爱情悬崖
+发行时间:2003.7
+### 懦夫
+发行时间:2003.7
+### 梯田
+发行时间:2003.7
+### 双刀
+发行时间:2003.7
+## 魔杰座
+### 给我一首歌的时间
+发行时间:2008.10
+### 乔克叔叔
+发行时间:2008.10
+### 时光机
+发行时间:2008.10
+### 说好的幸福呢
+发行时间:2008.10
+### 稻香
+发行时间:2008.10
+### 龙战骑士
+发行时间:2008.10
+### 花海
+发行时间:2008.10
+### 蛇舞
+发行时间:2008.10
+### 兰亭序
+发行时间:2008.10
+### 流浪诗人
+发行时间:2008.10
+### 魔术先生
+发行时间:2008.10
+## 八度空间
+### 火车叨位去
+发行时间:2002.7
+### 回到过去
+发行时间:2002.7
+### 爷爷泡的茶
+发行时间:2002.7
+### 半兽人
+发行时间:2002.7
+### 暗号
+发行时间:2002.7
+### 最后的战役
+发行时间:2002.7
+### 龙拳
+发行时间:2002.7
+### 米兰的小铁匠
+发行时间:2002.7
+### 分裂
+发行时间:2002.7
+### 半岛铁盒
+发行时间:2002.7
+## 十二新作
+### 哪里都是你
+发行时间:2012.12
+### 公公偏头痛
+发行时间:2012.12
+### 手语
+发行时间:2012.12
+### 乌克丽丽
+发行时间:2012.12
+### 红尘客栈
+发行时间:2012.12
+### 梦想启动
+发行时间:2012.12
+### 四季列车
+发行时间:2012.12
+### 大笨钟
+发行时间:2012.12
+### 爱你没差
+发行时间:2012.12
+### 傻笑
+发行时间:2012.12
+### 明明就
+发行时间:2012.12
+### 比较大的大提琴
+发行时间:2012.12
+## 七里香
+### 七里香
+发行时间:2004.8
+### 外婆
+发行时间:2004.8
+### 困兽之斗
+发行时间:2004.8
+### 我的地盘
+发行时间:2004.8
+### 借口
+发行时间:2004.8
+### 园游会
+发行时间:2004.8
+### 止战之殇
+发行时间:2004.8
+### 乱舞春秋
+发行时间:2004.8
+### 将军
+发行时间:2004.8
+### 搁浅
+发行时间:2004.8
+## 惊叹号
+### 惊叹号
+发行时间:2011.11
+### 公主病
+发行时间:2011.11
+### 琴伤
+发行时间:2011.11
+### 周杰伦
+发行时间:2011.11
+### 皮影戏
+发行时间:2011.11
+### 超跑女神
+发行时间:2011.11
+### 水手怕水
+发行时间:2011.11
+### 世界未末日
+发行时间:2011.11
+### 迷魂曲
+发行时间:2011.11
+### 你好吗
+发行时间:2011.11
+### 疗伤烧肉粽
+发行时间:2011.11
+## 跨时代
+### 免费教学录影带
+发行时间:2010.5
+### 爱的飞行日记
+发行时间:2010.5
+### 嘻哈空姐
+发行时间:2010.5
+### 超人不会飞
+发行时间:2010.5
+### 自导自演
+发行时间:2010.5
+### 跨时代
+发行时间:2010.5
+### 说了再见
+发行时间:2010.5
+### 烟花易冷
+发行时间:2010.5
+### 我落泪
+发行时间:2010.5
+### 雨下一整晚
+发行时间:2010.5
+### 好久不见
+发行时间:2010.5
+## 我很忙
+### 最长的电影
+发行时间:2007.11
+### 甜甜的
+发行时间:2007.11
+### 青花瓷
+发行时间:2007.11
+### 我不配
+发行时间:2007.11
+### 牛仔很忙
+发行时间:2007.11
+### 无双
+发行时间:2007.11
+### 彩虹
+发行时间:2007.11
+### 阳光宅男
+发行时间:2007.11
+### 蒲公英的约定
+发行时间:2007.11
+### 扯
+发行时间:2007.11
+## 依然范特西
+### 听妈妈的话
+发行时间:2006.9
+### 菊花台
+发行时间:2006.9
+### 退后
+发行时间:2006.9
+### 本草纲目
+发行时间:2006.9
+### 夜的第七章
+发行时间:2006.9
+### 迷迭香
+发行时间:2006.9
+### 千里之外
+发行时间:2006.9
+### 心雨
+发行时间:2006.9
+### 红模仿
+发行时间:2006.9
+### 白色风车
+发行时间:2006.9
+## Jay
+### 反方向的钟
+发行时间:2000.11
+### 印地安老斑鸠
+发行时间:2000.11
+### 完美主义
+发行时间:2000.11
+### 黑色幽默
+发行时间:2000.11
+### 可爱女人
+发行时间:2000.11
+### 伊斯坦堡
+发行时间:2000.11
+### 斗牛
+发行时间:2000.11
+### 娘子
+发行时间:2000.11
+### 龙卷风
+发行时间:2000.11
+### 星晴
+发行时间:2000.11
+## 非专辑
+## 范特西
+### 简单爱
+发行时间:2001.9
+### 对不起
+发行时间:2001.9
+### 忍者
+发行时间:2001.9
+### 上海一九四三
+发行时间:2001.9
+### 开不了口
+发行时间:2001.9
+### 周杰伦
+发行时间:2001.9
+### 威廉古堡
+发行时间:2001.9
+### 爱在西元前
+发行时间:2001.9
+### 双截棍
+发行时间:2001.9
+### 安静
+发行时间:2001.9
+## 哎呦，不錯哦
+### 阳明山
+发行时间:2014.12
+### 窃爱
+发行时间:2014.12
+### 算什么男人
+发行时间:2014.12
+### 天涯过客
+发行时间:2014.12
+### 怎么了
+发行时间:2014.12
+### 一口气全念对
+发行时间:2014.12
+### 我要夏天
+发行时间:2014.12
+### 手写的从前
+发行时间:2014.12
+### 鞋子特大号
+发行时间:2014.12
+### 听爸爸的话
+发行时间:2014.12
+### 美人鱼
+发行时间:2014.12
+### 听见下雨的声音
+发行时间:2014.12
+## 周杰伦的床边故事
+### 床边故事
+发行时间:2016.6
+### 说走就走
+发行时间:2016.6
+### 一点点
+发行时间:2016.6
+### 前世情人
+发行时间:2016.6
+### 英雄
+发行时间:2016.6
+### 不该
+发行时间:2016.6
+### 土耳其冰淇淋
+发行时间:2016.6
+### 告白气球
+发行时间:2016.6
+### Now You See Me
+发行时间:2016.6
+### 爱情废柴
+发行时间:2016.6
+## 最伟大的作品
+### Intro
+发行时间:2022.7
+### 最伟大的作品
+发行时间:2022.7
+### 还在流浪
+发行时间:2022.7
+### 说好不哭
+发行时间:2022.7
+### 红颜如霜
+发行时间:2022.7
+### 不爱我就拉倒
+发行时间:2022.7
+### Mojito
+发行时间:2022.7
+### 错过的烟火
+发行时间:2022.7
+### 等你下课
+发行时间:2022.7
+### 粉色海洋
+发行时间:2022.7
+### 倒影
+发行时间:2022.7
+### 我是如此相信
+发行时间:2022.7

rag/kb/周杰伦/周杰伦基本资料.md

@@ -0,0 +1,153 @@
+# 周杰伦  
+周杰伦（1979年1月18日—），台湾创作男歌手、演员、词曲作家及制作人，被誉为“中文流行天王”，对华语乐坛及全球华人社群具有深远影响。  
+
+## 基本信息  
+| 项目         | 内容                                                                 |  
+|--------------|----------------------------------------------------------------------|  
+| **英文名**   | Jay Chou                                                             |  
+| **昵称**     | 周董、小公举[[1]]                                                    |  
+| **出生**     | 1979年1月18日（46岁），台湾台北县林口乡（现新北市林口区）[[2]]       |  
+| **居住地**   | 台湾台北市大安区                                                     |  
+| **职业**     | 歌手、词曲作家、音乐制作人、演员、导演、企业家                       |  
+| **语言**     | 国语、台语、英语                                                     |  
+| **母校**     | 淡江中学                                                             |  
+| **宗教信仰** | 基督新教[[3]]                                                        |  
+| **配偶**     | 昆凌（2015年结婚）[[4]]                                              |  
+| **儿女**     | 2女1子                                                               |  
+| **出道日期** | 2000年11月6日，出道作品《杰伦》                                      |  
+
+## 早年经历  
+1. **家庭背景**  
+   - 父亲周耀中为生物教师，母亲叶惠美为美术教师，14岁父母离异后随母亲生活[[14][15][16][17]]。  
+   - 澄清《爸，我回来了》并非指涉父母家暴，而是社会现象的感慨[[18]]。  
+
+2. **教育与成长**  
+   - 3岁学钢琴，国中热爱篮球，高中就读淡江中学音乐科，奠定音乐基础[[19][20][21]]。  
+   - 大学联考落榜，因僵直性脊椎炎免服兵役[[22][23]]。  
+
+3. **音乐启蒙**  
+   - 受张学友、肖邦、李恕权、史帝夫·汪达影响，创作风格融合古典与流行[[24][25]]。  
+
+## 音乐事业  
+### 重要阶段  
+#### 1997年—2001年：起步与突破  
+- 1997年通过《超级新人王》被吴宗宪发掘，开启作曲生涯。
+- 1998年发表首支创作歌曲《三暝三日》（吴宗宪演唱）。  
+- 1999年与方文山合作《落雨声》（江蕙演唱），奠定合作基础[[21]]。  
+
+#### 2000年—2003年：奠定地位  
+- 2000年首张专辑《Jay》融合R&B、Hip-Hop与中国风，开创“周式曲风”[[21]]。  
+- 2001年《范特西》获金曲奖5项大奖，成为首位获“最佳专辑制作人”的新人[[26][27]]。  
+- 2003年《叶惠美》主打歌《以父之名》引发“周杰伦日”，全球8亿人同步收听[[28]]。  
+
+#### 2004年—2007年：国际影响力  
+- 2004年《七里香》获世界音乐大奖“大中华区最畅销艺人”[[29]]。  
+- 2005年《十一月的肖邦》尝试MV执导，2007年成立杰威尔音乐[[30]]。  
+
+#### 2008年—2012年：跨界与巡演  
+- 2008年《魔杰座》在日本武道馆开唱，成为第二位华人歌手[[31]]。  
+- 2011年与科比合作雪碧广告曲《天地一斗》，进军国际[[33]]。  
+
+#### 2014年—2019年：持续创新  
+- 2016年《周杰伦的床边故事》主打歌《告白气球》创公告牌纪录[[34]]。  
+- 2019年《说好不哭》与五月天阿信合唱，YouTube播放量破千万[[35]]。  
+
+#### 2020年至今：新专辑与合作  
+- 2022年《最伟大的作品》获IFPI全球最畅销专辑冠军[[36][37]]。  
+- 2023年《圣诞星》MV 13小时破140万次播放[[41]]。  
+
+## 电影事业  
+### 代表作品  
+1. **《头文字D》（2005年）**  
+   - 首登大银幕，获香港金像奖、金马奖双料最佳新演员[[48][49]]。  
+
+2. **《满城尽带黄金甲》（2006年）**  
+   - 饰演二王子元杰，提名香港金像奖最佳男配角[[56][63]]。  
+
+3. **《不能说的·秘密》（2007年）**  
+   - 自编自导自演，融合音乐创作，获韩国首映好评[[68]]。  
+
+4. **《青蜂侠》（2011年）**  
+   - 好莱坞处女作，演唱片尾曲《双截棍》[[65][67]]。  
+
+5. **《天台爱情》（2013年）**  
+   - 第二部自导自演电影，获张艺谋赞赏[[68]]。  
+
+## 副业  
+### 潮流事业  
+- 2006年创立潮流品牌**PHANTACi**，2016年推出电商平台**J Concept星品库**[[70][72]]。  
+
+### 电子竞技  
+- 2016年收购台北暗杀星战队，更名**J Team**，并开设“魔杰电竞馆”[[73]]。  
+
+### 餐饮与投资  
+- 开设“Mr.J藤原豆腐店”“J POT HOTPOT火锅料理”等餐饮品牌[[74]]。  
+- 入股中国数码文化集团，参与耳机品牌**1MORE万魔**[[75]]。  
+
+## 个人生活  
+### 婚姻与家庭  
+- 2015年与昆凌结婚，育有两女一子，三次婚礼分别在英国、台湾、澳大利亚举办[[83][84][85]]。  
+- 夫妻均为基督徒，婚��于英国塞尔比修道院举行[[89][92]]。  
+
+### 绯闻与争议  
+- 早年与蔡依林、侯佩岑等传绯闻，均否认[[75][76][77]]。  
+- 与昆凌恋情于2014年公开，称其“善良、贴心”[[78][80]]。  
+
+## 荣誉与奖项  
+### 重要荣誉  
+- 2003年《时代》杂志亚洲版封面人物[[21]]。  
+- 2009年CNN“亚洲最具影响力25人”[[12]]。  
+- 2010年小行星**257248**命名为“周杰伦星”[[38]]。  
+- 2023年《最伟大的作品》全球专辑销量冠军[[108]]。  
+
+### 主要奖项  
+- **金曲奖**：15座（含最佳专辑、最佳作曲人等）[[10][111][112]]。  
+- **世界音乐大奖**：4次“大中华区最畅销艺人”[[29][32]]。  
+- **香港金像奖**：最佳新演员（《头文字D》）[[48]]。  
+
+## 作品列表  
+### 音乐专辑  
+《Jay》（2000年）  
+《范特西》（2001年）  
+《八度空间》（2002年）  
+《叶惠美》（2003年）  
+《七里香》（2004年）  
+《十一月的肖邦》（2005年）  
+《依然范特西》（2006年）  
+《我很忙》（2007年）  
+《魔杰座》（2008年）  
+《跨时代》（2010年）  
+《惊叹号》（2011年）  
+《十二新作》（2012年）  
+《哎呦，不错哦》（2014年）  
+《周杰伦的床边故事》（2016年）  
+《最伟大的作品》（2022年）  
+
+### 书籍作品  
+| 出版日期   | 书名                                  | 作者       | 出版社               |  
+|------------|---------------------------------------|------------|----------------------|  
+| 2002年     | 《半岛铁盒》                          | 方文山     | 华人版图出版社       |  
+| 2004年     | 《D调的华丽》                         | 周杰伦     | 华人版图出版社       |  
+| 2007年     | 《不能说的秘密电影创作琴谱》          | 周杰伦     | 华人版图出版社       |  
+
+### 演出作品（电影）  
+| 年份 | 片名                 | 角色       | 备注                           |  
+|------|----------------------|------------|--------------------------------|  
+| 2005 | 《头文字D》          | 藤原拓海   | 首部主演电影，获双料新人奖     |  
+| 2006 | 《满城尽带黄金甲》   | 元杰       | 国际发行，提名香港金像奖       |  
+| 2007 | 《不能说的·秘密》    | 叶湘伦     | 自编自导自演                   |  
+| 2011 | 《青蜂侠》           | 加藤（Kato）| 好莱坞处女作，演唱片尾曲       |  
+
+## 注释与参考资料  
+1. 注1：现今的新北市林口区。  
+2. 注2：淡江中学音乐科为新设科系，因错过华冈艺校报名而就读[[21]]。  
+
+[1] 周杰倫綽號大揭密 被叫「小公舉」是因.... 自由时报. 2016-01-20.  
+[2] 藝人專區 周杰倫 JAY CHOU. 杰威尔音乐有限公司. 2024-12-14.  
+[3] 周杰倫昆凌喜迎小王子 牧師祝福:「多子多孫，生養眾多!」. 基督日报. 2017-06-23.  
+
+
+## 外部链接  
+- [杰威尔音乐官方网站](https://www.jvrmusic.com)  
+- [周杰伦YouTube频道](https://www.youtube.com/user/JayChtV)  
+- [周杰伦Instagram](https://www.instagram.com/jaychou)  

tts_api.py

@@ -0,0 +1,274 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import soundfile as sf
+from xcodec2.modeling_xcodec2 import XCodec2Model
+import numpy as np
+import ChatTTS
+import re
+DEFAULT_TTS_MODEL_NAME = "HKUSTAudio/LLasa-1B"
+DEMO_EXAMPLES = [
+    ["太乙真人.wav", "对，这就是我万人敬仰的太乙真人，虽然有点婴儿肥，但也掩不住我逼人的帅气。"],
+    ["邓紫棋.wav", "特别大的不同，因为以前在香港是过年的时候,我们可能见到的亲戚都是爸爸那边的亲戚"],
+    ["雷军.wav", "这是个好问题，我把来龙去脉给你简单讲，就是这个社会对小米有很多的误解，有很多的误解，呃，也能理解啊，就是小米这个模式呢"],
+    ["周杰伦.wav", "但如果你这兴趣可以得到很大的回响，那会更开心"],
+    ["Taylor Swift.wav", "It's actually uh, it's a concept record, but it's my first directly autobiographical album in a while because the last album that I put out was, uh, a rework."]
+]
+class TTSapi:
+    def __init__(self, 
+                 model_name=DEFAULT_TTS_MODEL_NAME, 
+                 codec_model_name="HKUST-Audio/xcodec2", 
+                 device=torch.device("cuda:0")):
+
+        self.reload(model_name, codec_model_name, device)
+        
+    def reload(self, 
+               model_name=DEFAULT_TTS_MODEL_NAME, 
+               codec_model_name="HKUST-Audio/xcodec2",
+               device=torch.device("cuda:0")):
+        if 'llasa' in model_name.lower():
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model = AutoModelForCausalLM.from_pretrained(model_name)
+            self.model.eval().to(device)
+    
+            self.codec_model = XCodec2Model.from_pretrained(codec_model_name)
+            self.codec_model.eval().to(device)
+            self.device = device
+            self.codec_model_name = codec_model_name
+            self.sr = 16000
+        elif 'chattts' in model_name.lower():
+            self.model = ChatTTS.Chat()
+            self.model.load(compile=False) # Set to True for better performance but would l significantly reduce the loading speed
+            self.sr = 24000
+            self.punctuation = r'[,，.。？?!！~～；;]'
+        else:
+            raise ValueError(f'不支持的TTS模型：{model_name}')
+        
+        self.model_name = model_name
+        
+    def ids_to_speech_tokens(self, speech_ids):
+        speech_tokens_str = []
+        for speech_id in speech_ids:
+            speech_tokens_str.append(f"<|s_{speech_id}|>")
+        return speech_tokens_str
+
+    def extract_speech_ids(self, speech_tokens_str):
+        speech_ids = []
+        for token_str in speech_tokens_str:
+            if token_str.startswith('<|s_') and token_str.endswith('|>'):
+                num_str = token_str[4:-2]
+
+                num = int(num_str)
+                speech_ids.append(num)
+            else:
+                print(f"Unexpected token: {token_str}")
+        return speech_ids
+
+ 
+    def forward(self, input_text, speech_prompt=None, save_path='wavs/generated/gen.wav'):
+        #TTS start!
+        with torch.no_grad():
+            if 'chattts' in self.model_name.lower():
+                # rand_spk = chat.sample_random_speaker()
+                # print(rand_spk) # save it for later timbre recovery
+
+                # params_infer_code = ChatTTS.Chat.InferCodeParams(
+                #     spk_emb = rand_spk, # add sampled speaker 
+                #     temperature = .3,   # using custom temperature
+                #     top_P = 0.7,        # top P decode
+                #     top_K = 20,         # top K decode
+                # )
+                break_num = max(min(len(re.split(self.punctuation, input_text)), 7), 2)
+                params_refine_text = ChatTTS.Chat.RefineTextParams(
+                    prompt=f'[oral_2][laugh_0][break_{break_num}]',
+                )
+                wavs = self.model.infer([input_text],
+                                params_refine_text=params_refine_text,
+                            )
+                gen_wav_save = wavs[0]
+                sf.write(save_path, gen_wav_save, 24000)
+                
+            else:   
+                if speech_prompt:
+                    # only 16khz speech support!
+                    prompt_wav, sr = sf.read(speech_prompt)   # you can find wav in Files
+                    prompt_wav = torch.from_numpy(prompt_wav).float().unsqueeze(0)
+                    
+                    # Encode the prompt wav
+                    vq_code_prompt = self.codec_model.encode_code(input_waveform=prompt_wav)
+                    print("Prompt Vq Code Shape:", vq_code_prompt.shape )   
+
+                    vq_code_prompt = vq_code_prompt[0,0,:]
+                    # Convert int 12345 to token <|s_12345|>
+                    speech_ids_prefix = self.ids_to_speech_tokens(vq_code_prompt)
+                else:
+                    speech_ids_prefix = ''
+                formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
+
+                # Tokenize the text ( and the speech prefix)
+                chat = [
+                    {"role": "user", "content": "Convert the text to speech:" + formatted_text},
+                    {"role": "assistant", "content": "<|SPEECH_GENERATION_START|>" + ''.join(speech_ids_prefix)}
+                ]
+
+                input_ids = self.tokenizer.apply_chat_template(
+                    chat, 
+                    tokenize=True, 
+                    return_tensors='pt', 
+                    continue_final_message=True
+                )
+                input_ids = input_ids.to(self.device)
+                speech_end_id = self.tokenizer.convert_tokens_to_ids('<|SPEECH_GENERATION_END|>')
+
+                # Generate the speech autoregressively
+                outputs = self.model.generate(
+                    input_ids,
+                    max_length=2048,  # We trained our model with a max length of 2048
+                    eos_token_id= speech_end_id ,
+                    do_sample=True,    
+                    top_p=1,           #  Adjusts the diversity of generated content
+                    temperature=1,   #  Controls randomness in output
+                )
+                # Extract the speech tokens
+                generated_ids = outputs[0][input_ids.shape[1] - len(speech_ids_prefix):-1]
+
+                speech_tokens = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)   
+
+                # Convert  token <|s_23456|> to int 23456 
+                speech_tokens = self.extract_speech_ids(speech_tokens)
+
+                speech_tokens = torch.tensor(speech_tokens).to(self.device).unsqueeze(0).unsqueeze(0)
+
+                # Decode the speech tokens to speech waveform
+                gen_wav = self.codec_model.decode_code(speech_tokens) 
+                
+                # if only need the generated part
+                if speech_prompt:
+                    gen_wav = gen_wav[:,:,prompt_wav.shape[1]:]
+                
+                gen_wav_save = gen_wav[0, 0, :].cpu().numpy()
+                sf.write(save_path, gen_wav_save, 16000)
+            # gen_wav_save = np.clip(gen_wav_save, -1, 1)
+            # gen_wav_save = (gen_wav_save * 32767).astype(np.int16)
+            return gen_wav_save
+
+
+if __name__ == '__main__':
+    # Llasa-8B shows better text understanding ability.
+
+    # input_text = " He shouted, 'Everyone, please gather 'round! Here's the plan: 1) Set-up at 9:15 a.m.; 2) Lunch at 12:00 p.m. (please RSVP!); 3) Playing — e.g., games, music, etc. — from 1:15 to 4:45; and 4) Clean-up at 5 p.m.'"
+    # prompt_text ="对，这就是我万人敬仰的太乙真人，虽然有点婴儿肥，但也掩不住我逼人的帅气。"
+    # input_text = prompt_text + '嘻嘻，臭宝儿你真可爱，我好喜欢你呀。'
+    # save_root = 'wavs/generated/'
+    # save_path = save_root + 'test.wav'
+    # speech_ref = 'wavs/ref/太乙真人.wav'
+    # # speech_ref = None
+    # # 帘外雨潺潺，春意阑珊。罗衾不耐五更寒。梦里不知身是客，一晌贪欢。独自莫凭栏，无限江山。别时容易见时难。流水落花春去也，天上人间。
+    # llasa_tts = TTSapi()
+    # gen = llasa_tts.forward(input_text, speech_prompt=speech_ref, save_path=save_path)
+    # print(gen.shape)
+    import gradio as gr
+    synthesiser = TTSapi()
+    TTS_LOADED = True
+    def predict(config):
+        global TTS_LOADED, synthesiser
+        print(f"待合成文本：{config['msg']}")
+        print(f"选中TTS模型：{config['tts_model']}")
+        print(f"参考音频路径：{config['ref_audio']}")
+        print(f"参考音频文本：{config['ref_audio_transcribe']}")
+        text = config['msg']
+        try:
+            if len(text) == 0:
+                audio_output = np.array([0], dtype=np.int16)
+                print("输入为空，无法合成语音") 
+            else:
+                if not TTS_LOADED:
+                    print('TTS模型首次加载...')
+                    gr.Info("初次加载TTS模型，请稍候..", duration=63)
+                    synthesiser = TTSapi(model_name=config['tts_model'])#, device="cuda:2")
+                    TTS_LOADED = True
+                    print('加载完毕...')
+                # 检查当前模型是否是所选
+                if config['tts_model'] != synthesiser.model_name:
+                    print(f'当前TTS模型{synthesiser.model_name}非所选，重新加载')
+                    synthesiser.reload(model_name=config['tts_model'])
+
+                # 如果提供了参考音频，则需把参考音频的文本加在response_content前面作为前缀
+                if config['ref_audio']:
+                    prompt_text = config['ref_audio_transcribe']
+                    if prompt_text is None:
+                        # prompt_text = ...
+                        raise NotImplementedError('暂时必须提供文本')  # TODO：考虑后续加入ASR模型
+                    text = prompt_text + text
+                    
+                audio_output = synthesiser.forward(text, speech_prompt=config['ref_audio'])
+            
+        except Exception as e:
+            print('!!!!!!!!')
+            print(e)
+            print('!!!!!!!!')
+        
+        return (synthesiser.sr if synthesiser else 16000, audio_output)
+
+    with gr.Blocks(title="TTS Demo", theme=gr.themes.Soft(font=["sans-serif", "Arial"])) as demo:
+        gr.Markdown("""
+        # Personalized TTS Demo
+        ## 使用步骤
+        * 上传你想要合成的目标说话人的语音，10s左右即可，并在下面填入对应的文本。或直接点击下方示例
+        * 输入你想要合成的文字，点击合成语音按钮，稍等片刻即可
+        
+        """)
+        with gr.Row():
+            with gr.Column():
+                # TTS模型选择
+                tts_model = gr.Dropdown(
+                    label="选择TTS模型",
+                    choices=["ChatTTS", "HKUSTAudio/LLasa-1B", "HKUSTAudio/LLasa-3B", "HKUSTAudio/LLasa-8B"],
+                    value=DEFAULT_TTS_MODEL_NAME,
+                    interactive=True,
+                    visible=False  # 给产品演示，暂时不展示模型选择
+                )
+                
+                # 参考音频上传
+                ref_audio = gr.Audio(
+                    label="上传参考音频",
+                    type="filepath",
+                    interactive=True
+                )
+                ref_audio_transcribe = gr.Textbox(label="参考音频对应文本", visible=True)
+                # 创建示例数据
+                examples = gr.Examples(
+                    examples=DEMO_EXAMPLES,
+                    inputs=[ref_audio, ref_audio_transcribe],
+                    fn=predict
+                )
+        
+            with gr.Column():
+                audio_player = gr.Audio(
+                    label="听听我声音~",
+                    type="numpy",
+                    interactive=False
+                )
+                msg = gr.Textbox(label="输入文本", placeholder="请输入想要合成的文本")
+                submit_btn = gr.Button("合成语音", variant="primary")
+                        
+        current_config = gr.State({
+            "msg": None,
+            "tts_model": DEFAULT_TTS_MODEL_NAME,
+            "ref_audio": None,
+            "ref_audio_transcribe": None
+        })
+        gr.on(triggers=[msg.change, tts_model.change, ref_audio.change, 
+                            ref_audio_transcribe.change], 
+                fn=lambda text, model, audio, ref_text: {"msg": text, "tts_model": model, "ref_audio": audio, 
+                                                                "ref_audio_transcribe": ref_text},
+                inputs=[msg, tts_model, ref_audio, ref_audio_transcribe],
+                outputs=current_config
+                )
+        submit_btn.click(
+            predict,
+            [current_config],
+            [audio_player],
+            queue=False
+        )
+    demo.launch(share=False, server_name='0.0.0.0', server_port=7863, inbrowser=True)
+    

utils.py

@@ -0,0 +1,197 @@
+import re
+import requests
+import chardet
+import config as cfg 
+from bs4 import BeautifulSoup
+from pathlib import Path
+from transformers import AutoTokenizer
+from duckduckgo_search import DDGS
+
+
+def log_in(uid, state):
+    state['chat_history'] = []
+    state['thinking_history'] = ''
+    state['uid'] = uid
+    if uid!=0:
+        response = f"Your Log In UID: {uid}"
+    else:
+        response = f"You Are Not Logged In Yet, Use Public Directory"
+    user_dir = Path(cfg.USER_DIR) / str(uid)
+    user_dir.mkdir(parents=True, exist_ok=True)
+    state['user_dir'] = user_dir
+    
+    # 加载历史会话
+    state['available_history'] = []
+    for json_file in user_dir.rglob("*.json"):
+        state['available_history'].append(json_file.stem)
+        
+    return response, state
+
+def clean_response(response_content):   
+    response_content = re.sub(r'\*\*|__', '', response_content)
+    response_content = re.sub(r'\\\(|\\\)|\\\[|\\\]', '', response_content)
+    response_content = re.sub(r'\\boxed\{([^}]*)\}', r'\1', response_content)
+    response_content = re.sub(r'\\\\', '', response_content)
+    response_content = re.sub(r'\n\s*\n', '\n\n', response_content)
+    response_content = re.sub(r'\s+', ' ', response_content)
+    return response_content.strip()
+
+def parse_output(response_content):
+    cleaned_content = clean_response(response_content)
+    if "<think>" in cleaned_content and "</think>" in cleaned_content:
+        split_pattern = r'<think>|</think>'
+        parts = re.split(split_pattern, cleaned_content)
+        return parts[1], parts[2]
+    return None, cleaned_content
+
+
+def parse_chat_history(chat_history):
+    """从保存的历史会话中解析出chatbot可以识别的格式
+    chat_history示例：
+        [
+            {
+                "role": "user",
+                "content": "hello"
+            },
+            {
+                "role": "assistant",
+                "content": " Hello! How can I assist you today? 😊"
+            }
+        ]
+    Args:
+        chat_history (list): _description_
+    """
+    from gradio import Warning
+    try:
+        assert len(chat_history) % 2 == 0
+    except AssertionError:
+        Warning('历史会话可能有遗失，用户和AI的消息数不匹配，截断最后一条消息...')
+        chat_history = chat_history[:-1]
+    
+    if len(chat_history) == 0:
+        Warning('历史会话为空或无法匹配，加载历史会话失败...')
+        return []
+
+    messages = []
+    responses = []
+    
+    for conversation in chat_history:
+        if conversation['role'] == 'user':
+            messages.append(conversation['content'])
+        elif conversation['role'] == 'assistant':
+            responses.append(conversation['content'])
+    
+    if len(messages) != len(responses):
+        Warning('用户和AI的消息无法匹配，加载历史会话失败...')
+        return []
+    
+    return list(zip(messages, responses))
+
+
+def web_search(query: str, max_results: int = 3):
+    """获取网络搜索结果并提取关键内容"""
+    try:
+        # 获取搜索结果链接
+        with DDGS() as ddgs:
+            results = [r for r in ddgs.text(query, max_results=max_results)]
+        
+        # 提取网页正文
+        web_contents = []
+        for result in results:
+            try:
+                response = requests.get(result['href'], timeout=5)
+                encoding = chardet.detect(response.content)['encoding']
+                if response.encoding != encoding:
+                    response.encoding = encoding
+                soup = BeautifulSoup(response.text, 'html.parser')
+                # 提取主要文本内容（可根据网站结构调整）
+                main_content = soup.find('main') or soup.find('article') or soup.body
+                web_contents.append({
+                    'title': result['title'],
+                    'content': main_content.get_text(separator=' ', strip=True)[:1000]  # 限制长度
+                })
+            except Exception as e:
+                print('该结果搜索异常...', e)
+                continue
+        return web_contents
+    except Exception as e:
+        print('网络搜索异常，返回空...', e)
+        return []
+
+
+def parse_net_search(search_res):
+    res = []
+    for item in search_res:
+        if len(item['content']) > 0:
+            res.append(f"标题:\n{item['title']}\n内容:\n{item['content']}\n")
+    return res
+
+
+def wash_up_content(doc_score):
+    """_summary_
+
+    Args:
+        doc (string): 整个文档的内容，可能有多行，用换行符分割
+    """
+    if isinstance(doc_score, tuple):
+        doc, score = doc_score
+    else:
+        doc = doc_score
+        score = None
+    res = list(filter(lambda x: len(x) > 0, doc.split('\n')))
+    prefix = '✅<br>"' if score is None else '✅Recall with score:{:.3f}<br>'.format(score)
+    res[0] = prefix + res[0]
+    return res
+
+            
+MODEL_HF_MAPPING = {
+    "qwen2.5:14b-instruct": "Qwen/Qwen2.5-14B-Instruct",
+    "qwen2.5:32b-instruct": "Qwen/Qwen2.5-32B-Instruct",
+    "qwen2.5:7b-instruct": "Qwen/Qwen2.5-7B-Instruct",
+    "qwen2.5:3b-instruct": "Qwen/Qwen2.5-3B-Instruct",
+    "qwen2.5:0.5b-instruct": "Qwen/Qwen2.5-0.5B-Instruct",
+    "qwen2.5:0.5b": "Qwen/Qwen2.5-0.5B",
+    "qwen2.5:32b": "Qwen/Qwen2.5-32B",
+    "qwen3:32b": "Qwen/Qwen3-32B",
+    "qwen3:14b": "Qwen/Qwen3-14B",
+    "qwen3:4b": "Qwen/Qwen3-4B",
+    "qwen3:8b": "Qwen/Qwen3-8B",
+    "qwen3:30b-a3b": "Qwen/Qwen3-30B-A3B",
+    "qwq": "Qwen/QwQ-32B",
+    "deepseek-r1:14b":"deepseek-ai/DeepSeek-R1-Distill-Qwen-14B",
+    "deepseek-r1:7b": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
+    "deepseek-r1:32b": "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
+    
+}
+
+def load_tokenizer(model_name):
+    hf_model_name = MODEL_HF_MAPPING.get(model_name, model_name)
+    return AutoTokenizer.from_pretrained(hf_model_name, use_fast=True)
+
+def messages_to_prompt(messages):
+    # 按模型要求的格式拼接（此处为示例，需根据实际模型调整）
+    prompt = ""
+    for msg in messages:
+        prompt += f"{msg['role']}: {msg['content']}\n"
+    return prompt.strip()
+
+def count_tokens_local(messages, tokenizer):
+    prompt = messages_to_prompt(messages)
+    return len(tokenizer(prompt, return_tensors=None, truncation=False)["input_ids"])
+
+
+def concate_metadata(metadata):
+    """把Document对象的metadata的各个键值拼接起来
+
+    Args:
+        metadata (dict): _description_
+    """
+    return '\n'.join([f"{k}: {v}" for k, v in metadata.items()])
+
+if __name__ == "__main__":
+    query = "今天几号"
+    ret = web_search(query)
+    for item in ret:
+        print(item)
+    ret = parse_net_search(ret)
+    print(ret)