Upload 5 files

flow_inference.py

@@ -0,0 +1,142 @@
+import torch
+import torchaudio
+import numpy as np
+import re
+from hyperpyyaml import load_hyperpyyaml
+import uuid
+from collections import defaultdict
+
+
+def fade_in_out(fade_in_mel, fade_out_mel, window):
+    device = fade_in_mel.device
+    fade_in_mel, fade_out_mel = fade_in_mel.cpu(), fade_out_mel.cpu()
+    mel_overlap_len = int(window.shape[0] / 2)
+    fade_in_mel[..., :mel_overlap_len] = fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len] + \
+                                         fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:]
+    return fade_in_mel.to(device)
+
+
+class AudioDecoder:
+    def __init__(self, config_path, flow_ckpt_path, hift_ckpt_path, device="cuda"):
+        self.device = device
+
+        with open(config_path, 'r') as f:
+            self.scratch_configs = load_hyperpyyaml(f)
+
+        # Load models
+        self.flow = self.scratch_configs['flow']
+        self.flow.load_state_dict(torch.load(flow_ckpt_path, map_location=self.device))
+        self.hift = self.scratch_configs['hift']
+        self.hift.load_state_dict(torch.load(hift_ckpt_path, map_location=self.device))
+
+        # Move models to the appropriate device
+        self.flow.to(self.device)
+        self.hift.to(self.device)
+        self.mel_overlap_dict = defaultdict(lambda: None)
+        self.hift_cache_dict = defaultdict(lambda: None)
+        self.token_min_hop_len = 2 * self.flow.input_frame_rate
+        self.token_max_hop_len = 4 * self.flow.input_frame_rate
+        self.token_overlap_len = 5
+        self.mel_overlap_len = int(self.token_overlap_len / self.flow.input_frame_rate * 22050 / 256)
+        self.mel_window = np.hamming(2 * self.mel_overlap_len)
+        # hift cache
+        self.mel_cache_len = 1
+        self.source_cache_len = int(self.mel_cache_len * 256)
+        # speech fade in out
+        self.speech_window = np.hamming(2 * self.source_cache_len)
+
+    def token2wav(self, token, uuid, prompt_token=torch.zeros(1, 0, dtype=torch.int32),
+                  prompt_feat=torch.zeros(1, 0, 80), embedding=torch.zeros(1, 192), finalize=False):
+        tts_mel = self.flow.inference(token=token.to(self.device),
+                                      token_len=torch.tensor([token.shape[1]], dtype=torch.int32).to(self.device),
+                                      prompt_token=prompt_token.to(self.device),
+                                      prompt_token_len=torch.tensor([prompt_token.shape[1]], dtype=torch.int32).to(
+                                          self.device),
+                                      prompt_feat=prompt_feat.to(self.device),
+                                      prompt_feat_len=torch.tensor([prompt_feat.shape[1]], dtype=torch.int32).to(
+                                          self.device),
+                                      embedding=embedding.to(self.device))
+
+        # mel overlap fade in out
+        if self.mel_overlap_dict[uuid] is not None:
+            tts_mel = fade_in_out(tts_mel, self.mel_overlap_dict[uuid], self.mel_window)
+        # append hift cache
+        if self.hift_cache_dict[uuid] is not None:
+            hift_cache_mel, hift_cache_source = self.hift_cache_dict[uuid]['mel'], self.hift_cache_dict[uuid]['source']
+            tts_mel = torch.concat([hift_cache_mel, tts_mel], dim=2)
+
+        else:
+            hift_cache_source = torch.zeros(1, 1, 0)
+        # _tts_mel=tts_mel.contiguous()
+        # keep overlap mel and hift cache
+        if finalize is False:
+            self.mel_overlap_dict[uuid] = tts_mel[:, :, -self.mel_overlap_len:]
+            tts_mel = tts_mel[:, :, :-self.mel_overlap_len]
+            tts_speech, tts_source = self.hift.inference(mel=tts_mel, cache_source=hift_cache_source)
+
+            self.hift_cache_dict[uuid] = {'mel': tts_mel[:, :, -self.mel_cache_len:],
+                                          'source': tts_source[:, :, -self.source_cache_len:],
+                                          'speech': tts_speech[:, -self.source_cache_len:]}
+            # if self.hift_cache_dict[uuid] is not None:
+            #     tts_speech = fade_in_out(tts_speech, self.hift_cache_dict[uuid]['speech'], self.speech_window)
+            tts_speech = tts_speech[:, :-self.source_cache_len]
+
+        else:
+            tts_speech, tts_source = self.hift.inference(mel=tts_mel, cache_source=hift_cache_source)
+            del self.hift_cache_dict[uuid]
+            del self.mel_overlap_dict[uuid]
+            # if uuid in self.hift_cache_dict.keys() and self.hift_cache_dict[uuid] is not None:
+            #     tts_speech = fade_in_out(tts_speech, self.hift_cache_dict[uuid]['speech'], self.speech_window)
+        return tts_speech, tts_mel
+
+    def offline_inference(self, token):
+        this_uuid = str(uuid.uuid1())
+        tts_speech, tts_mel = self.token2wav(token, uuid=this_uuid, finalize=True)
+        return tts_speech.cpu()
+
+    def stream_inference(self, token):
+        token.to(self.device)
+        this_uuid = str(uuid.uuid1())
+
+        # Prepare other necessary input tensors
+        llm_embedding = torch.zeros(1, 192).to(self.device)
+        prompt_speech_feat = torch.zeros(1, 0, 80).to(self.device)
+        flow_prompt_speech_token = torch.zeros(1, 0, dtype=torch.int32).to(self.device)
+
+        tts_speechs = []
+        tts_mels = []
+
+        block_size = self.flow.encoder.block_size
+        prev_mel = None
+
+        for idx in range(0, token.size(1), block_size):
+            # if idx>block_size: break
+            tts_token = token[:, idx:idx + block_size]
+
+            print(tts_token.size())
+
+            if prev_mel is not None:
+                prompt_speech_feat = torch.cat(tts_mels, dim=-1).transpose(1, 2)
+                flow_prompt_speech_token = token[:, :idx]
+
+            if idx + block_size >= token.size(-1):
+                is_finalize = True
+            else:
+                is_finalize = False
+
+            tts_speech, tts_mel = self.token2wav(tts_token, uuid=this_uuid,
+                                                 prompt_token=flow_prompt_speech_token.to(self.device),
+                                                 prompt_feat=prompt_speech_feat.to(self.device), finalize=is_finalize)
+
+            prev_mel = tts_mel
+            prev_speech = tts_speech
+            print(tts_mel.size())
+
+            tts_speechs.append(tts_speech)
+            tts_mels.append(tts_mel)
+
+        # Convert Mel spectrogram to audio using HiFi-GAN
+        tts_speech = torch.cat(tts_speechs, dim=-1).cpu()
+
+        return tts_speech.cpu()
+

