Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import whisperx | |
| from transformers import AutoTokenizer | |
| from transformers import AutoModelForCausalLM | |
| from transformers import CLIPVisionModel, CLIPImageProcessor | |
| import peft | |
| import gradio as gr | |
| device = 'cpu' | |
| model_name = "microsoft/phi-2" | |
| whisper_model = whisperx.load_model('small', device='cpu', compute_type='float32') | |
| image_processor = CLIPImageProcessor.from_pretrained('openai/clip-vit-base-patch32') | |
| clip_model = CLIPVisionModel.from_pretrained('openai/clip-vit-base-patch32') | |
| tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=False) | |
| tokenizer.pad_token = tokenizer.eos_token | |
| tokenizer.bos_token = tokenizer.eos_token | |
| phi2_model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| trust_remote_code=True, | |
| device_map = 'cpu' | |
| ) | |
| phi2_model.config.use_cache = False | |
| def CLIP_embeddings(image): | |
| _ = clip_model.requires_grad_(False) | |
| image = image_processor(images=image, return_tensors="pt") | |
| image_out = clip_model(image['pixel_values'].to(device=clip_model.device), output_hidden_states=True) | |
| return features(image_out) | |
| def embed_audio(file_name): | |
| result = whisper_model.transcribe(file_name) | |
| res_text = '' | |
| for segment in result['segments']: | |
| res_text = res_text + segment['text'] | |
| return res_text.strip() | |
| def features(image_out): | |
| image_features = image_out.hidden_states[-1] | |
| return image_features[:, 1:, :] | |
| def embed_text(text): | |
| input_tokens = tokenizer(text, return_tensors="pt", return_attention_mask=False) | |
| return phi2_model.get_input_embeddings()(input_tokens.input_ids) | |
| class ResBlock(nn.Module): | |
| def __init__(self, input_size): | |
| super().__init__() | |
| self.pre_norm = nn.LayerNorm(input_size) | |
| self.proj = nn.Sequential( | |
| nn.Linear(input_size, input_size), | |
| nn.GELU(), | |
| nn.Linear(input_size, input_size) | |
| ) | |
| def forward(self, x): | |
| x = self.pre_norm(x) | |
| return x + self.proj(x) | |
| class Projection_Model(nn.Module): | |
| def __init__( | |
| self, | |
| dim_input_CLIP = 768, | |
| dim_input_Phi2 = 2560 | |
| ): | |
| super(Projection_Model, self).__init__() | |
| self.projection_img = nn.Linear( | |
| dim_input_CLIP, dim_input_Phi2, bias=False | |
| ) | |
| self.resblock = ResBlock(dim_input_Phi2) | |
| def forward(self, x): | |
| x = self.projection_img(x) | |
| return self.resblock(x) | |
| model = Projection_Model() | |
| model.projection_img.load_state_dict(torch.load("projection.pth", map_location='cpu')) | |
| model.resblock.load_state_dict(torch.load("block.pth", map_location='cpu')) | |
| def embeddings_image(image): | |
| clip_embeddings = CLIP_embeddings(image) | |
| return model(clip_embeddings) | |
| user = "TharunSivamani" | |
| model_name = "qlora-phi2" | |
| model_id = f"{user}/{model_name}" | |
| phi2_model_peft = peft.PeftModel.from_pretrained(phi2_model, model_id) | |
| def inference(image=None, audio=None, text=None): | |
| if len(text) == 0: | |
| text = None | |
| if image is None and audio is None and text is None: | |
| return None | |
| context = tokenizer("Context: ", return_tensors="pt", return_attention_mask=False) | |
| input_embeds = phi2_model_peft.get_input_embeddings()(context.input_ids) | |
| if image is not None: | |
| query = text | |
| image_embeds = embeddings_image(image) | |
| input_embeds = torch.cat((input_embeds, image_embeds), dim=1) | |
| if audio is not None: | |
| audio_transcribed = embed_audio(audio) | |
| audio_embeds = embed_text(audio_transcribed) | |
| input_embeds = torch.cat((input_embeds, audio_embeds), dim=1) | |
| if text is not None: | |
| query = text | |
| text_embeds = embed_text(text) | |
| input_embeds = torch.cat((input_embeds, text_embeds), dim=1) | |
| question = tokenizer(" Question: " + query, return_tensors="pt", return_attention_mask=False) | |
| question_embeds = phi2_model_peft.get_input_embeddings()(question.input_ids) | |
| input_embeds = torch.cat((input_embeds, question_embeds), dim=1) | |
| answer = tokenizer(" Answer: ", return_tensors="pt", return_attention_mask=False) | |
| answer_embeds = phi2_model_peft.get_input_embeddings()(answer.input_ids) | |
| input_embeds = torch.cat((input_embeds, answer_embeds), dim=1) | |
| result = phi2_model_peft.generate(inputs_embeds=input_embeds, bos_token_id = tokenizer.bos_token_id) | |
| final_ans = tokenizer.batch_decode(result)[0] | |
| final_ans = final_ans.split(tokenizer.eos_token) | |
| if final_ans[0] == '': | |
| return final_ans[1] | |
| else: | |
| return final_ans[0] | |
| demo = gr.Interface( | |
| fn = inference, | |
| inputs = [ | |
| gr.Image(label="Image Input"), | |
| gr.Audio(label="Audio Input", sources=["microphone", "upload"], type="filepath"), | |
| gr.Textbox(label="Text Input"), | |
| ], | |
| outputs = [ | |
| gr.Textbox(label='Answer'), | |
| ], | |
| ) | |
| demo.launch() |