mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-09 01:52:41 +00:00
435 lines
14 KiB
Python
435 lines
14 KiB
Python
# Modified from https://github.com/QwenLM/Qwen2.5-Omni/blob/main/web_demo.py
|
|
import io
|
|
import sys
|
|
from argparse import ArgumentParser
|
|
|
|
import gradio as gr
|
|
import gradio.processing_utils as processing_utils
|
|
import numpy as np
|
|
import sherpa_onnx
|
|
import soundfile as sf
|
|
import torch
|
|
import whisper
|
|
from cosyvoice.cli.cosyvoice import CosyVoice
|
|
from gradio_client import utils as client_utils
|
|
from model import SPEECH_LLM, EncoderProjector
|
|
from peft import LoraConfig, get_peft_model
|
|
from train import DEFAULT_SPEECH_TOKEN, add_model_arguments
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config
|
|
from whisper_encoder_forward_monkey_patch import replace_whisper_encoder_forward
|
|
|
|
# https://github.com/FunAudioLLM/CosyVoice/tree/main/third_party
|
|
sys.path.append("/workspace/CosyVoice/third_party/Matcha-TTS")
|
|
|
|
|
|
def get_model(params, device="cuda"):
|
|
"""Load and prepare the speech-to-speech model."""
|
|
if params.remove_whisper_encoder_input_length_restriction:
|
|
replace_whisper_encoder_forward()
|
|
|
|
whisper_model = whisper.load_model(params.speech_encoder_path_or_name, "cpu")
|
|
speech_encoder = whisper_model.encoder
|
|
speech_encoder_dim = whisper_model.dims.n_audio_state
|
|
tokenizer = AutoTokenizer.from_pretrained(params.llm_path_or_name)
|
|
|
|
if params.use_flash_attn:
|
|
attn_implementation = "flash_attention_2"
|
|
else:
|
|
attn_implementation = "eager"
|
|
|
|
llm = AutoModelForCausalLM.from_pretrained(
|
|
params.llm_path_or_name,
|
|
attn_implementation=attn_implementation,
|
|
torch_dtype=torch.float16,
|
|
)
|
|
if params.use_lora:
|
|
lora_config = LoraConfig(
|
|
r=64,
|
|
lora_alpha=16,
|
|
target_modules=[
|
|
"q_proj",
|
|
"k_proj",
|
|
"v_proj",
|
|
"o_proj",
|
|
"up_proj",
|
|
"gate_proj",
|
|
"down_proj",
|
|
],
|
|
task_type="CAUSAL_LM",
|
|
)
|
|
llm = get_peft_model(llm, lora_config)
|
|
llm.print_trainable_parameters()
|
|
|
|
special_tokens_dict = {"additional_special_tokens": [DEFAULT_SPEECH_TOKEN]}
|
|
tokenizer.add_special_tokens(special_tokens_dict)
|
|
llm.config.pad_token_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
|
|
llm.config.bos_token_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
|
llm.config.eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
|
|
|
|
llm.config.default_speech_token_id = tokenizer.convert_tokens_to_ids(
|
|
DEFAULT_SPEECH_TOKEN
|
|
)
|
|
|
|
encoder_projector = EncoderProjector(
|
|
speech_encoder_dim, llm.config.hidden_size, params.encoder_projector_ds_rate
|
|
)
|
|
|
|
# codec_vocab_size = 4096 + 4
|
|
codec_vocab_size = 6561 + 4
|
|
config = Qwen2Config(
|
|
vocab_size=codec_vocab_size,
|
|
hidden_size=1024,
|
|
num_hidden_layers=12,
|
|
num_attention_heads=16,
|
|
num_key_value_heads=16,
|
|
intermediate_size=2048,
|
|
max_position_embeddings=4096,
|
|
)
|
|
codec_lm = AutoModelForCausalLM.from_config(
|
|
config=config,
|
|
attn_implementation=attn_implementation,
|
|
torch_dtype=torch.float16,
|
|
)
|
|
codec_lm.resize_token_embeddings(codec_vocab_size)
|
|
codec_lm.vocab_size = codec_vocab_size
|
|
codec_lm.config.pad_token_id = codec_vocab_size - 1
|
|
codec_lm.config.eos_token_id = codec_vocab_size - 2
|
|
codec_lm.config.bos_token_id = codec_vocab_size - 3
|
|
codec_lm.config.mask_token_id = codec_vocab_size - 4
|
|
|
|
model = SPEECH_LLM(
|
|
speech_encoder,
|
|
llm,
|
|
encoder_projector,
|
|
codec_lm,
|
|
codec_lm_padding_side="left" if params.use_flash_attn else "right",
|
|
)
|
|
|
|
checkpoint = torch.load(f"{params.checkpoint_path}", map_location="cpu")
|
|
model.load_state_dict(checkpoint, strict=False)
|
|
|
|
model.to(device)
|
|
model.eval()
|
|
return model, tokenizer
|
|
|
|
|
|
def audio_decode_cosyvoice(audio_tokens, codec_decoder):
|
|
"""
|
|
Generate audio from tokens with optional tone and prompt embedding.
|
|
|
|
Args:
|
|
audio_tokens (list): List of audio tokens to be processed.
|
|
codec_decoder: Codec decoder for generating audio.
|
|
|
|
Returns:
|
|
torch.Tensor: Generated audio waveform.
|
|
"""
|
|
flow_embedding = codec_decoder.frontend.spk2info["中文女"]["embedding"]
|
|
flow_prompt_speech_token = torch.zeros(1, 0, dtype=torch.int32)
|
|
prompt_speech_feat = torch.zeros(1, 0, 80)
|
|
tts_mel, _ = codec_decoder.model.flow.inference(
|
|
token=audio_tokens.to(codec_decoder.model.device),
|
|
token_len=torch.tensor([audio_tokens.shape[1]], dtype=torch.int32).to(
|
|
codec_decoder.model.device
|
|
),
|
|
prompt_token=flow_prompt_speech_token.to(codec_decoder.model.device),
|
|
prompt_token_len=torch.tensor(
|
|
[flow_prompt_speech_token.shape[1]], dtype=torch.int32
|
|
).to(codec_decoder.model.device),
|
|
prompt_feat=prompt_speech_feat.to(codec_decoder.model.device),
|
|
prompt_feat_len=torch.tensor(
|
|
[prompt_speech_feat.shape[1]], dtype=torch.int32
|
|
).to(codec_decoder.model.device),
|
|
embedding=flow_embedding.to(codec_decoder.model.device),
|
|
flow_cache=torch.zeros(1, 80, 0, 2).to(codec_decoder.model.device),
|
|
)
|
|
|
|
audio_hat, _ = codec_decoder.model.hift.inference(
|
|
speech_feat=tts_mel, cache_source=torch.zeros(1, 1, 0)
|
|
)
|
|
|
|
return audio_hat
|
|
|
|
|
|
def preprocess(
|
|
messages,
|
|
tokenizer,
|
|
):
|
|
"""Preprocesses the data for supervised fine-tuning."""
|
|
texts = []
|
|
TEMPLATE = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if loop.last %}{{''}}{% else %}{{ '<|im_end|>\n' }}{% endif %}{% endfor %}"
|
|
for i, msg in enumerate(messages):
|
|
texts.append(
|
|
tokenizer.apply_chat_template(
|
|
msg,
|
|
tokenize=True,
|
|
add_generation_prompt=False,
|
|
chat_template=TEMPLATE,
|
|
padding="longest",
|
|
truncation=False,
|
|
)
|
|
)
|
|
max_len_texts = max([len(text) for text in texts])
|
|
if tokenizer.padding_side == "right":
|
|
texts = [
|
|
text + [tokenizer.pad_token_id] * (max_len_texts - len(text))
|
|
for text in texts
|
|
]
|
|
else:
|
|
texts = [
|
|
[tokenizer.pad_token_id] * (max_len_texts - len(text)) + text
|
|
for text in texts
|
|
]
|
|
|
|
input_ids = torch.tensor(texts, dtype=torch.int)
|
|
|
|
attention_mask = input_ids.ne(tokenizer.pad_token_id)
|
|
|
|
return input_ids, attention_mask
|
|
|
|
|
|
def _launch_demo(args, model, tokenizer, token2wav_model, asr_model):
|
|
def format_history(history: list):
|
|
messages = []
|
|
for item in history:
|
|
if isinstance(item["content"], str):
|
|
messages.append({"role": item["role"], "content": item["content"]})
|
|
return messages
|
|
|
|
def decode(
|
|
model,
|
|
token2wav_model,
|
|
tokenizer,
|
|
feature,
|
|
messages,
|
|
):
|
|
"""Decode one
|
|
Returns:
|
|
pass
|
|
"""
|
|
|
|
dtype = torch.float32
|
|
device = model.llm.device
|
|
|
|
feature = feature.to(device, dtype=dtype)
|
|
|
|
input_ids, attention_mask = preprocess([messages], tokenizer)
|
|
|
|
generated_ids, audio_tokens = model.decode_with_speech_output(
|
|
feature, input_ids.to(device, dtype=torch.long), attention_mask.to(device)
|
|
)
|
|
|
|
hyps = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
|
|
|
|
yield {"type": "text", "data": hyps[0]}
|
|
|
|
audio_tokens = [token for token in audio_tokens if token < 4096]
|
|
audio_tokens = torch.tensor(audio_tokens, dtype=torch.int32).unsqueeze(0)
|
|
audio_hat = audio_decode_cosyvoice(audio_tokens, token2wav_model)
|
|
audio = audio_hat.squeeze(0).cpu().numpy()
|
|
audio = np.array(audio * 32767).astype(np.int16)
|
|
wav_io = io.BytesIO()
|
|
sf.write(wav_io, audio, samplerate=22050, format="WAV")
|
|
wav_io.seek(0)
|
|
wav_bytes = wav_io.getvalue()
|
|
audio_path = processing_utils.save_bytes_to_cache(
|
|
wav_bytes, "audio.wav", cache_dir=demo.GRADIO_CACHE
|
|
)
|
|
|
|
yield {"type": "audio", "data": audio_path}
|
|
|
|
def media_predict(audio, history):
|
|
# First yield
|
|
yield (
|
|
None, # microphone
|
|
history, # media_chatbot
|
|
gr.update(visible=False), # submit_btn
|
|
gr.update(visible=True), # stop_btn
|
|
)
|
|
print(2333, history, audio)
|
|
history.append({"role": "user", "content": (audio,)})
|
|
history.append({"role": "user", "content": f"{DEFAULT_SPEECH_TOKEN}"})
|
|
history.append({"role": "assistant", "content": ""})
|
|
formatted_history = format_history(
|
|
history=history
|
|
) # only keep string text format
|
|
|
|
assert audio is not None
|
|
audio_transcript = get_transcript(
|
|
audio,
|
|
asr_model,
|
|
)
|
|
history[-2]["content"] = audio_transcript
|
|
|
|
fbank = whisper.log_mel_spectrogram(audio, device=model.llm.device)
|
|
fbank = fbank.unsqueeze(0)
|
|
assert fbank.ndim == 3
|
|
|
|
for chunk in decode(
|
|
model, token2wav_model, tokenizer, fbank, formatted_history
|
|
):
|
|
if chunk["type"] == "text":
|
|
history[-1]["content"] = chunk["data"]
|
|
yield (
|
|
None, # microphone
|
|
history, # media_chatbot
|
|
gr.update(visible=False), # submit_btn
|
|
gr.update(visible=True), # stop_btn
|
|
)
|
|
if chunk["type"] == "audio":
|
|
history.append(
|
|
{"role": "assistant", "content": gr.Audio(chunk["data"])}
|
|
)
|
|
|
|
# Final yield
|
|
yield (
|
|
None, # microphone
|
|
history, # media_chatbot
|
|
gr.update(visible=True), # submit_btn
|
|
gr.update(visible=False), # stop_btn
|
|
)
|
|
|
|
with gr.Blocks() as demo:
|
|
with gr.Tab("Online"):
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
microphone = gr.Audio(sources=["microphone"], type="filepath")
|
|
submit_btn = gr.Button("Submit", variant="primary")
|
|
stop_btn = gr.Button("Stop", visible=False)
|
|
clear_btn = gr.Button("Clear History")
|
|
with gr.Column(scale=2):
|
|
media_chatbot = gr.Chatbot(height=650, type="messages")
|
|
|
|
def clear_history():
|
|
return [], gr.update(value=None)
|
|
|
|
submit_event = submit_btn.click(
|
|
fn=media_predict,
|
|
inputs=[
|
|
microphone,
|
|
media_chatbot,
|
|
],
|
|
outputs=[microphone, media_chatbot, submit_btn, stop_btn],
|
|
)
|
|
stop_btn.click(
|
|
fn=lambda: (gr.update(visible=True), gr.update(visible=False)),
|
|
inputs=None,
|
|
outputs=[submit_btn, stop_btn],
|
|
cancels=[submit_event],
|
|
queue=False,
|
|
)
|
|
clear_btn.click(
|
|
fn=clear_history, inputs=None, outputs=[media_chatbot, microphone]
|
|
)
|
|
|
|
demo.queue(default_concurrency_limit=100, max_size=100).launch(
|
|
max_threads=100,
|
|
ssr_mode=False,
|
|
share=args.share,
|
|
inbrowser=args.inbrowser,
|
|
server_port=args.server_port,
|
|
server_name=args.server_name,
|
|
)
|
|
|
|
|
|
def _get_args():
|
|
parser = ArgumentParser()
|
|
|
|
parser.add_argument(
|
|
"--checkpoint-path",
|
|
type=str,
|
|
default=None,
|
|
help="Checkpoint name or path, default to %(default)r",
|
|
)
|
|
parser.add_argument(
|
|
"--token2wav-path",
|
|
type=str,
|
|
default=None,
|
|
help="Token2Wav path, default to %(default)r",
|
|
)
|
|
parser.add_argument(
|
|
"--asr-model-dir",
|
|
type=str,
|
|
default=None,
|
|
help="ASR model dir, default to %(default)r",
|
|
)
|
|
parser.add_argument(
|
|
"--flash-attn2",
|
|
action="store_true",
|
|
default=False,
|
|
help="Enable flash_attention_2 when loading the model.",
|
|
)
|
|
parser.add_argument(
|
|
"--share",
|
|
action="store_true",
|
|
default=False,
|
|
help="Create a publicly shareable link for the interface.",
|
|
)
|
|
parser.add_argument(
|
|
"--inbrowser",
|
|
action="store_true",
|
|
default=False,
|
|
help="Automatically launch the interface in a new tab on the default browser.",
|
|
)
|
|
parser.add_argument(
|
|
"--server-port", type=int, default=8001, help="Demo server port."
|
|
)
|
|
parser.add_argument(
|
|
"--server-name", type=str, default="127.0.0.1", help="Demo server name."
|
|
)
|
|
add_model_arguments(parser)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def read_wave(wave_filename: str):
|
|
"""
|
|
Args:
|
|
wave_filename:
|
|
Path to a wave file. It should be single channel and can be of type
|
|
32-bit floating point PCM. Its sample rate does not need to be 24kHz.
|
|
|
|
Returns:
|
|
Return a tuple containing:
|
|
- A 1-D array of dtype np.float32 containing the samples,
|
|
which are normalized to the range [-1, 1].
|
|
- Sample rate of the wave file.
|
|
"""
|
|
|
|
samples, sample_rate = sf.read(wave_filename, dtype="float32")
|
|
assert (
|
|
samples.ndim == 1
|
|
), f"Expected single channel, but got {samples.ndim} channels."
|
|
|
|
samples_float32 = samples.astype(np.float32)
|
|
|
|
return samples_float32, sample_rate
|
|
|
|
|
|
def get_transcript(audio_path, recognizer):
|
|
samples, sample_rate = read_wave(audio_path)
|
|
s = recognizer.create_stream()
|
|
s.accept_waveform(sample_rate, samples)
|
|
recognizer.decode_streams([s])
|
|
return s.result.text
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = _get_args()
|
|
model, tokenizer = get_model(args)
|
|
token2wav = CosyVoice(
|
|
args.token2wav_path, load_jit=False, load_trt=False, fp16=False
|
|
)
|
|
|
|
asr_model = sherpa_onnx.OfflineRecognizer.from_paraformer(
|
|
paraformer=f"{args.asr_model_dir}/model.int8.onnx",
|
|
tokens=f"{args.asr_model_dir}/tokens.txt",
|
|
num_threads=2,
|
|
sample_rate=16000,
|
|
feature_dim=80,
|
|
decoding_method="greedy_search",
|
|
debug=False,
|
|
)
|
|
|
|
_launch_demo(args, model, tokenizer, token2wav, asr_model)
|