model_server.py

@@ -0,0 +1,116 @@
+"""
+A model worker executes the model.
+"""
+import argparse
+import json
+import uuid
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+from transformers import AutoModel, AutoTokenizer
+import torch
+import uvicorn
+
+from transformers.generation.streamers import BaseStreamer
+from threading import Thread
+from queue import Queue
+
+
+class TokenStreamer(BaseStreamer):
+    def __init__(self, skip_prompt: bool = False, timeout=None):
+        self.skip_prompt = skip_prompt
+
+        # variables used in the streaming process
+        self.token_queue = Queue()
+        self.stop_signal = None
+        self.next_tokens_are_prompt = True
+        self.timeout = timeout
+
+    def put(self, value):
+        if len(value.shape) > 1 and value.shape[0] > 1:
+            raise ValueError("TextStreamer only supports batch size 1")
+        elif len(value.shape) > 1:
+            value = value[0]
+
+        if self.skip_prompt and self.next_tokens_are_prompt:
+            self.next_tokens_are_prompt = False
+            return
+
+        for token in value.tolist():
+            self.token_queue.put(token)
+
+    def end(self):
+        self.token_queue.put(self.stop_signal)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value = self.token_queue.get(timeout=self.timeout)
+        if value == self.stop_signal:
+            raise StopIteration()
+        else:
+            return value
+
+
+class ModelWorker:
+    def __init__(self, model_path, device='cuda'):
+        self.device = device
+        self.glm_model = AutoModel.from_pretrained(model_path, trust_remote_code=True,
+                                                   device_map=device,low_cpu_mem_usage=True,load_in_8bit=True).eval()
+        self.glm_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+    @torch.inference_mode()
+    def generate_stream(self, params):
+        tokenizer, model = self.glm_tokenizer, self.glm_model
+
+        prompt = params["prompt"]
+
+        temperature = float(params.get("temperature", 1.0))
+        top_p = float(params.get("top_p", 1.0))
+        max_new_tokens = int(params.get("max_new_tokens", 256))
+
+        inputs = tokenizer([prompt], return_tensors="pt")
+        inputs = inputs.to(self.device)
+        streamer = TokenStreamer(skip_prompt=True)
+        thread = Thread(target=model.generate,
+                        kwargs=dict(**inputs, max_new_tokens=int(max_new_tokens),
+                                    temperature=float(temperature), top_p=float(top_p),
+                                    streamer=streamer))
+        thread.start()
+        for token_id in streamer:
+            yield (json.dumps({"token_id": token_id, "error_code": 0}) + "\n").encode()
+
+    def generate_stream_gate(self, params):
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except Exception as e:
+            print("Caught Unknown Error", e)
+            ret = {
+                "text": "Server Error",
+                "error_code": 1,
+            }
+            yield (json.dumps(ret)+ "\n").encode()
+
+
+app = FastAPI()
+
+
+@app.post("/generate_stream")
+async def generate_stream(request: Request):
+    params = await request.json()
+
+    generator = worker.generate_stream_gate(params)
+    return StreamingResponse(generator)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=10000)
+    parser.add_argument("--model-path", type=str, default="glm-4-voice-9b-int8")
+    args = parser.parse_args()
+
+    worker = ModelWorker(args.model_path)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

