Update nllb_mrl.py

nllb_mrl.py

@@ -15,12 +15,14 @@ class MatryoshkaNllbClipConfig(PretrainedConfig):
         clip_model_name: str = "",
         target_resolution: int = -1,
         mrl_resolutions: List[int] = [],
+        preprocess_cfg: Union[dict, None] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.clip_model_name = clip_model_name
         self.target_resolution = target_resolution
         self.mrl_resolutions = mrl_resolutions
+        self.preprocess_cfg = preprocess_cfg
 
 
 class MatryoshkaLayer(nn.Module):
@@ -50,10 +52,14 @@ 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"],
         )
-        pp_cfg = PreprocessCfg(**self.model.visual.preprocess_cfg)
         self.transform = image_transform_v2(
             pp_cfg,
             is_train=False,
@@ -106,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,