x

app.py

@@ -1,7 +1,45 @@
 import gradio as gr
+import sys
+from viitor_voice.inference.transformers_engine import TransformersEngine
+import spaces
 
-def greet(name):
-    return "Hello " + name + "!!"
+if __name__ == '__main__':
+    # Initialize your OfflineInference class with the appropriate paths
+    offline_inference = TransformersEngine("ZzWater/viitor-voice-mix")
+
+
+    @spaces.GPU
+    def clone_batch(text_list, prompt_audio, prompt_text):
+        print(prompt_audio.name)
+        try:
+            audios = offline_inference.batch_infer(
+                text_list=[text_list],
+                prompt_audio_path=prompt_audio.name,  # Use uploaded file's path
+                prompt_text=prompt_text,
+            )
+            return 24000, audios[0].cpu().numpy()[0].astype('float32')
+        except Exception as e:
+            return str(e)
+
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# TTS Inference Interface")
+        with gr.Tab("Batch Clone"):
+            gr.Markdown("### Batch Clone TTS")
+
+            text_list_clone = gr.Textbox(label="Input Text List (Comma-Separated)",
+                                         placeholder="Enter text1, text2, text3...")
+            prompt_audio = gr.File(label="Upload Prompt Audio")
+            prompt_text = gr.Textbox(label="Prompt Text", placeholder="Enter the prompt text")
+
+            clone_button = gr.Button("Run Batch Clone")
+            clone_output = gr.Audio(label="Generated Audios", type="numpy")
+
+            clone_button.click(
+                fn=clone_batch,
+                inputs=[text_list_clone, prompt_audio, prompt_text],
+                outputs=clone_output
+            )
+
+    demo.launch()
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()

requirements.txt

@@ -0,0 +1,9 @@
+requests
+accelerate==1.1.1
+datasets==3.1.0
+transformers
+tokenizers
+snac
+torch==2.4.0
+torchaudio==2.4.0
+soundfile

viitor_voice/inference/common.py

@@ -0,0 +1,90 @@
+import os
+import re
+from io import BytesIO
+from urllib.parse import urlparse
+import requests
+import torchaudio
+
+
+def load_audio(source):
+    def is_url(path):
+        try:
+            result = urlparse(path)
+            return all([result.scheme, result.netloc])
+        except Exception:
+            return False
+
+    if is_url(source):
+        # 从 URL 加载音频
+        response = requests.get(source)
+        response.raise_for_status()  # 检查请求状态
+        audio_data = BytesIO(response.content)  # 转为类文件对象
+    else:
+        # 从本地文件加载音频
+        if not os.path.exists(source):
+            raise FileNotFoundError(f"File not found: {source}")
+        audio_data = source  # 本地路径可以直接传递给 torchaudio.load
+
+    # 使用 torchaudio 加载音频
+    waveform, sample_rate = torchaudio.load(audio_data)
+    return waveform, sample_rate
+
+
+pattern = re.compile(r"<\|speech-(\d+)\|>")
+
+
+def combine_sequences(first_elements, second_elements, third_elements):
+    group_size = 7
+    sequence = []
+
+    second_index = 0
+    third_index = 0
+
+    for first in first_elements:
+        group = [None] * group_size
+
+        # Assign the first element
+        group[0] = first
+
+        # Assign the second and fifth elements if they exist
+        if second_index < len(second_elements):
+            group[1] = second_elements[second_index]
+            second_index += 1
+        if second_index < len(second_elements):
+            group[4] = second_elements[second_index]
+            second_index += 1
+
+        # Assign the remaining elements from third_elements if they exist
+        for j in [2, 3, 5, 6]:
+            if third_index < len(third_elements):
+                group[j] = third_elements[third_index]
+                third_index += 1
+
+        # Remove None values at the end of the group if the group is incomplete
+        sequence.extend([x for x in group if x is not None])
+
+    return sequence
+
+
+def split_sequence(sequence):
+    group_size = 7
+    first_elements = []
+    second_elements = []
+    third_elements = []
+
+    # Iterate over the sequence in chunks of 7
+    for i in range(0, len(sequence), group_size):
+        group = sequence[i:i + group_size]
+
+        # Add elements to the respective lists based on their position in the group
+        if len(group) >= 1:
+            first_elements.append(group[0])
+        if len(group) >= 5:
+            second_elements.extend([group[1], group[4]])
+        if len(group) >= 7:
+            third_elements.extend([group[2], group[3], group[5], group[6]])
+        else:
+            third_elements.extend(group[2:])
+
+    return first_elements, second_elements, third_elements
+