quantifization.py

@@ -0,0 +1,27 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+device = "cuda:0"
+
+tokenizer = AutoTokenizer.from_pretrained("glm-4-voice-9b", trust_remote_code=True)
+
+tokenizer.chat_template = "{{role}}: {{content}}"
+
+query = "你好"
+
+inputs = tokenizer.apply_chat_template([{"role": "user", "content": query}],
+add_generation_prompt=True,
+tokenize=True,
+return_tensors="pt",
+return_dict=True
+)
+
+inputs = inputs.to(device)
+model = AutoModelForCausalLM.from_pretrained(
+"glm-4-voice-9b",
+low_cpu_mem_usage=True,
+trust_remote_code=True,
+load_in_8bit=True
+).eval()
+model.save_pretrained("glm-4-voice-9b-int8")
+tokenizer.save_pretrained("glm-4-voice-9b-int8")

requirements.txt

@@ -0,0 +1,36 @@
+conformer==0.3.2
+deepspeed==0.14.2; sys_platform == 'linux'
+diffusers==0.27.2
+fastapi==0.115.3
+fastapi-cli==0.0.4
+gdown==5.1.0
+gradio==5.3.0
+grpcio==1.57.0
+grpcio-tools==1.57.0
+huggingface_hub==0.25.2
+hydra-core==1.3.2
+HyperPyYAML==1.2.2
+inflect==7.3.1
+librosa==0.10.2
+lightning==2.2.4
+matplotlib==3.7.5
+modelscope==1.15.0
+networkx==3.1
+numpy==1.24.4
+omegaconf==2.3.0
+onnxruntime-gpu==1.16.0; sys_platform == 'linux'
+onnxruntime==1.16.0; sys_platform == 'darwin' or sys_platform == 'windows'
+openai-whisper==20231117
+protobuf==4.25
+pydantic==2.7.0
+rich==13.7.1
+Requests==2.32.3
+safetensors==0.4.5
+soundfile==0.12.1
+tensorboard==2.14.0
+transformers==4.44.1
+uvicorn==0.32.0
+wget==3.2
+WeTextProcessing==1.0.3
+torch==2.3.0
+torchaudio==2.3.0

web_demo.py

