# 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() # # 1234232434232 # # 1234232434232 # # 1234232434232 # # 1234232434232 # # 1234232434232 # # 1234232434232 # # 5069 # # 3670