viitor_voice/inference/transformers_engine.py

@@ -0,0 +1,64 @@
+import numpy as np
+import torch
+import torchaudio
+from snac import SNAC
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from viitor_voice.inference.common import combine_sequences, load_audio, pattern, split_sequence
+
+
+class TransformersEngine:
+    def __init__(self, model_path, device='cuda'):
+        self.device = device
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).to(device)
+        self.snac_model = SNAC.from_pretrained('hubertsiuzdak/snac_24khz').eval().to(device)
+
+    def batch_infer(self, text_list, prompt_audio_path, prompt_text, flattened_snac_encode=None):
+        if flattened_snac_encode is None:
+            prompt_audio, sr = load_audio(prompt_audio_path)
+            if sr != 24000:
+                prompt_audio = torchaudio.functional.resample(prompt_audio, sr, 24000)
+
+            snac_encode = self.snac_model.encode(prompt_audio[None,].to(self.device))
+            first_elements, second_elements, third_elements = \
+                snac_encode[0].cpu().numpy().tolist(), snac_encode[1].cpu().numpy().tolist(), snac_encode[
+                    2].cpu().numpy().tolist()
+            flattened_snac_encode = combine_sequences(first_elements[0], second_elements[0], third_elements[0])
+        prompt_snac_texts = ''.join(
+            ['<|speech-{}|>'.format(i) if j % 7 != 0 else '<|SEP_AUDIO|><|speech-{}|>'.format(i) for
+             j, i in
+             enumerate(flattened_snac_encode)])
+
+        prompts = [
+            '<|START_TEXT|>' + prompt_text + x + '<|END_TEXT|>' + '<|START_AUDIO|>' + prompt_snac_texts + '<|SEP_AUDIO|>'
+            for x in text_list]
+        prompt_ids_list = self.tokenizer(prompts, add_special_tokens=False).input_ids
+        results = []
+        for prompt_ids in prompt_ids_list:
+            prompt_ids = torch.tensor([prompt_ids], dtype=torch.int64).to(self.device)
+            output_ids = self.model.generate(prompt_ids, eos_token_id=156008, no_repeat_ngram_size=0, num_beams=1,
+                                             do_sample=False, repetition_penalty=1.3,
+                                             suppress_tokens=list(range(151641)))
+            output_ids = output_ids[0, prompt_ids.shape[-1]:].cpu().numpy().tolist()
+            generated_text = self.tokenizer.batch_decode([output_ids], skip_special_tokens=False)
+            snac_tokens = pattern.findall(generated_text)
+            snac_tokens = [int(x) for x in snac_tokens]
+            results.append(snac_tokens)
+        audios = self.batch_decode_audios(results)
+        return audios
+
+    def batch_decode_audios(self, snac_tokens_list):
+        audios = []
+        with torch.no_grad():
+            for snac_tokens in snac_tokens_list:
+                try:
+                    first_elements, second_elements, third_elements = split_sequence(snac_tokens)
+                    codes = [torch.from_numpy(np.array(x).astype(np.int32)[None,]).to(self.device) for x in
+                             [first_elements, second_elements, third_elements]]
+                    audio_hat_all = self.snac_model.decode(codes)[0].cpu()
+                    audios.append(audio_hat_all.to(torch.float32))
+                except:
+                    audios.append('error')
+                    print('error')
+        return audios
+