Update nllb_mrl.py

nllb_mrl.py

@@ -3,8 +3,8 @@ from typing import List, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from huggingface_hub import PyTorchModelHubMixin
-from open_clip import create_model_and_transforms, get_tokenizer
+from open_clip import create_model, get_tokenizer
+from open_clip.transform import PreprocessCfg, image_transform_v2
 from PIL import Image
 from transformers import PretrainedConfig, PreTrainedModel
 
@@ -13,14 +13,12 @@ class MatryoshkaNllbClipConfig(PretrainedConfig):
     def __init__(
         self,
         clip_model_name: str = "",
-        clip_model_version: str = "",
         target_resolution: int = -1,
         mrl_resolutions: List[int] = [],
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.clip_model_name = clip_model_name
-        self.clip_model_version = clip_model_version
         self.target_resolution = target_resolution
         self.mrl_resolutions = mrl_resolutions
 
@@ -46,14 +44,19 @@ class MatryoshkaLayer(nn.Module):
 
 class MatryoshkaNllbClip(PreTrainedModel):
     config_class = MatryoshkaNllbClipConfig
-    
+
     def __init__(self, config: MatryoshkaNllbClipConfig, device):
         super().__init__(config)
         if isinstance(device, str):
             device = torch.device(device)
         self.config = config
-        self.model, _, self.transform = create_model_and_transforms(
-            config.clip_model_name, config.clip_model_version, output_dict=True
+        self.model = create_model(
+            config.clip_model_name, output_dict=True
+        )
+        pp_cfg = PreprocessCfg(**self.model.visual.preprocess_cfg)
+        self.transform = image_transform_v2(
+            pp_cfg,
+            is_train=False,
         )
         self._device = device
         self.model.to(device)