inference_vision.py

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()