Update app.py

app.py

@@ -25,10 +25,10 @@ vilt_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetun
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 git_model_base.to(device)
-blip_model_base.to(device)
+# blip_model_base.to(device)
 #git_model_large.to(device)
 #blip_model_large.to(device)
-vilt_model.to(device)
+# vilt_model.to(device)
 
 def generate_answer_git(processor, model, image, question):
     # prepare image
@@ -42,10 +42,16 @@ def generate_answer_git(processor, model, image, question):
     generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50, return_dict_in_generate=True, output_scores=True)
     print('scores:')
     print(generated_ids.scores)
-    scoresList0 = torch.softmax(generated_ids.scores[0], dim=1).flatten().tolist()
-    print(scoresList0)
-    scoresList1 = torch.softmax(generated_ids.scores[1], dim=1).flatten().tolist()
-    print(scoresList1)
+    # scoresList0 = torch.softmax(generated_ids.scores[0], dim=1).flatten().tolist()
+    # print(scoresList0)
+    # scoresList1 = torch.softmax(generated_ids.scores[1], dim=1).flatten().tolist()
+    # print(scoresList1)
+    idx = generated_ids.scores[0].argmax(-1).item()
+    idx1 = generated_ids.scores[1].argmax(-1).item()
+    print(idx, idx1)
+    ans = model.config.id2label[idx]
+    ans1 = model.config.id2label[idx1]
+    print(ans, ans1)
     print('sequences:')
     print(generated_ids.sequences)
     print(generated_ids)
@@ -82,13 +88,13 @@ def generate_answers(image, question):
 
     # answer_git_large = generate_answer_git(git_processor_large, git_model_large, image, question)
 
-    answer_blip_base = generate_answer_blip(blip_processor_base, blip_model_base, image, question)
+    # answer_blip_base = generate_answer_blip(blip_processor_base, blip_model_base, image, question)
 
     # answer_blip_large = generate_answer_blip(blip_processor_large, blip_model_large, image, question)
 
-    answer_vilt = generate_answer_vilt(vilt_processor, vilt_model, image, question)
+    # answer_vilt = generate_answer_vilt(vilt_processor, vilt_model, image, question)
 
-    return answer_git_base, answer_blip_base, answer_vilt
+    return answer_git_base
 
    
 examples = [["cats.jpg", "How many cats are there?"], ["stop_sign.png", "What's behind the stop sign?"], ["astronaut.jpg", "What's the astronaut riding on?"]]