move image generation init to build_blocks

Author: francoishernandez

eole/models/model.py

@@ -416,7 +416,6 @@ def from_config(
         model = cls.build_blocks(
             model_config, vocabs, running_config=running_config
         )  # corresponds to build_task_specific_model
-        model._maybe_init_image_generation(model_config, running_config)
         # TODO: handle this better at some point?
         model.share_decoder_embeddings = model_config.share_decoder_embeddings
         # generator -> shall it be called within build_blocks?
@@ -976,7 +975,7 @@ def build_blocks(cls, model_config, vocabs, running_config=None):
             share_embeddings=model_config.share_embeddings,
         )
         decoder = build_decoder(model_config, running_config=running_config)
-        return cls(
+        model = cls(
             encoder=encoder,
             decoder=decoder,
             adapter=adapter,
@@ -988,7 +987,9 @@ def build_blocks(cls, model_config, vocabs, running_config=None):
             image_end_token_id=model_config.encoder.image_end_token_id,
             vit_position_embeddings=model_config.vit_position_embeddings,
         )
+        model._maybe_init_image_generation(model_config, running_config)
         # from there, the base blocks exist, and the rest is done in the from_opt from base class
+        return model
 
     def embed_vision_language_features(self, src, images):
         # TODO: test with batch > 1?