Update linear_mapping.py

linear_mapping.py

@@ -2,6 +2,7 @@ from config import LinearMappingConfig
 from transformers import (
     GPT2TokenizerFast, GPT2LMHeadModel, AutoModel,
     CLIPVisionModel, AutoProcessor, BatchEncoding,
+    AutoConfig, CLIPVisionConfig
 )
 from transformers.models.gpt2.modeling_gpt2 import GPT2DoubleHeadsModelOutput
 import torch
@@ -104,9 +105,11 @@ class ImagePrefix(nn.Module):
 
     def __init__(self, config: LinearMappingConfig):
         super().__init__()
-        self.encoder = AutoModel.from_pretrained(config.image_model)
-        if "clip" in config.image_model:
-            self.encoder = CLIPVisionModel.from_pretrained(config.image_model)
+        clip_config = CLIPVisionConfig.from_pretrained(config.image_model)
+
+        self.encoder = CLIPVisionModel(clip_config)
+        if config.image_from_pretrained:
+            self.encoder = self.encoder.from_pretrained(config.image_model)
 
         if config.freeze_image_model:
             for param in self.encoder.parameters():
@@ -128,7 +131,9 @@ class LinearMapping(nn.Module):
     def __init__(self, config: LinearMappingConfig):
         super().__init__()
         self.image_prefix = ImagePrefix(config)
-        self.language_model = GPT2LMHeadModel.from_pretrained(config.text_model)
+        self.language_model = GPT2LMHeadModel(AutoConfig.from_pretrained(config.text_model))
+        if config.text_from_pretrained:
+            self.language_model = self.language_model.from_pretrained(config.text_model)
         self.processor = LinearMappingProcessor(config)
         self.tokenizer = self.processor.tokenizer
         self.image_processor = self.processor.image_processor