Update app.py

app.py

@@ -1,41 +1,60 @@
-import gradio as gr
 import torch
 import whisperx
-from model import MainQLoraModel
+import gradio as gr
+from peft import PeftModel
 from configs import get_config_phase2
-from transformers import AutoTokenizer, AutoProcessor
+from transformers import AutoTokenizer, AutoProcessor, CLIPVisionModel, AutoModelForCausalLM
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+clip_model = CLIPVisionModel.from_pretrained(config.get("clip_model_name"))
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    config.get("phi2_model_name"),
+    low_cpu_mem_usage=True,
+    return_dict=True,
+    torch_dtype=torch.float32,
+    trust_remote_code=True
+)
+
+
+ckpts = "ckpts/Qlora_adaptor/"
+phi2_model = PeftModel.from_pretrained(base_model, ckpts)
+phi2_model = phi2_model.merge_and_unload().to(device)
+
+projection_layer = torch.nn.Linear(config.get("clip_embed"), config.get("phi_embed"))
+projection_layer.load_state_dict(torch.load('./ckpts/model_phase2.pth', map_location=config.get("device")))
 
-# get config
-config = get_config_phase2() 
 # tokenizer
 tokenizer  = AutoTokenizer.from_pretrained(config.get("phi2_model_name"), trust_remote_code=True)
 processor  = AutoProcessor.from_pretrained(config.get("clip_model_name"), trust_remote_code=True)
-llmModel = MainQLoraModel(tokenizer, config).to(config.get("device"))
-audio_model = whisperx.load_model('tiny', 'cpu', compute_type="float16")
+
+audio_model = whisperx.load_model('tiny', 'cpu', compute_type="float32")
 
 
 def generate_answers(img=None, aud = None, q = None, max_tokens = 30):
+    print(img, aud, q)
           
     batch_size = 1
     start_iq = tokenizer.encode("<iQ>")
     end_iq = tokenizer.encode("</iQ>")
     start_iq_embeds = torch.tensor(start_iq).repeat(batch_size, 1)
     end_iq_embeds = torch.tensor(end_iq).repeat(batch_size, 1)
-    start_iq_embeds = llmModel.phi2_model.model.model.embed_tokens(start_iq_embeds.to(config.get("device")))
-    end_iq_embeds = llmModel.phi2_model.model.model.embed_tokens(end_iq_embeds.to(config.get("device")))
+    start_iq_embeds = phi2_model.model.embed_tokens(start_iq_embeds.to(config.get("device")))
+    end_iq_embeds = phi2_model.model.embed_tokens(end_iq_embeds.to(config.get("device")))
     
     inputs_embeddings = []
     inputs_embeddings.append(start_iq_embeds)
 
-    predicted_caption = torch.full((batch_size, max_tokens), llmModel.EOS_TOKEN_ID, dtype=torch.long, device=config.get('device'))
+    predicted_caption = torch.full((batch_size, max_tokens), 50256, dtype=torch.long, device=config.get('device'))
     
-    if images is not None:
-        images = processor(images=img, return_tensors="pt").to(config.get("device"))
+    if img is not None:
+        images = processor(images=img, return_tensors="pt")['pixel_values'].to(device)
         images = {'pixel_values': images.to(config.get("device"))}
-        clip_outputs = llmModel.clip_model(**images)
+        clip_outputs = clip_model(**images)
         # remove cls token
         images = clip_outputs.last_hidden_state[:, 1:, :]
-        image_embeddings = llmModel.projection_layer(images).to(torch.float16)
+        image_embeddings = projection_layer(images).to(torch.float32)
         inputs_embeddings.append(image_embeddings)
     
     if aud is not None:
@@ -44,13 +63,14 @@ def generate_answers(img=None, aud = None, q = None, max_tokens = 30):
         for seg in trans['segments']:
             audio_res += seg['text']
         audio_res = audio_res.strip()
+        print(audio_res)
         audio_tokens = tokenizer(q,return_tensors="pt", return_attention_mask=False)['input_ids']
-        audio_embeds = llmModel.phi2_model.model.model.embed_tokens(audio_tokens.to(config.get("device")))
+        audio_embeds = phi2_model.model.embed_tokens(audio_tokens.to(config.get("device")))
         inputs_embeddings.append(audio_embeds)
         
     if q is not None:
         ques = tokenizer(q, return_tensors="pt", return_attention_mask=False)['input_ids']
-        q_embeds = llmModel.phi2_model.model.model.embed_tokens(ques.to(config.get("device")))
+        q_embeds = phi2_model.model.embed_tokens(ques.to(config.get("device")))
         inputs_embeddings.append(q_embeds)
         
     inputs_embeddings.append(end_iq_embeds)
@@ -58,11 +78,12 @@ def generate_answers(img=None, aud = None, q = None, max_tokens = 30):
     combined_embeds  = torch.cat(inputs_embeddings, dim=1)
 
     for pos in range(max_tokens - 1):
-        model_output_logits = llmModel.phi2_model.forward(inputs_embeds = combined_embeds)['logits']
+        model_output_logits = phi2_model.forward(inputs_embeds = combined_embeds)['logits']
+        print(model_output_logits.shape)
         predicted_word_token_logits = model_output_logits[:, -1, :].unsqueeze(1)
         predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1)
         predicted_caption[:, pos] = predicted_word_token.view(1,-1).to('cpu')
-        next_token_embeds = llmModel.phi2_model.model.model.embed_tokens(predicted_word_token)
+        next_token_embeds = phi2_model.model.embed_tokens(predicted_word_token)
         combined_embeds   = torch.cat([combined_embeds, next_token_embeds], dim=1)
     predicted_captions_decoded = tokenizer.batch_decode(predicted_caption,ignore_index = 50256)[0]
     return predicted_captions_decoded