Fix project_encoder

vit_gpt2/configuration_vit_gpt2.py

@@ -47,9 +47,9 @@ class ViTGPT2Config(PretrainedConfig):
         _project_encoder = getattr(self.text_config, "project_encoder", None)
         if project_encoder is not None and _project_encoder is not None:
             assert project_encoder == _project_encoder
-        elif project_encoder:
+        elif project_encoder is not None:
             _project_encoder = project_encoder
-        elif _project_encoder:
+        elif _project_encoder is not None:
             project_encoder = _project_encoder
         else:
             project_encoder = False

vit_gpt2/modeling_flax_vit_gpt2_lm.py

@@ -534,6 +534,7 @@ class FlaxViTGPT2LMForConditionalGeneration(FlaxViTGPT2LMPreTrainedModel):
                 vision_pretrained_model_name_or_path, *vision_model_args, **vision_kwargs
             )
 
+        project_encoder = kwargs.pop("project_encoder", None)
         if text_model is None:
             assert (
                 text_pretrained_model_name_or_path is not None
@@ -542,6 +543,8 @@ class FlaxViTGPT2LMForConditionalGeneration(FlaxViTGPT2LMPreTrainedModel):
             if "config" not in text_kwargs:
                 text_config = GPT2Config.from_pretrained(text_pretrained_model_name_or_path)
                 text_config.project_encoder = text_kwargs.pop("project_encoder", None)
+                if project_encoder is not None:
+                    text_config.project_encoder = project_encoder
                 text_kwargs["config"] = text_config
 
             text_kwargs["config"].add_cross_attention = True
@@ -553,8 +556,6 @@ class FlaxViTGPT2LMForConditionalGeneration(FlaxViTGPT2LMPreTrainedModel):
 
         # instantiate config with corresponding kwargs
         dtype = kwargs.pop("dtype", jnp.float32)
-        project_encoder = kwargs.pop("project_encoder", None)
-
         config = ViTGPT2Config.from_vision_text_configs(
             vision_model.config, text_model.config, project_encoder=project_encoder, **kwargs
         )