Update app.py

app.py

@@ -5,6 +5,7 @@ from transformers import AutoTokenizer,BitsAndBytesConfig, AutoModelForCausalLM,
 import torch
 from peft import PeftModel
 import torch.nn as nn
+import whisperx
 
 clip_model_name = "openai/clip-vit-base-patch32"
 phi_model_name  = "microsoft/phi-2"
@@ -15,6 +16,8 @@ IMAGE_TOKEN_ID = 23893 # token for word comment
 device = "cuda" if torch.cuda.is_available() else "cpu"
 clip_embed = 768
 phi_embed  = 2560
+compute_type = "float16"
+audio_batch_size = 16
 
 class SimpleResBlock(nn.Module):
     def __init__(self, phi_embed):
@@ -33,8 +36,8 @@ class SimpleResBlock(nn.Module):
 clip_model = CLIPVisionModel.from_pretrained(clip_model_name).to(device)
 projection = torch.nn.Linear(clip_embed, phi_embed).to(device)
 resblock = SimpleResBlock(phi_embed).to(device)
-
 phi_model = AutoModelForCausalLM.from_pretrained(phi_model_name,trust_remote_code=True).to(device)
+audio_model = whisperx.load_model("large-v2", device, compute_type=compute_type)
 
 # load weights
 model_to_merge = PeftModel.from_pretrained(phi_model,'./model_chkpt/lora_adaptor')
@@ -53,9 +56,19 @@ def model_generate_ans(img,val_q):
         clip_val_outputs = clip_model(**image_processed).last_hidden_state[:,1:,:]
         val_image_embeds = projection(clip_val_outputs)
         val_image_embeds = resblock(val_image_embeds).to(torch.float16)
-    
         img_token_tensor = torch.tensor(IMAGE_TOKEN_ID).to(device)
         img_token_embeds = merged_model.model.embed_tokens(img_token_tensor).unsqueeze(0).unsqueeze(0)
+
+        # audio
+        # audio  = whisperx.load_audio(audio_file)
+        # result = audio_model.transcribe(audio, batch_size=audio_batch_size)
+
+        # audio_txt = []
+        # for s in result["segments"]:
+        #    audio_txt.append(s['text'])
+        #    print(s['text'])
+
+        # audio_text = "".join(audio_txt)
     
         val_q_tokenised = tokenizer(val_q, return_tensors="pt", return_attention_mask=False)['input_ids'].squeeze(0).to(device)
         val_q_embeds    = merged_model.model.embed_tokens(val_q_tokenised).unsqueeze(0)