clean up and structure

Author: francoishernandez

eole/bin/convert/HF_mappings.py

@@ -37,7 +37,7 @@
 MODEL_OVERRIDES = {
     "LlamaForCausalLM": {},  # default
     "MistralForCausalLM": {},
-    "Qwen2ForCausalLM": { # for bagel, but we need to add some conditions to keep supporting real qwen2...
+    "Bagel": {  # bagel's arch is actually Qwen2, but requires specific mapping
         "decoder_layer_prefix": "language_model.model.layers.",
         "decoder.layer_norm.weight": "language_model.model.norm.weight",
         "decoder.layer_norm_moe_gen.weight": "language_model.model.norm_moe_gen.weight",
@@ -95,12 +95,23 @@
         "config": {
             "add_qkvbias": True,
             "add_final_linear_bias": False,
-            # "ffn_layernorm": True,
+            "adapter": "bagel",
+            "vit_position_embeddings": True,
             "decoder": {
                 "query_norm": True,
                 "key_norm": True,
             },
-        }
+            "encoder": {
+                "mlp_activation_fn": "gelu-tanh",
+                "add_ffnbias": True,
+                "add_final_linear_bias": True,
+                "add_qkvbias": True,
+                "layer_norm": "standard",
+                "patch_conv_bias": True,
+                "patch_conv_linear": True,
+                "layernorm_pre": False,  # implies post layernorm
+            },
+        },
     },
     "Qwen3ForCausalLM": {
         "decoder": {
@@ -412,7 +423,6 @@
         "Gemma2ForCausalLM": "gemma-rms",
         "M2M100ForConditionalGeneration": "standard",
         "Gemma3ForConditionalGeneration": "gemma-rms",
-        "Qwen2ForCausalLM": "rms",
     },
 )
 
@@ -446,7 +456,7 @@
         "Mistral3ForConditionalGeneration": VisionTransformerLMModelConfig,
         "Gemma3ForConditionalGeneration": VisionTransformerLMModelConfig,
         "M2M100ForConditionalGeneration": TransformerModelConfig,
-        "Qwen2ForCausalLM": VisionTransformerLMModelConfig,
+        "Bagel": VisionTransformerLMModelConfig,
     },
 )
 

eole/bin/convert/convert_HF.py

