Update nllb_mrl.py

nllb_mrl.py

@@ -16,12 +16,14 @@ class MatryoshkaNllbClipConfig(PretrainedConfig):
         clip_model_name: str = "",
         target_resolution: int = -1,
         mrl_resolutions: List[int] = [],
+        preprocess_cfg: Union[PreprocessCfg, 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):
@@ -54,10 +56,8 @@ class MatryoshkaNllbClip(PreTrainedModel):
         self.model = create_model(
             config.clip_model_name, output_dict=True
         )
-        preprocess_cfg = get_pretrained_cfg(config.clip_model_name, "v1")
-        pp_cfg = PreprocessCfg(preprocess_cfg)
         self.transform = image_transform_v2(
-            pp_cfg,
+            config.preprocess_cfg,
             is_train=False,
         )
         self._device = device