Update README.md

README.md

@@ -34,308 +34,7 @@ By combining autoregression and flow matching, MonoSpeech establishes a foundati
 
 ## 2. Quick Start
 
-Please refer to [**Github Repository**](https://github.com/gwh22/MonoSpeech)
+Please refer to [**Github Repository**](https://github.com/gwh22/Univoice)
 
 
 
-## 3. Usage
-For Zero-shot TTS :
-```py
-import argparse
-import json
-import multiprocessing as mp
-import os
-import socket
-from typing import List, Optional
-from tqdm import tqdm
-import random
-
-import transformers
-import torch
-import torchaudio
-import torch.distributed as dist
-import numpy as np
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
-from transformers import pipeline
-
-from monospeech.monospeech_model import MonoSpeech
-from monospeech.constants import *
-
-from monospeech.utils import MelSpec, make_pad_mask, MelSpec_bigvGAN, MelSpec_Taco
-from monospeech.tensor_util import spec_to_figure, spec_to_figure_single
-
-def setup_seed(seed):
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-    random.seed(seed)
-    torch.backends.cudnn.deterministic = True
-
-@torch.no_grad()
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--ckpt_path", type=str, required=True)
-    parser.add_argument("--llm_path", type=str, required=True)
-    parser.add_argument("--cfg_scale", type=float, required=True)
-    
-    args = parser.parse_args()
-
-    rank = int(os.environ["LOCAL_RANK"])
-    world_size = int(os.environ["WORLD_SIZE"])
-
-    dist.init_process_group("nccl", rank=rank, world_size=world_size)
-    torch.cuda.set_device(rank)
-    setup_seed(42) # random seed default=42
-
-    # load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(args.llm_path, add_bos_token=True, add_eos_token=True)
-    
-    # load model
-    model_config = AutoConfig.from_pretrained(args.llm_path)
-    model_config.learn_sigma = True
-    model_config.tokenizer_max_length = 1024
-    model_config.tokenizer_padding_side = 'right'
-    model_config.use_flash_attn = False
-    # model_config.attn_implementation="flash_attention_2" if model_config.use_flash_attn==True else "eager"
-    model_config.use_pos_embed = True
-    model_config.decoder_t_embed = "add_before_speech_tokens"
-    model_config.use_adaln_final_layer = True
-    model_config.use_bi_attn_img_tokens = True   # or False for causal DiT
-    model_config.add_pos_embed_each_layer = False
-    model_config.use_hybrid_attn_mask = False
-    model_config.audio_encoder_path = 'hf_ckpts/whisper-large-v3'
-    model_config.speaker_encoder_path = 'hf_ckpts/wav2vec2-large-xlsr-53'
-    model = MonoSpeech(
-        model_config,
-        llm_path = args.llm_path,
-        tokenizer = tokenizer,
-        cfg_scale = args.cfg_scale,
-    )
-    ckpt_type = args.ckpt_path.split(".")[-1]
-    if ckpt_type == "safetensors":
-        from safetensors.torch import load_file
-        checkpoint = load_file(args.ckpt_path, device='cuda')
-    else:
-        checkpoint = torch.load(args.ckpt_path, map_location='cuda')
-    model.load_state_dict(checkpoint)
-    model.eval().cuda()
-
-    
-    # wav_path for speaker
-    wav_path = "data/LJ001-0001.wav"
-    
-    audio, source_sample_rate = torchaudio.load(wav_path)
-    if audio.shape[0] > 1: # mono
-        audio = torch.mean(audio, dim=0, keepdim=True)
-    if source_sample_rate != 22050:   # whisper---16KHZ
-        resampler = torchaudio.transforms.Resample(source_sample_rate, 22050)
-        audio = resampler(audio)
-    mel_spectrogram = MelSpec_bigvGAN(
-        n_fft=1024,
-        hop_length=256,
-        win_length=1024,
-        n_mel_channels=80,
-        target_sample_rate=22050,
-    )
-    mel_spec = mel_spectrogram(audio)
-    mel_spec = [mel_spec.squeeze(0).to('cuda')] # (D,T)
-    speechs = [[]]
-    flags = [[0]]
-
-
-    # duration set by yourself
-    duration = 6
-    target_len = [int(duration*22050//256)]  # mel_spec[0].shape[1]. 
-    text = ["At once the goat gave a leap, escaped from the soldiers and with bowed head rushed upon the Boolooroo".lower()]
-
-    temp = torch.randn(1).to('cuda')
-    with torch.inference_mode():
-        mel_out, mel_gt = model.sample(
-            input_ids=temp,
-            attention_mask=temp,
-            labels=temp,
-            mel_spec=mel_spec,
-            speechs=speechs,
-            flags=flags,
-            target_len=target_len,
-            text=text,
-            wav_path=[wav_path],
-        )
-    text_name = '_'.join(text[0].strip().split())
-    os.makedirs('infers', exist_ok=True)
-    # bigvagn vocoder
-    from BigVGAN import bigvgan
-    vocoder = bigvgan.BigVGAN.from_pretrained('hf_ckpts/bigvgan_22k', use_cuda_kernel=False)
-    vocoder.remove_weight_norm()
-    vocoder = vocoder.eval().to('cuda')
-
-    # generate waveform from mel
-    with torch.inference_mode():
-        wav_gen = vocoder(mel_out.transpose(0,1).unsqueeze(0)) # wav_gen is FloatTensor with shape [B(1), 1, T_time] and values in [-1, 1]
-    wav_gen_float = wav_gen.squeeze(0).cpu()
-    # wav_gen_int16 = (wav_gen_float * 32767.0).numpy().astype('int16') # wav_gen is now np.ndarray with shape [1, T_time] and int16 dtype
-
-    torchaudio.save(f'infers/{text_name}.wav', wav_gen_float, 22050)
-
-
-if __name__ == "__main__":
-    main()
-```
-
-For ASR :
-```py
-import argparse
-import json
-import multiprocessing as mp
-import os
-import socket
-from typing import List, Optional
-
-import transformers
-import random
-import numpy as np
-import torch
-import torchaudio
-import torch.distributed as dist
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
-from transformers import pipeline
-
-from monospeech.monospeech_model import MonoSpeech
-from monospeech.constants import *
-
-
-def setup_seed(seed):
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-    random.seed(seed)
-    torch.backends.cudnn.deterministic = True
-
-def preprocess_inputs(tokenizer: transformers.PreTrainedTokenizer, inputs: List[str], speechs: List[torch.Tensor], max_length=512, device='cuda'):
-    """
-    Currently, only support batch size 1.
-    """
-    assert len(inputs) == 1
-
-    input_ids, attention_mask = tokenizer(
-        inputs,
-        max_length=max_length,
-        truncation=True,
-        add_special_tokens=False,
-        return_tensors="pt",
-    ).values()
-    
-    if len(speechs) > 0:
-        im_start_token_id = tokenizer.convert_tokens_to_ids(DEFAULT_SPEECH_START_TOKEN)
-        im_end_token_id = tokenizer.convert_tokens_to_ids(DEFAULT_SPEECH_END_TOKEN)
-        speech_token_id = tokenizer.convert_tokens_to_ids(DEFAULT_SPEECH_TOKEN)
-        for cur_input_ids in input_ids:
-            for idx in torch.where(cur_input_ids == im_start_token_id):
-                if cur_input_ids[idx + 1] == tokenizer.pad_token_id:
-                    cur_input_ids[idx + 1] = speech_token_id
-
-        attention_mask = input_ids.ne(tokenizer.pad_token_id)
-
-        flags = [[1]]
-    else:
-        flags = []
-
-    return {
-        'input_ids': input_ids.to(device),
-        'attention_mask': attention_mask.to(device),
-        'speechs': [speechs],
-        'flags': flags,
-        't': torch.tensor([0]).to(device),
-    }
-
-
-@torch.no_grad()
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--ckpt_path", type=str, required=True)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    parser.add_argument("--top_p", type=float, default=0.9)
-    parser.add_argument("--top_k", type=int, default=50)
-    parser.add_argument("--num_beams", type=int, default=1)
-    parser.add_argument("--llm_path", type=str, required=True)
-    args = parser.parse_args()
-
-    rank = int(os.environ["LOCAL_RANK"])
-    world_size = int(os.environ["WORLD_SIZE"])
-
-    dist.init_process_group("nccl", rank=rank, world_size=world_size)
-    torch.cuda.set_device(rank)
-    setup_seed(42) # random seed default=42
-
-    # load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(args.llm_path, add_bos_token=True, add_eos_token=True)
-    
-    # # load model
-    model_config = AutoConfig.from_pretrained(args.llm_path)
-    model_config.learn_sigma = True
-    model_config.tokenizer_max_length = 1024
-    model_config.tokenizer_padding_side = 'right'
-    model_config.use_flash_attn = False
-    # model_config.attn_implementation="flash_attention_2" if model_config.use_flash_attn==True else "eager"
-    model_config.use_pos_embed = True
-    model_config.decoder_t_embed = "add_before_speech_tokens"
-    model_config.use_adaln_final_layer = True
-    model_config.use_bi_attn_img_tokens = True   # or False for causal DiT
-    model_config.add_pos_embed_each_layer = False
-    model_config.use_hybrid_attn_mask = False
-    model_config.audio_encoder_path = 'hf_ckpts/whisper-large-v3'
-    model_config.speaker_encoder_path = 'hf_ckpts/wav2vec2-large-xlsr-53'
-    model = MonoSpeech(
-        model_config,
-        llm_path = args.llm_path,
-        tokenizer = tokenizer,
-        cfg_scale = 1,
-    )
-    ckpt_type = args.ckpt_path.split(".")[-1]
-    if ckpt_type == "safetensors":
-        from safetensors.torch import load_file
-        checkpoint = load_file(args.ckpt_path, device='cuda')
-    else:
-        checkpoint = torch.load(args.ckpt_path, map_location='cuda')
-    model.load_state_dict(checkpoint)
-    model.eval().cuda()
-
-    feature_extracter = transformers.WhisperFeatureExtractor.from_pretrained('hf_ckpts/whisper-large-v3')
-
-    # asr wav_path
-    wav_path = "data/LJ001-0001.wav"
-    audio, source_sample_rate = torchaudio.load(wav_path)
-    if audio.shape[0] > 1: # mono
-        audio = torch.mean(audio, dim=0, keepdim=True)
-    if source_sample_rate != 16000:   # whisper---16KHZ
-        resampler = torchaudio.transforms.Resample(source_sample_rate, 16000)
-        audio = resampler(audio)
-    
-    mel_spec = feature_extracter(audio.numpy(), sampling_rate=16000).input_features[0]
-    mel_spec = torch.tensor(mel_spec, dtype=torch.float32)
-    # speechs and prompt
-    speechs = [mel_spec.to('cuda')]
-    prompt = f"{DEFAULT_SPEECH_START_TOKEN}{DEFAULT_PAD_TOKEN}{DEFAULT_SPEECH_END_TOKEN}\n"
-    inputs = [f"{tokenizer.bos_token}{prompt}"]
-
-    inputs_dict = preprocess_inputs(tokenizer, inputs, speechs)
-
-    with torch.inference_mode():
-        output_ids = model.generate(
-            input_ids=inputs_dict['input_ids'],
-            attention_mask=inputs_dict['attention_mask'],
-            speechs=inputs_dict['speechs'],
-            flags=inputs_dict['flags'],
-            t=inputs_dict['t'],
-            temperature=args.temperature,
-            top_p=args.top_p,
-            top_k=args.top_k,
-            num_beams=args.num_beams,
-        )
-    output_ids = output_ids.replace("\n"," ").replace("<|im_end|>","")
-    print(output_ids)
-
-
-if __name__ == "__main__":
-    main()
-```