Update app.py

app.py

@@ -2,7 +2,7 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoProcessor, AutoModelForImageTextRetrieval
+from transformers import AutoProcessor, AutoModel, CLIPVisionModel
 import torch.nn.functional as F
 
 #---------------------------------
@@ -13,16 +13,18 @@ def load_biomedclip_model():
     """Loads the BiomedCLIP model and tokenizer."""
     biomedclip_model_name = 'microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224'
     processor = AutoProcessor.from_pretrained(biomedclip_model_name)
-    model = AutoModelForImageTextRetrieval.from_pretrained(biomedclip_model_name).cuda().eval()
-    return model, processor
+    config = AutoModel.from_pretrained(biomedclip_model_name).config
+    vision_model = CLIPVisionModel.from_pretrained(config.vision_config._name_or_path, torch_dtype=torch.float16).cuda().eval()
+    text_model = AutoModel.from_pretrained(config.text_config._name_or_path).cuda().eval()
+    return vision_model, text_model, processor
 
-def compute_similarity(image, text, biomedclip_model, biomedclip_processor):
+def compute_similarity(image, text, vision_model, text_model, biomedclip_processor):
     """Computes similarity scores using BiomedCLIP."""
     with torch.no_grad():
-        inputs = biomedclip_processor(text=text, images=image, return_tensors="pt", padding=True).to(biomedclip_model.device)
-        outputs = biomedclip_model(**inputs)
-        image_embeds = outputs.image_embeds
-        text_embeds = outputs.text_embeds
+        inputs = biomedclip_processor(text=text, images=image, return_tensors="pt", padding=True).to(text_model.device)
+        text_embeds = text_model(**inputs).last_hidden_state[:,0,:] # Extract the [CLS] token
+        image_inputs = biomedclip_processor(images=image, return_tensors="pt").to(vision_model.device)
+        image_embeds = vision_model(**image_inputs).last_hidden_state[:,0,:] # Extract the image embedding
     image_embeds = F.normalize(image_embeds, dim=-1)
     text_embeds = F.normalize(text_embeds, dim=-1)
     similarity = (text_embeds @ image_embeds.transpose(-1, -2)).squeeze()
@@ -55,12 +57,12 @@ def gradio_ask(user_message, chatbot, chat_state):
     return '', chatbot, chat_state
 
 @spaces.GPU
-def gradio_answer(chatbot, chat_state, img_list, biomedclip_model, biomedclip_processor, similarity_output):
+def gradio_answer(chatbot, chat_state, img_list, vision_model, text_model, biomedclip_processor, similarity_output):
     """Computes and displays similarity scores."""
     if not img_list:
         return chatbot, chat_state, img_list, similarity_output
 
-    similarity_score = compute_similarity(img_list[0], chatbot[-1][0], biomedclip_model, biomedclip_processor)
+    similarity_score = compute_similarity(img_list[0], chatbot[-1][0], vision_model, text_model, biomedclip_processor)
     print(f'Similarity Score is: {similarity_score}')
 
     similarity_text = f"Similarity Score: {similarity_score:.3f}"
@@ -77,7 +79,7 @@ examples_list=[
                 ]
 
 # Load models and related resources outside of the Gradio block for loading on startup
-biomedclip_model, biomedclip_processor = load_biomedclip_model()
+vision_model, text_model, biomedclip_processor = load_biomedclip_model()
 
 with gr.Blocks() as demo:
     gr.Markdown(title)
@@ -100,7 +102,7 @@ with gr.Blocks() as demo:
     upload_button.click(upload_img, [image, text_input, chat_state, similarity_output], [image, text_input, upload_button, chat_state, img_list, similarity_output])
 
     text_input.submit(gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state]).then(
-        gradio_answer, [chatbot, chat_state, img_list, biomedclip_model, biomedclip_processor, similarity_output], [chatbot, chat_state, img_list, similarity_output]
+        gradio_answer, [chatbot, chat_state, img_list, vision_model, text_model, biomedclip_processor, similarity_output], [chatbot, chat_state, img_list, similarity_output]
     )
     clear.click(gradio_reset, [chat_state, img_list, similarity_output], [chatbot, image, text_input, upload_button, chat_state, img_list, similarity_output], queue=False)