@@ -0,0 +1,264 @@
+import json
+import os.path
+import tempfile
+import sys
+import re
+import uuid
+import requests
+from argparse import ArgumentParser
+
+import torchaudio
+from transformers import WhisperFeatureExtractor, AutoTokenizer, AutoModel
+from speech_tokenizer.modeling_whisper import WhisperVQEncoder
+
+# import gc
+
+
+sys.path.insert(0, "./cosyvoice")
+sys.path.insert(0, "./third_party/Matcha-TTS")
+
+from speech_tokenizer.utils import extract_speech_token
+
+import gradio as gr
+import torch
+
+audio_token_pattern = re.compile(r"<\|audio_(\d+)\|>")
+
+from flow_inference import AudioDecoder
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default="8888")
+    parser.add_argument("--flow-path", type=str, default="./glm-4-voice-decoder")
+    parser.add_argument("--model-path", type=str, default="./glm-4-voice-9b-int8")
+    parser.add_argument("--tokenizer-path", type=str, default="./glm-4-voice-tokenizer")
+    args = parser.parse_args()
+
+    flow_config = os.path.join(args.flow_path, "config.yaml")
+    flow_checkpoint = os.path.join(args.flow_path, 'flow.pt')
+    hift_checkpoint = os.path.join(args.flow_path, 'hift.pt')
+    glm_tokenizer = None
+    device = "cuda"
+    audio_decoder: AudioDecoder = None
+    whisper_model, feature_extractor = None, None
+
+
+    def initialize_fn():
+        global audio_decoder, feature_extractor, whisper_model, glm_model, glm_tokenizer
+        if audio_decoder is not None:
+            return
+
+        # GLM
+        glm_tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
+
+        # Flow & Hift
+        audio_decoder = AudioDecoder(config_path=flow_config, flow_ckpt_path=flow_checkpoint,
+                                     hift_ckpt_path=hift_checkpoint,
+                                     device=device)
+
+        # Speech tokenizer
+        whisper_model = WhisperVQEncoder.from_pretrained(args.tokenizer_path).eval().to(device)
+        feature_extractor = WhisperFeatureExtractor.from_pretrained(args.tokenizer_path)
+
+
+    def clear_fn():
+        return [], [], '', '', '', None, None
+
+
+    def inference_fn(
+            temperature: float,
+            top_p: float,
+            max_new_token: int,
+            input_mode,
+            audio_path: str | None,
+            input_text: str | None,
+            history: list[dict],
+            previous_input_tokens: str,
+            previous_completion_tokens: str,
+    ):
+
+        # gc.collect()
+        # torch.cuda.empty_cache()
+        # torch.cuda.ipc_collect()
+
+        if input_mode == "audio":
+            assert audio_path is not None
+            history.append({"role": "user", "content": {"path": audio_path}})
+            audio_tokens = extract_speech_token(
+                whisper_model, feature_extractor, [audio_path]
+            )[0]
+            if len(audio_tokens) == 0:
+                raise gr.Error("No audio tokens extracted")
+            audio_tokens = "".join([f"<|audio_{x}|>" for x in audio_tokens])
+            audio_tokens = "<|begin_of_audio|>" + audio_tokens + "<|end_of_audio|>"
+            user_input = audio_tokens
+            system_prompt = "User will provide you with a speech instruction. Do it step by step. First, think about the instruction and respond in a interleaved manner, with 13 text token followed by 26 audio tokens. "
+
+        else:
+            assert input_text is not None
+            history.append({"role": "user", "content": input_text})
+            user_input = input_text
+            system_prompt = "User will provide you with a text instruction. Do it step by step. First, think about the instruction and respond in a interleaved manner, with 13 text token followed by 26 audio tokens."
+
+
+        # Gather history
+        inputs = previous_input_tokens + previous_completion_tokens
+        inputs = inputs.strip()
+        if "<|system|>" not in inputs:
+            inputs += f"<|system|>\n{system_prompt}"
+        inputs += f"<|user|>\n{user_input}<|assistant|>streaming_transcription\n"
+
+        with torch.no_grad():
+            response = requests.post(
+                "http://localhost:10000/generate_stream",
+                data=json.dumps({
+                    "prompt": inputs,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "max_new_tokens": max_new_token,
+                }),
+                stream=True
+            )
+            text_tokens, audio_tokens = [], []
+            audio_offset = glm_tokenizer.convert_tokens_to_ids('<|audio_0|>')
+            end_token_id = glm_tokenizer.convert_tokens_to_ids('<|user|>')
+            complete_tokens = []
+            prompt_speech_feat = torch.zeros(1, 0, 80).to(device)
+            flow_prompt_speech_token = torch.zeros(1, 0, dtype=torch.int64).to(device)
+            this_uuid = str(uuid.uuid4())
+            tts_speechs = []
+            tts_mels = []
+            prev_mel = None
+            is_finalize = False
+            block_size = 10
+            for chunk in response.iter_lines():
+                token_id = json.loads(chunk)["token_id"]
+                if token_id == end_token_id:
+                    is_finalize = True
+                if len(audio_tokens) >= block_size or (is_finalize and audio_tokens):
+                    block_size = 20
+                    tts_token = torch.tensor(audio_tokens, device=device).unsqueeze(0)
+
+                    if prev_mel is not None:
+                        prompt_speech_feat = torch.cat(tts_mels, dim=-1).transpose(1, 2)
+
+                    tts_speech, tts_mel = audio_decoder.token2wav(tts_token, uuid=this_uuid,
+                                                                  prompt_token=flow_prompt_speech_token.to(device),
+                                                                  prompt_feat=prompt_speech_feat.to(device),
+                                                                  finalize=is_finalize)
+                    prev_mel = tts_mel
+
+                    tts_speechs.append(tts_speech.squeeze())
+                    tts_mels.append(tts_mel)
+                    yield history, inputs, '', '', (22050, tts_speech.squeeze().cpu().numpy()), None
+                    flow_prompt_speech_token = torch.cat((flow_prompt_speech_token, tts_token), dim=-1)
+                    audio_tokens = []
+                if not is_finalize:
+                    complete_tokens.append(token_id)
+                    if token_id >= audio_offset:
+                        audio_tokens.append(token_id - audio_offset)
+                    else:
+                        text_tokens.append(token_id)
+        tts_speech = torch.cat(tts_speechs, dim=-1).cpu()
+        complete_text = glm_tokenizer.decode(complete_tokens, spaces_between_special_tokens=False)
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            torchaudio.save(f, tts_speech.unsqueeze(0), 22050, format="wav")
+        history.append({"role": "assistant", "content": {"path": f.name, "type": "audio/wav"}})
+        history.append({"role": "assistant", "content": glm_tokenizer.decode(text_tokens, ignore_special_tokens=False)})
+        yield history, inputs, complete_text, '', None, (22050, tts_speech.numpy())
+
+
+    def update_input_interface(input_mode):
+        if input_mode == "audio":
+            return [gr.update(visible=True), gr.update(visible=False)]
+        else:
+            return [gr.update(visible=False), gr.update(visible=True)]
+
+
+    # Create the Gradio interface
+    with gr.Blocks(title="GLM-4-Voice Demo", fill_height=True) as demo:
+        with gr.Row():
+            temperature = gr.Number(
+                label="Temperature",
+                value=0.2
+            )
+
+            top_p = gr.Number(
+                label="Top p",
+                value=0.8
+            )
+
+            max_new_token = gr.Number(
+                label="Max new tokens",
+                value=2000,
+            )
+
+        chatbot = gr.Chatbot(
+            elem_id="chatbot",
+            bubble_full_width=False,
+            type="messages",
+            scale=1,
+        )
+
+        with gr.Row():
+            with gr.Column():
+                input_mode = gr.Radio(["audio", "text"], label="Input Mode", value="audio")
+                audio = gr.Audio(label="Input audio", type='filepath', show_download_button=True, visible=True)
+                text_input = gr.Textbox(label="Input text", placeholder="Enter your text here...", lines=2, visible=False)
+
+            with gr.Column():
+                submit_btn = gr.Button("Submit")
+                reset_btn = gr.Button("Clear")
+                output_audio = gr.Audio(label="Play", streaming=True,
+                                        autoplay=True, show_download_button=False)
+                complete_audio = gr.Audio(label="Last Output Audio (If Any)", show_download_button=True)
+
+
+
+        gr.Markdown("""## Debug Info""")
+        with gr.Row():
+            input_tokens = gr.Textbox(
+                label=f"Input Tokens",
+                interactive=False,
+            )
+
+            completion_tokens = gr.Textbox(
+                label=f"Completion Tokens",
+                interactive=False,
+            )
+
+        detailed_error = gr.Textbox(
+            label=f"Detailed Error",
+            interactive=False,
+        )
+
+        history_state = gr.State([])
+
+        respond = submit_btn.click(
+            inference_fn,
+            inputs=[
+                temperature,
+                top_p,
+                max_new_token,
+                input_mode,
+                audio,
+                text_input,
+                history_state,
+                input_tokens,
+                completion_tokens,
+            ],
+            outputs=[history_state, input_tokens, completion_tokens, detailed_error, output_audio, complete_audio]
+        )
+
+        respond.then(lambda s: s, [history_state], chatbot)
+
+        reset_btn.click(clear_fn, outputs=[chatbot, history_state, input_tokens, completion_tokens, detailed_error, output_audio, complete_audio])
+        input_mode.input(clear_fn, outputs=[chatbot, history_state, input_tokens, completion_tokens, detailed_error, output_audio, complete_audio]).then(update_input_interface, inputs=[input_mode], outputs=[audio, text_input])
+
+    initialize_fn()
+    # Launch the interface
+    demo.launch(
+        server_port=args.port,
+        server_name=args.host
+    )