import os
import torch
from litgpt.generate.base import next_token_image_batch
import soundfile as sf
from utils.snac_utils import layershift, reconscruct_snac, reconstruct_tensors, get_time_str
from utils.snac_utils import get_snac, generate_audio_data
import clip
import inference
from tqdm import tqdm
from inference import OmniInference, load_model, load_audio, download_model
from inference import text_vocabsize, padded_text_vocabsize, get_text_stream
from PIL import Image
torch.set_printoptions(sci_mode=False)
_image = inference._image
_eoimage = inference._eoimage
_pad_t = inference._pad_t
_input_t = inference._input_t
_answer_t = inference._answer_t
_eot = inference._eot
_eoa = inference._eoa
_pad_a = inference._pad_a
_input_a = inference._input_a
_answer_a = inference._answer_a
def get_input_ids_ImageQA_ATBatch(mel, leng, whispermodel, device):
with torch.no_grad():
mel = mel.unsqueeze(0).to(device)
audio_feature = whispermodel.embed_audio(mel)[0][:leng]
audio_len = audio_feature.size(0)
input_ids = []
input_ids_item = [[] for i in range(8)]
for i in range(7):
input_ids_item[i] = [layershift(_image,i)] + [layershift(_pad_a,i)] * 50 + [layershift(_eoimage,i)]
input_ids_item[i] += [layershift(_input_a,i)]+[layershift(_pad_a,i)]*(audio_len)+[layershift(_eoa,i)]
input_ids_item[i] += [layershift(_answer_a,i)]
input_ids_item[-1] = [_pad_t]* (52 + 2 + audio_len) + [_answer_t]
input_ids_item = [torch.tensor(item) for item in input_ids_item]
input_ids.append(input_ids_item)
input_ids_item = [[] for i in range(8)]
for i in range(7):
input_ids_item[i] = [layershift(_image,i)] + [layershift(_pad_a,i)] * 50 + [layershift(_eoimage,i)]
input_ids_item[i] += [layershift(_input_a,i)]+[layershift(_pad_a,i)]*(audio_len)+[layershift(_eoa,i)] + [layershift(_pad_a,i)]
input_ids_item[-1] = [_pad_t]* (52 + 2 + audio_len) + [_answer_t]
input_ids_item = [torch.tensor(item) for item in input_ids_item]
input_ids.append(input_ids_item)
stacked_inputids = [[] for _ in range(8)]
for i in range(2):
for j in range(8):
stacked_inputids[j].append(input_ids[i][j])
stacked_inputids = [torch.stack(tensors) for tensors in stacked_inputids]
return torch.stack([audio_feature,audio_feature]), stacked_inputids
def load_clip_model(ckpt_dir, device):
clip_model_path = ckpt_dir + "/ViT-B-32.pt"
if not os.path.exists(clip_model_path):
clip_model_path = "ViT-B/32"
clipmodel, clippreprocess = clip.load(clip_model_path, device=device)
return clipmodel, clippreprocess
class OmniVisionInference(OmniInference):
def __init__(self, ckpt_dir='./checkpoint', device='cuda:0'):
self.device = device
if not os.path.exists(ckpt_dir):
print(f"checkpoint directory {ckpt_dir} not found, downloading from huggingface")
download_model(ckpt_dir)
self.fabric, self.model, self.text_tokenizer, self.snacmodel, self.whispermodel = load_model(ckpt_dir, device)
self.clipmodel, self.clippreprocess = load_clip_model(ckpt_dir, device)
def warm_up(self,
audio_sample='./data/samples/vision_qa_audio.wav',
image_sample='./data/samples/vision_qa_image.jpg'
):
for _ in self.run_vision_AA_batch_stream(audio_sample, image_sample,
save_path="./data/samples/vision_qa_output.wav",
warm_up=True):
pass
@torch.inference_mode()
def run_vision_AA_batch_stream(self, audio_path, image_path,
stream_stride=4,
max_returned_tokens=2048,
temperature=0.9,
top_k=1,
top_p=1.0,
eos_id_a=_eoa,
eos_id_t=_eot,
pad_id=_pad_t,
save_path=None,
warm_up=False
):
with self.fabric.init_tensor():
self.model.set_kv_cache(batch_size=2)
model = self.model
mel, leng = load_audio(audio_path)
img = Image.open(image_path)
audio_feature, input_ids = get_input_ids_ImageQA_ATBatch(mel, leng, self.whispermodel, self.device)
ima = self.clippreprocess(img).unsqueeze(0).to(self.device)
ima_feature = self.clipmodel.encode_image(ima).squeeze(0).to(self.device)
ima_feature = torch.stack([ima_feature.clone(),ima_feature.clone()]).to(self.device)
leng = [leng,leng]
task = ['ImageQA_A','ImageQA_AT']
T = input_ids[0].size(1)
assert max_returned_tokens > T, f"max_returned_tokens {max_returned_tokens} should be greater than audio length {T}"
if model.max_seq_length < max_returned_tokens - 1:
raise NotImplementedError(
f"max_seq_length {model.max_seq_length} needs to be >= {max_returned_tokens - 1}"
)
list_output = [[] for i in range(8)]
tokens_A , token_T = next_token_image_batch(
model,
audio_feature.to(torch.float32).to(self.device),
ima_feature.to(torch.float32).to(self.device) ,
input_ids ,
whisper_lens = leng ,
task = task,
input_pos = torch.arange(0, T, device=self.device),
temperature=temperature,
top_k=top_k,
top_p=top_p
)
for i in range(7): list_output[i].append(tokens_A[i].tolist()[0])
list_output[7].append(token_T.tolist()[0])
text_end = False
index = 1
nums_generate = stream_stride
begin_generate = False
current_index = 0
input_pos = torch.tensor([T], device=self.device)
model_input_ids = [[] for i in range(8)]
for i in range(7):
tokens_A[i] = tokens_A[i].clone() + padded_text_vocabsize+ i * 4160
model_input_ids[i].append(tokens_A[i].clone().to(self.device).to(torch.int32))
model_input_ids[i].append(torch.tensor([layershift(4097,i)],device=self.device))
model_input_ids[i] = torch.stack(model_input_ids[i])
model_input_ids[-1].append(token_T.clone().to(torch.int32))
model_input_ids[-1].append(token_T.clone().to(torch.int32))
model_input_ids[-1] = torch.stack(model_input_ids[-1])
text_index = 0
is_text_end = False
for _ in tqdm(range(2, max_returned_tokens - T + 1)):
tokens_A , token_T = next_token_image_batch(model, None , None ,
input_ids = model_input_ids,
whisper_lens= None,
task = None,
input_pos = input_pos,
temperature=temperature,
top_k=top_k,
top_p=top_p)
if text_end:
token_T = torch.tensor([_pad_t], device=self.device)
if tokens_A[-1] == eos_id_a:
break
if token_T == eos_id_t:
text_end = True
for i in range(7): list_output[i].append(tokens_A[i].tolist()[0])
list_output[7].append(token_T.tolist()[0])
if index == 7:
begin_generate = True
if begin_generate:
current_index += 1
if current_index == nums_generate:
current_index = 0
snac = get_snac(list_output,index,nums_generate)
audio_stream = generate_audio_data(snac, self.snacmodel, self.device)
if is_text_end:
text_stream = ""
else:
text_stream, text_index, is_text_end = get_text_stream(list_output, text_index, self.text_tokenizer)
yield (audio_stream, text_stream)
if warm_up:
break
input_pos = input_pos.add_(1)
model_input_ids = [[] for i in range(8)]
for i in range(7):
tokens_A[i] = tokens_A[i].clone() + padded_text_vocabsize+ i * 4160
model_input_ids[i].append(tokens_A[i].clone().to(self.device).to(torch.int32))
model_input_ids[i].append(torch.tensor([layershift(4097,i)],device=self.device))
model_input_ids[i] = torch.stack(model_input_ids[i])
model_input_ids[-1].append(token_T.clone().to(torch.int32))
model_input_ids[-1].append(token_T.clone().to(torch.int32))
model_input_ids[-1] = torch.stack(model_input_ids[-1])
index += 1
text_tokens = list_output[-1]
if text_vocabsize in text_tokens:
text_tokens = text_tokens[:text_tokens.index(text_vocabsize)]
res_text = self.text_tokenizer.decode(torch.tensor(text_tokens))
print(f"text output: {res_text}")
if save_path is not None:
audiolist = reconscruct_snac(list_output)
audio = reconstruct_tensors(audiolist)
with torch.inference_mode():
audio_hat = self.snacmodel.decode(audio)
sf.write(save_path, audio_hat.squeeze().cpu().numpy(), 24000)
model.clear_kv_cache()
def test_vision_infer():
client = OmniVisionInference()
client.warm_up()
input_audio_path = './data/samples/vision_qa_audio.wav'
input_image_path = './data/samples/vision_qa_image.jpg'
res_text = ""
for audio_stream, text_stream in client.run_vision_AA_batch_stream(
input_audio_path,
input_image_path,
save_path="./vision_qa_output.wav"
):
res_text += text_stream
print(f"text_output: {res_text}")
if __name__ == "__main__":
test_vision_infer()