| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from transformers import PreTrainedConfig, SiglipVisionConfig |
| 4 | from transformers.models.qwen2 import Qwen2Config |
| 5 | |
| 6 | |
| 7 | class BagelConfig(PreTrainedConfig): |
| 8 | """Configuration class for BAGEL model.""" |
| 9 | |
| 10 | model_type = "bagel" |
| 11 | |
| 12 | def __init__( |
| 13 | self, |
| 14 | visual_gen: bool = True, |
| 15 | visual_und: bool = True, |
| 16 | llm_config: dict | Qwen2Config | None = None, |
| 17 | vit_config: dict | SiglipVisionConfig | None = None, |
| 18 | vae_config: dict | None = None, |
| 19 | latent_patch_size: int = 2, |
| 20 | max_latent_size: int = 32, |
| 21 | vit_max_num_patch_per_side: int = 70, |
| 22 | connector_act: str = "gelu_pytorch_tanh", |
| 23 | interpolate_pos: bool = False, |
| 24 | timestep_shift: float = 1.0, |
| 25 | **kwargs, |
| 26 | ): |
| 27 | super().__init__(**kwargs) |
| 28 | self.visual_gen = visual_gen |
| 29 | self.visual_und = visual_und |
| 30 | |
| 31 | # Convert dict configs to proper config objects |
| 32 | if isinstance(llm_config, dict): |
| 33 | self.llm_config = Qwen2Config(**llm_config) |
| 34 | else: |
| 35 | self.llm_config = llm_config or Qwen2Config() |
| 36 | |
| 37 | if isinstance(vit_config, dict): |
| 38 | self.vit_config = SiglipVisionConfig(**vit_config) |
| 39 | else: |
| 40 | self.vit_config = vit_config or SiglipVisionConfig() |
| 41 | |
| 42 | self.vae_config = vae_config or {"z_channels": 16, "downsample": 8} |
| 43 | self.latent_patch_size = latent_patch_size |
| 44 | self.max_latent_size = max_latent_size |
| 45 | self.vit_max_num_patch_per_side = vit_max_num_patch_per_side |
| 46 | self.connector_act = connector_act |
| 47 | self.interpolate_pos = interpolate_pos |
| 48 | self.timestep_shift = timestep_shift |
| 49 | |
| 50 | @property |
| 51 | def hidden_size(self) -> int: |
| 52 | """Return the hidden size of the language model.""" |
| 53 | return self.llm_config.hidden_size |