Update nllb_mrl.py

nllb_mrl.py

@@ -4,7 +4,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from open_clip import create_model, get_tokenizer
-from open_clip.pretrained import get_pretrained_cfg
 from open_clip.transform import PreprocessCfg, image_transform_v2
 from PIL import Image
 from transformers import PretrainedConfig, PreTrainedModel
@@ -16,7 +15,7 @@ class MatryoshkaNllbClipConfig(PretrainedConfig):
         clip_model_name: str = "",
         target_resolution: int = -1,
         mrl_resolutions: List[int] = [],
-        preprocess_cfg: Union[PreprocessCfg, None] = None,
+        preprocess_cfg: Union[dict, None] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -53,11 +52,16 @@ class MatryoshkaNllbClip(PreTrainedModel):
         if isinstance(device, str):
             device = torch.device(device)
         self.config = config
-        self.model = create_model(
-            config.clip_model_name, output_dict=True
+        self.model = create_model(config.clip_model_name, output_dict=True)
+        pp_cfg = PreprocessCfg(
+            size=config.preprocess_cfg["size"],
+            mean=config.preprocess_cfg["mean"],
+            std=config.preprocess_cfg["std"],
+            interpolation=config.preprocess_cfg["interpolation"],
+            resize_mode=config.preprocess_cfg["resize_mode"],
         )
         self.transform = image_transform_v2(
-            config.preprocess_cfg,
+            pp_cfg,
             is_train=False,
         )
         self._device = device
@@ -108,7 +112,7 @@ class MatryoshkaNllbClip(PreTrainedModel):
                     )
                 features = self.matryoshka_layer.layers[str(resolution)](features)
             return F.normalize(features, dim=-1) if normalize else features
-    
+
     def encode_text(
         self,
         text,