Spaces:
Running on L40S
Running on L40S
File size: 3,738 Bytes
de9e682 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import os
import argparse
import torch
import torchaudio
from tqdm import tqdm
from vllm.v1.engine.processor import Processor
from vllm.engine.llm_engine import LLMEngine
# Bypass checkout
Processor._validate_model_input = lambda *args, **kwargs: None
LLMEngine._validate_token_prompt = lambda *args, **kwargs: None
from vllm import __version__ as vllm_version
from vllm import LLM, SamplingParams
from megatron.tokenizer import build_tokenizer
from mucodec.generate_1rvq import Tango
class Args:
def __init__(self):
pass
class vllmInf:
def __init__(self, model_path, vocal_file, tokenizer="Qwen2Tokenizer", extra_vocab_size=16384):
args = Args()
args.vocab_file = vocal_file
args.load = model_path
args.extra_vocab_size = extra_vocab_size
args.patch_tokenizer_type = tokenizer
self.tokenizer = build_tokenizer(args)
self.text_offset = len(self.tokenizer.tokenizer.get_vocab())
self.max_tokens = 8192
self.llm = LLM(
model=model_path,
trust_remote_code=True,
max_model_len=self.max_tokens,
dtype="bfloat16",
)
def run(self, audios: list[list[int]]):
batch_token_ids = []
max_tokens = self.max_tokens
for audio in audios:
audio = audio + self.text_offset
sentence_ids = [self.tokenizer.sep_token_id] + audio.tolist() + [self.tokenizer.tokenizer.sep_token_id]
max_tokens = min(max_tokens, self.max_tokens - len(sentence_ids))
batch_token_ids.append(sentence_ids)
sampling_params = SamplingParams(
n=1,
max_tokens=max_tokens,
top_p=0.1,
temperature=0.1
)
if vllm_version == "0.8.5":
outputs = self.llm.generate(
prompt_token_ids=batch_token_ids,
sampling_params=sampling_params
)
else:
inputs = [
{"prompt_token_ids": token_ids}
for token_ids in batch_token_ids
]
outputs = self.llm.generate(
prompts=inputs,
sampling_params=sampling_params
)
lyrics = []
for output in outputs:
generate_ids = output.outputs[0].token_ids
lyrics.append(self.tokenizer.detokenize(generate_ids))
return lyrics
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='')
parser.add_argument("-i", dest="input_dir")
parser.add_argument("-q", dest="qwen_ckpt", default="SongPrep-7B/")
parser.add_argument("-c", dest="codec_ckpt", default="SongPrep-7B/mucodec.safetensors")
args = parser.parse_args()
vocal_file = "conf/vocab_type.yaml"
qwen_path = args.qwen_ckpt
codec_path = args.codec_ckpt
input_dir = args.input_dir
# codec
input_audio_paths = []
input_audios = []
tango = Tango(model_path=codec_path)
for audio_path in tqdm(os.listdir(input_dir)):
if not audio_path.endswith(".wav"):
continue
src_wave, fs = torchaudio.load(os.path.join(input_dir, audio_path))
if (fs != 48000):
src_wave = torchaudio.functional.resample(src_wave, fs, 48000)
code = tango.sound2code(src_wave)
input_audios.append(code[0][0].cpu().numpy())
input_audio_paths.append(audio_path)
del tango
torch.cuda.empty_cache()
# batch transcription
vllm_inf = vllmInf(qwen_path, vocal_file)
lyrics = vllm_inf.run(input_audios)
# display
for audio_path, lyric in zip(input_audio_paths, lyrics):
print(f"====={audio_path}=====")
print(lyric)
print("\n")
|