Update the model weights loading method

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
-from transformers import AutoModel, AutoTokenizer
+from huggingface_hub import hf_hub_download
+import open_clip
 from PIL import Image
 import json
 import os
@@ -11,14 +12,29 @@ from tqdm import tqdm
 ASSETS_DIR = "assets"
 EXAMPLES_DIR = "examples"
 PATH_TEXT_PROMPTS = os.path.join(ASSETS_DIR, "prompt_fetal_view.json")
-MODEL_NAME = "numansaeed/fetalclip-model"  # Replace with your actual model name
+PATH_FETALCLIP_CONFIG = os.path.join(ASSETS_DIR, "FetalCLIP_config.json")
+MODEL_NAME = "numansaeed/fetalclip-model"
 
 def load_model():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
-    # Load model and tokenizer from Hugging Face
-    model = AutoModel.from_pretrained(MODEL_NAME)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    # Load and register model configuration
+    with open(PATH_FETALCLIP_CONFIG, "r") as file:
+        config_fetalclip = json.load(file)
+    open_clip.factory._MODEL_CONFIGS["FetalCLIP"] = config_fetalclip
+    
+    # Download model weights from Hugging Face Hub
+    weights_path = hf_hub_download(
+        repo_id=MODEL_NAME,
+        filename="FetalCLIP_weights.pt"
+    )
+    
+    # Load the FetalCLIP model and preprocessing transforms
+    model, _, preprocess_test = open_clip.create_model_and_transforms(
+        "FetalCLIP",
+        pretrained=weights_path
+    )
+    tokenizer = open_clip.get_tokenizer("FetalCLIP")
     
     model = model.float()
     model.eval()
@@ -35,8 +51,8 @@ def load_model():
         for plane, prompts in tqdm(text_prompts.items()):
             list_plane.append(plane)
             
-            prompts = tokenizer(prompts, return_tensors="pt", padding=True).to(device)
-            text_features = model(**prompts).last_hidden_state.mean(dim=1)
+            prompts = tokenizer(prompts).to(device)
+            text_features = model.encode_text(prompts)
             text_features /= text_features.norm(dim=-1, keepdim=True)
             
             text_features = text_features.mean(dim=0).unsqueeze(0)
@@ -45,10 +61,10 @@ def load_model():
             list_text_features.append(text_features)
     text_features = torch.stack(list_text_features)[:,0]
     
-    return model, tokenizer, text_features, list_plane, device
+    return model, preprocess_test, text_features, list_plane, device
 
 # Load model and text features at startup
-model, tokenizer, text_features, list_plane, device = load_model()
+model, preprocess_test, text_features, list_plane, device = load_model()
 
 def process_image(image, model_type, top_k):
     if image is None:
@@ -57,11 +73,11 @@ def process_image(image, model_type, top_k):
     try:
         # Preprocess image
         img = make_image_square_with_zero_padding(Image.fromarray(image))
-        img = tokenizer(images=img, return_tensors="pt").to(device)
+        img = preprocess_test(img).unsqueeze(0).to(device)
         
         # Get image features
         with torch.no_grad():
-            image_features = model(**img).last_hidden_state.mean(dim=1)
+            image_features = model.encode_image(img)
             image_features /= image_features.norm(dim=-1, keepdim=True)
         
         # Calculate similarity scores