@@ -119,7 +119,7 @@ def download_file_from_hub(file_name, required=True):
 
         # Fetch required and optional files
         paths = {
-            "config_path": get_file_fn("llm_config.json", required=False), # hard patch for bagel
+            "config_path": get_file_fn("llm_config.json", required=False) or get_file_fn("config.json", required=True),
             "tokenizer_config_json": get_file_fn("tokenizer_config.json", required=True),
             "generation_config_json": get_file_fn("generation_config.json", required=False),
             "tokenizer_model": get_file_fn("tokenizer.model", required=False)
@@ -128,7 +128,8 @@ def download_file_from_hub(file_name, required=True):
             "wmap_path": get_file_fn("model.safetensors.index.json", required=False)
             or get_file_fn("pytorch_model.bin.index.json", required=False),
             "model_path": get_file_fn("model.safetensors", required=False)
-            or get_file_fn("pytorch_model.bin", required=False) or get_file_fn("ema.safetensors", required=False),
+            or get_file_fn("pytorch_model.bin", required=False)
+            or get_file_fn("ema.safetensors", required=False),
             "special_tokens_json": get_file_fn("special_tokens_map.json", required=False),
             "vision_config_path": get_file_fn("vit_config.json", required=False),
             "ae_model_path": get_file_fn("ae.safetensors", required=False),
@@ -162,6 +163,8 @@ def __getattr__(self, name):
 
     @property
     def arch(self):
+        if self.model_dir == "ByteDance-Seed/BAGEL-7B-MoT":
+            return "Bagel"
         return self.config["architectures"][0]
 
     @property
@@ -280,8 +283,6 @@ def build_config_dict(hf):
         other_config = config  # save what is not text/vision for later use
         config = config.get("text_config", config)
 
-    print("VISION_CONFIG:", vision_config)
-
     model_config = {}
     training_config = {}
 
@@ -360,28 +361,22 @@ def build_config_dict(hf):
         model_config["projector_activation_fn"] = other_config.get("projector_hidden_act", "gelu")
         model_config["spatial_merge_size"] = other_config.get("spatial_merge_size", None)
 
-    if arch == "Qwen2ForCausalLM":
-        model_config["adapter"] = "bagel"
+    if arch == "Bagel":
         model_config["encoder"] = {
-            "mlp_activation_fn": "gelu-tanh",  # no up_proj it seems
             "hidden_size": vision_config.get("hidden_size", 1152),
-            # "image_size": vision_config["image_size"],
-            "image_size": 1024,
+            "image_size": 1024,  # 980 for VIT (vit_config.json), 1024 for VAE
             "patch_size": vision_config["patch_size"],
             "heads": vision_config["num_attention_heads"],
             "heads_kv": vision_config["num_attention_heads"],
-            "layers": 26, # 27 in config, but actually 26 in safetensors...
+            "layers": 26,  # 27 in config, but actually 26 in safetensors
             "transformer_ff": vision_config["intermediate_size"],
             # siglip style learned position embeddings (like gemma3)
             "position_encoding_type": PositionEncodingType.Learned,
             "n_positions": (vision_config["image_size"] // vision_config["patch_size"]) ** 2,
-            "add_ffnbias": True,
-            "add_final_linear_bias": True,
-            "add_qkvbias": True,
-            "layer_norm": "standard",
-            "patch_conv_bias": True,
-            "layernorm_pre": False,  # implies post layernorm
             "image_token_id": 151654,
+            "image_start_token_id": 151652,
+            "image_end_token_id": 151653,
+            "max_patches_per_side": 70,
         }
 
     if arch == "Gemma3ForConditionalGeneration":
@@ -679,7 +674,7 @@ def build_shards(model_config, hf, args, params):
         eole_safetensor = {}
 
         def build_first_shard(hf, eole_safetensor):
-            # let's add AE here
+            # let's add AE here (visual autoencoder for image generation)
             if hf.ae_model_path is not None:
                 ae_checkpoint = hf.get_load_ckpt(*os.path.split(hf.ae_model_path))
                 ae_params = safetensors.torch.load_file(ae_checkpoint)

eole/config/inference.py

@@ -105,31 +105,12 @@ class ImageGenerationConfig(Config):
         description="Height of the generated image. "
         "This will only work if the model is trained for image generation.",
     )
-    cfg_text_scale: float | None = Field(
-        default=1.0,
-        description="Classifier-free guidance scale for text input. "
-    )
-    cfg_image_scale: float | None = Field(
-        default=1.0,
-        description="Classifier-free guidance scale for image input. "
-    )
-    cfg_interval_min: float | None = Field(
-        default=0.0,
-        description="Minimum classifier-free guidance interval. "
-    )
-    cfg_interval_max: float | None = Field(
-        default=1.0,
-        description="Maximum classifier-free guidance interval. "
-    )
-    timestep_shift: float | None = Field(
-        default=1.0,
-        description="Shift the timestep for image generation. "
-    )
-    num_timesteps: int | None = Field(
-        default=50,
-        description="Number of timesteps for image generation. "
-    )
-
+    cfg_text_scale: float | None = Field(default=1.0, description="Classifier-free guidance scale for text input. ")
+    cfg_image_scale: float | None = Field(default=1.0, description="Classifier-free guidance scale for image input. ")
+    cfg_interval_min: float | None = Field(default=0.0, description="Minimum classifier-free guidance interval. ")
+    cfg_interval_max: float | None = Field(default=1.0, description="Maximum classifier-free guidance interval. ")
+    timestep_shift: float | None = Field(default=1.0, description="Shift the timestep for image generation. ")
+    num_timesteps: int | None = Field(default=50, description="Number of timesteps for image generation. ")
 
 
 # in legacy opts, decoding config is separated (probably to be used elsewhere)

eole/config/models.py

@@ -371,10 +371,17 @@ class VisionEncoderConfig(TransformerConfig, EncoderConfig):
     num_channels: int | None = 3
     image_size: int | None = 1024
     patch_size: int | None = 16
+    max_patches_per_side: int | None = None
+    max_latent_size: int | None = 64  # bagel
+    latent_patch_size: int | None = 2
+    latent_channel: int | None = 16
     image_token_id: int | None = 10  # pixtral uses 10, gemma3 uses 262144
+    image_start_token_id: int | None = None
+    image_end_token_id: int | None = None
     mm_tokens_per_image: int | None = 256  # added for gemma3
     layernorm_pre: bool = True  # True for pixtral/mistral False for gemma3
     patch_conv_bias: bool = False  # False for pixtral/mistral True for gemma3
+    patch_conv_linear: bool = False  # False for pixtral/gemma3 True for bagel
 
 
 # use Field with default= + description would be more readable
@@ -771,6 +778,10 @@ class VisionTransformerLMModelConfig(TransformerConfig, BaseModelConfig):
 
     adapter: str | None = Field(default="llava", description="Adapter type to use in the model.")
 
+    vit_position_embeddings: bool = Field(
+        default=False, description="Additional position embeddings for images, introduced for Bagel."
+    )
+
     @model_validator(mode="before")
     @classmethod
     def encoder_decoder_type(cls, data: Any) -> Any:

eole/decoders/decoder.py

@@ -14,6 +14,10 @@ def __init__(self, attentional=True):
         # Decoder state
         self.state = {}
 
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
     @classmethod
     def from_config(cls, decoder_config, running_config=None, with_cross_attn=False):
         """Alternate constructor.