fount-vlm-nell-02 0.3.11__py3-none-any.whl

mlx_vlm/models/glm_ocr/processing.py

@@ -0,0 +1,208 @@
+from typing import List, Union
+
+import numpy as np
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.processing_utils import ProcessorMixin
+
+from ..base import install_auto_processor_patch
+
+
+class GlmOcrProcessor(ProcessorMixin):
+    """
+    Processor for GLM-OCR that wraps an image processor and tokenizer.
+
+    Handles:
+    - Image preprocessing via image_processor
+    - Token replacement for image/video placeholders based on grid dimensions
+    """
+
+    attributes = ["image_processor", "tokenizer"]
+    valid_kwargs = ["chat_template"]
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        chat_template=None,
+        **kwargs,
+    ):
+        self.tokenizer = tokenizer
+        self.image_processor = image_processor
+
+        self.image_token = "<|image|>"
+        self.video_token = "<|video|>"
+
+        if tokenizer is not None:
+            self.image_token = getattr(tokenizer, "image_token", "<|image|>")
+            self.video_token = getattr(tokenizer, "video_token", "<|video|>")
+
+            self.image_token_id = getattr(tokenizer, "image_token_id", None)
+            if self.image_token_id is None:
+                self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
+
+            self.video_token_id = getattr(tokenizer, "video_token_id", None)
+            if self.video_token_id is None:
+                self.video_token_id = tokenizer.convert_tokens_to_ids(self.video_token)
+        else:
+            self.image_token_id = None
+            self.video_token_id = None
+
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    def __call__(
+        self,
+        images=None,
+        text: Union[str, List[str]] = None,
+        videos=None,
+        **kwargs,
+    ) -> BatchFeature:
+        """
+        Process images/videos and text for the model.
+
+        Args:
+            images: Single image or list of images (PIL.Image, np.ndarray, etc.)
+            text: Single text or list of texts
+            videos: Video inputs (optional)
+            **kwargs: Additional arguments passed to image_processor and tokenizer
+
+        Returns:
+            BatchFeature with:
+                - input_ids: Token IDs with image/video placeholders expanded
+                - attention_mask: Attention mask
+                - pixel_values: Processed image/video patches
+                - image_grid_thw: Grid dimensions for each image
+                - video_grid_thw: Grid dimensions for each video (if videos provided)
+        """
+        image_inputs = {}
+        video_inputs = {}
+        image_grid_thw = None
+        video_grid_thw = None
+
+        padding = kwargs.pop("padding", False)
+        return_token_type_ids = kwargs.pop("return_token_type_ids", False)
+        return_tensors = kwargs.pop("return_tensors", None)
+
+        if images is not None and self.image_processor is not None:
+            image_inputs = self.image_processor(images=images)
+            image_grid_thw = image_inputs.get("image_grid_thw")
+
+        if videos is not None:
+            if hasattr(self, "video_processor") and self.video_processor is not None:
+                video_inputs = self.video_processor(videos=videos, **kwargs)
+                video_grid_thw = video_inputs.get("video_grid_thw")
+
+        if text is None:
+            text = [""]
+        elif not isinstance(text, list):
+            text = [text]
+
+        text = [t for t in text]
+
+        merge_size = getattr(self.image_processor, "merge_size", 2)
+        if hasattr(self.image_processor, "spatial_merge_size"):
+            merge_size = self.image_processor.spatial_merge_size
+        merge_length = merge_size**2
+
+        if image_grid_thw is not None:
+            index = 0
+            for i in range(len(text)):
+                while self.image_token in text[i]:
+                    grid = image_grid_thw[index]
+                    if hasattr(grid, "tolist"):
+                        grid = grid.tolist()
+                    num_image_tokens = int(np.prod(grid) // merge_length)
+
+                    text[i] = text[i].replace(
+                        self.image_token,
+                        "<|placeholder|>" * num_image_tokens,
+                        1,
+                    )
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+        if video_grid_thw is not None:
+            video_index = 0
+            for i in range(len(text)):
+                while self.video_token in text[i]:
+                    grid = video_grid_thw[video_index]
+                    if hasattr(grid, "tolist"):
+                        grid = grid.tolist()
+
+                    num_frames = grid[0]
+                    num_tokens_per_frame = int(
+                        np.prod(grid) // merge_length // num_frames
+                    )
+
+                    video_structure = ""
+                    for frame_idx in range(num_frames):
+                        frame_structure = self.image_token * num_tokens_per_frame
+                        video_structure += frame_structure
+
+                    text[i] = text[i].replace(self.video_token, video_structure, 1)
+                    video_index += 1
+
+        text_inputs = self.tokenizer(
+            text,
+            padding=padding,
+            return_token_type_ids=return_token_type_ids,
+            **kwargs,
+        )
+
+        return BatchFeature(
+            data={**text_inputs, **image_inputs, **video_inputs},
+            tensor_type=return_tensors,
+        )
+
+    def batch_decode(self, *args, **kwargs):
+        """Decode token IDs to text."""
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """Decode token IDs to text."""
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def apply_chat_template(self, *args, **kwargs):
+        """Apply chat template using the tokenizer."""
+        return self.tokenizer.apply_chat_template(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        """Return combined input names from tokenizer and image processor."""
+        tokenizer_input_names = (
+            self.tokenizer.model_input_names if self.tokenizer else []
+        )
+        image_processor_input_names = (
+            self.image_processor.model_input_names
+            if hasattr(self.image_processor, "model_input_names")
+            else []
+        )
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """Load processor from pretrained model path."""
+        from transformers import AutoImageProcessor, AutoTokenizer
+
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+
+        image_processor = AutoImageProcessor.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+
+        return cls(image_processor=image_processor, tokenizer=tokenizer, **kwargs)
+
+
+__all__ = ["GlmOcrProcessor"]
+
+# Register the processor with AutoProcessor for the glm_ocr model type
+install_auto_processor_patch("glm_ocr", GlmOcrProcessor)

mlx_vlm/models/glm_ocr/vision.py

@@ -0,0 +1,342 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .config import VisionConfig
+
+
+def check_array_shape(arr):
+    shape = arr.shape
+
+    if len(shape) == 4:
+        out_channels, kH, KW, _ = shape
+        return (out_channels >= kH) and (out_channels >= KW) and (kH == KW)
+    elif len(shape) == 5:
+        out_channels, kT, kH, KW, _ = shape
+        return (out_channels >= kH) and (out_channels >= KW) and (kH == KW)
+    else:
+        return False
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return mx.concatenate([-x2, x1], axis=-1)
+
+
+def apply_rotary_pos_emb_vision(
+    q: mx.array, k: mx.array, cos: mx.array, sin: mx.array
+) -> tuple:
+    orig_q_dtype = q.dtype
+    orig_k_dtype = k.dtype
+    q, k = q.astype(mx.float32), k.astype(mx.float32)
+    cos = mx.expand_dims(cos, axis=-2).astype(mx.float32)
+    sin = mx.expand_dims(sin, axis=-2).astype(mx.float32)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    q_embed = q_embed.astype(orig_q_dtype)
+    k_embed = k_embed.astype(orig_k_dtype)
+    return q_embed, k_embed
+
+
+class GlmOcrVisionRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+
+    def __call__(self, seqlen: int) -> mx.array:
+        inv_freq = 1.0 / (
+            self.theta ** (mx.arange(0, self.dim, 2, dtype=mx.float32) / self.dim)
+        )
+        seq = mx.arange(seqlen, dtype=inv_freq.dtype)
+        freqs = mx.outer(seq, inv_freq)
+        return freqs
+
+
+class GlmOcrVisionPatchEmbed(nn.Module):
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.patch_size = config.patch_size
+        self.temporal_patch_size = config.temporal_patch_size
+        self.in_channels = config.in_channels
+        self.embed_dim = config.hidden_size
+
+        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
+        self.proj = nn.Conv3d(
+            self.in_channels,
+            self.embed_dim,
+            kernel_size=kernel_size,
+            stride=kernel_size,
+            bias=True,
+        )
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        hidden_states = hidden_states.reshape(
+            -1,
+            self.in_channels,
+            self.temporal_patch_size,
+            self.patch_size,
+            self.patch_size,
+        ).moveaxis(1, 4)
+
+        hidden_states = self.proj(hidden_states)
+        hidden_states = hidden_states.reshape(-1, self.embed_dim)
+        return hidden_states
+
+
+class GlmOcrVisionPatchMerger(nn.Module):
+    def __init__(
+        self, dim: int, context_dim: int, hidden_act: str, bias: bool = False
+    ) -> None:
+        super().__init__()
+        self.proj = nn.Linear(dim, dim, bias=bias)
+        self.post_projection_norm = nn.LayerNorm(dim)
+        self.gate_proj = nn.Linear(dim, context_dim, bias=bias)
+        self.up_proj = nn.Linear(dim, context_dim, bias=bias)
+        self.down_proj = nn.Linear(context_dim, dim, bias=bias)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = self.proj(x)
+        x = nn.gelu(self.post_projection_norm(x))
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class GlmOcrVisionAttention(nn.Module):
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.dim = config.hidden_size
+        self.num_heads = config.num_heads
+        self.head_dim = self.dim // self.num_heads
+        self.scale = self.head_dim**-0.5
+
+        self.qkv = nn.Linear(
+            config.hidden_size, config.hidden_size * 3, bias=config.attention_bias
+        )
+        self.proj = nn.Linear(
+            config.hidden_size, config.hidden_size, bias=config.attention_bias
+        )
+
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        cu_seqlens: mx.array,
+        position_embeddings: tuple,
+    ) -> mx.array:
+        seq_length = hidden_states.shape[0]
+
+        qkv = self.qkv(hidden_states)
+        qkv = qkv.reshape(seq_length, 3, self.num_heads, -1)
+        qkv = qkv.transpose(1, 0, 2, 3)
+        q, k, v = mx.split(qkv, 3, axis=0)
+        q = q.squeeze(0)
+        k = k.squeeze(0)
+        v = v.squeeze(0)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        cos, sin = position_embeddings
+        q, k = apply_rotary_pos_emb_vision(q, k, cos, sin)
+
+        q = q.transpose(1, 0, 2)[None, ...]
+        k = k.transpose(1, 0, 2)[None, ...]
+        v = v.transpose(1, 0, 2)[None, ...]
+
+        lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+        split_indices = []
+        cumsum = 0
+        for i, length in enumerate(lengths[:-1]):
+            cumsum += length
+            split_indices.append(cumsum)
+
+        q_splits = mx.split(q, split_indices, axis=2)
+        k_splits = mx.split(k, split_indices, axis=2)
+        v_splits = mx.split(v, split_indices, axis=2)
+
+        attn_outputs = []
+        for q_chunk, k_chunk, v_chunk in zip(q_splits, k_splits, v_splits):
+            output = mx.fast.scaled_dot_product_attention(
+                q_chunk, k_chunk, v_chunk, scale=self.scale, mask=None
+            )
+            attn_outputs.append(output)
+
+        attn_output = mx.concatenate(attn_outputs, axis=2)
+        # Transpose from (batch, heads, seq, head_dim) to (batch, seq, heads, head_dim)
+        # then reshape to (seq, hidden_size)
+        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(seq_length, -1)
+        attn_output = self.proj(attn_output)
+        return attn_output
+
+
+class GlmOcrVisionMLP(nn.Module):
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=config.attention_bias
+        )
+        self.up_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=config.attention_bias
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size, config.hidden_size, bias=config.attention_bias
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class GlmOcrVisionBlock(nn.Module):
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.norm1 = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm2 = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attn = GlmOcrVisionAttention(config)
+        self.mlp = GlmOcrVisionMLP(config)
+
+    def __call__(
+        self, hidden_states: mx.array, cu_seqlens: mx.array, position_embeddings: tuple
+    ) -> mx.array:
+        hidden_states = hidden_states + self.attn(
+            self.norm1(hidden_states),
+            cu_seqlens=cu_seqlens,
+            position_embeddings=position_embeddings,
+        )
+        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+        return hidden_states
+
+
+class VisionModel(nn.Module):
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+        self.spatial_merge_size = config.spatial_merge_size
+        self.patch_size = config.patch_size
+
+        self.patch_embed = GlmOcrVisionPatchEmbed(config)
+
+        head_dim = config.hidden_size // config.num_heads
+        self.rotary_pos_emb = GlmOcrVisionRotaryEmbedding(head_dim // 2)
+
+        self.blocks = [GlmOcrVisionBlock(config) for _ in range(config.depth)]
+
+        self.merger = GlmOcrVisionPatchMerger(
+            dim=config.out_hidden_size,
+            context_dim=config.out_hidden_size * config.in_channels,
+            hidden_act=config.hidden_act,
+        )
+
+        self.downsample = nn.Conv2d(
+            in_channels=config.hidden_size,
+            out_channels=config.out_hidden_size,
+            kernel_size=config.spatial_merge_size,
+            stride=config.spatial_merge_size,
+            bias=True,
+        )
+
+        self.post_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def rot_pos_emb(self, grid_thw: mx.array):
+        pos_ids = []
+
+        for t, h, w in grid_thw.tolist():
+            hpos_ids = mx.expand_dims(mx.arange(h), 1)
+            hpos_ids = mx.repeat(hpos_ids, w, axis=1)
+            hpos_ids = hpos_ids.reshape(
+                h // self.spatial_merge_size,
+                self.spatial_merge_size,
+                w // self.spatial_merge_size,
+                self.spatial_merge_size,
+            )
+            hpos_ids = mx.transpose(hpos_ids, (0, 2, 1, 3))
+            hpos_ids = hpos_ids.flatten()
+
+            wpos_ids = mx.expand_dims(mx.arange(w), 0)
+            wpos_ids = mx.repeat(wpos_ids, h, axis=0)
+            wpos_ids = wpos_ids.reshape(
+                h // self.spatial_merge_size,
+                self.spatial_merge_size,
+                w // self.spatial_merge_size,
+                self.spatial_merge_size,
+            )
+            wpos_ids = mx.transpose(wpos_ids, (0, 2, 1, 3))
+            wpos_ids = wpos_ids.flatten()
+
+            stacked_pos_ids = mx.stack([hpos_ids, wpos_ids], axis=-1)
+            pos_ids.append(mx.tile(stacked_pos_ids, (t, 1)))
+
+        pos_ids = mx.concatenate(pos_ids, axis=0)
+        max_grid_size = mx.max(grid_thw[:, 1:])
+        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size.item())
+        rotary_pos_emb = rotary_pos_emb_full[pos_ids].reshape(pos_ids.shape[0], -1)
+
+        emb = mx.concatenate((rotary_pos_emb, rotary_pos_emb), axis=-1)
+        return (mx.cos(emb), mx.sin(emb)), pos_ids
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        grid_thw: mx.array,
+        output_hidden_states: Optional[bool] = None,
+    ) -> mx.array:
+        hidden_states = self.patch_embed(hidden_states)
+        position_embeddings, _ = self.rot_pos_emb(grid_thw)
+
+        seq_lens = grid_thw[:, 1] * grid_thw[:, 2]
+        repeats = grid_thw[:, 0]
+        repeated_values = []
+        for seq_len, repeat_count in zip(seq_lens.tolist(), repeats.tolist()):
+            repeated_values.extend([seq_len] * repeat_count)
+
+        cu_seqlens = mx.array(repeated_values).cumsum(axis=0)
+        cu_seqlens = mx.pad(cu_seqlens, (1, 0), constant_values=0)
+
+        for blk in self.blocks:
+            hidden_states = blk(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                position_embeddings=position_embeddings,
+            )
+
+        hidden_states = self.post_layernorm(hidden_states)
+
+        hidden_states = hidden_states.reshape(
+            -1,
+            self.spatial_merge_size,
+            self.spatial_merge_size,
+            hidden_states.shape[-1],
+        )
+        hidden_states = self.downsample(hidden_states).reshape(
+            -1, self.config.out_hidden_size
+        )
+
+        merged_hidden_states = self.merger(hidden_states)
+        return merged_hidden_states
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for k, v in weights.items():
+            if "position_ids" in k:
+                continue
+            elif "patch_embed.proj.weight" in k or "downsample.weight" in k:
+                if check_array_shape(v):
+                    sanitized_weights[k] = v
+                else:
+                    if v.ndim == 5:
+                        sanitized_weights[k] = v.transpose(0, 2, 3, 4, 1)
+                    elif v.ndim == 4:
+                        sanitized_weights[k] = v.transpose(0, 2, 3, 1)
+                    else:
+                        sanitized_weights[k] = v
+            else:
+                sanitized_weights[k] = v
+
+        return sanitized_weights

mlx_vlm/models/hunyuan_vl/__init__.py

@@ -0,0 +1,7 @@
+from .config import ModelConfig, TextConfig, VisionConfig
+from .hunyuan_vl import LanguageModel, Model, VisionModel
+from .processing_hunyuan_vl import (
+    HunYuanVLImageProcessor,
+    HunYuanVLProcessor,
+    ImageProcessor,
+)

mlx_vlm/models/hunyuan_vl/config.py

@@ -0,0 +1,136 @@
+import inspect
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+from ..base import BaseModelConfig
+
+
+@dataclass
+class VisionConfig(BaseModelConfig):
+    model_type: str = "hunyuan_vl"
+    hidden_size: int = 1152
+    out_hidden_size: int = 1024
+    num_hidden_layers: int = 27
+    num_attention_heads: int = 16
+    intermediate_size: int = 4304
+    patch_size: int = 16
+    num_channels: int = 3
+    spatial_merge_size: int = 2
+    attention_dropout: float = 0.0
+    hidden_dropout: float = 0.0
+    rms_norm_eps: float = 1e-5
+    interpolate_mode: str = "bilinear"
+    cat_extra_token: int = 1
+    img_max_token_num: int = 4096
+    max_vit_seq_len: int = 16384
+    add_patchemb_bias: bool = True
+    max_image_size: int = 2048
+    hidden_act: str = "gelu"
+
+
+@dataclass
+class TextConfig(BaseModelConfig):
+    model_type: str = "hunyuan_vl"
+    vocab_size: int = 120818
+    org_vocab_size: int = 120818
+    hidden_size: int = 1024
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 16
+    num_key_value_heads: Optional[int] = 8
+    head_dim: Optional[int] = 128
+    attention_head_dim: Optional[int] = 128
+    intermediate_size: int = 3584
+    hidden_act: str = "silu"
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    attention_dropout: float = 0.0
+    use_qk_norm: bool = True
+    rope_theta: float = 10000.0
+    rope_scaling: Optional[Dict[str, Union[float, int, bool, List[int]]]] = field(
+        default_factory=lambda: {
+            "alpha": 1000.0,
+            "beta_fast": 32,
+            "beta_slow": 1,
+            "factor": 1.0,
+            "mscale": 1.0,
+            "mscale_all_dim": 1.0,
+            "type": "xdrope",
+            "xdrope_section": [16, 16, 16, 16],
+        }
+    )
+    max_position_embeddings: int = 32768
+    rms_norm_eps: float = 1e-5
+    norm_type: str = "rms"
+    tie_word_embeddings: bool = True
+    use_cache: bool = True
+    initializer_range: float = 0.02
+    routed_scaling_factor: float = 1.0
+    dtype: str = "bfloat16"
+    bos_token_id: int = 120000
+    eos_token_id: int = 120020
+    eod_token_id: int = 120020
+    pad_token_id: int = -1
+    pad_id: int = 120002
+    sep_token_id: int = 0
+    text_start_id: int = 7
+    text_end_id: int = 8
+    num_experts: int = 1
+    pretraining_tp: int = 1
+    use_cla: bool = False
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+
+        if self.attention_head_dim is None:
+            self.attention_head_dim = self.head_dim
+
+
+@dataclass
+class ModelConfig(BaseModelConfig):
+    text_config: TextConfig = field(default_factory=TextConfig)
+    vision_config: VisionConfig = field(default_factory=VisionConfig)
+    model_type: str = "hunyuan_vl"
+    image_start_token_id: int = 120118
+    image_end_token_id: int = 120119
+    image_token_id: int = 120120
+    image_newline_token_id: int = 120121
+    bos_token_id: int = 120000
+    eos_token_id: int = 120020
+    pad_token_id: int = -1
+    pad_id: int = 120002
+    sep_token_id: int = 0
+    text_start_id: int = 7
+    text_end_id: int = 8
+    vocab_size: int = 120818
+    org_vocab_size: int = 120818
+    routed_scaling_factor: float = 1.0
+    norm_type: str = "rms"
+    dtype: str = "bfloat16"
+    use_cache: bool = True
+    tie_word_embeddings: bool = True
+
+    @classmethod
+    def from_dict(cls, params):
+        text_params = params.get("text_config", {})
+        vision_params = params.get("vision_config", {})
+
+        for key, value in params.items():
+            if key in TextConfig.__dataclass_fields__ and key not in text_params:
+                text_params[key] = value
+            if key in VisionConfig.__dataclass_fields__ and key not in vision_params:
+                vision_params[key] = value
+
+        return cls(
+            text_config=TextConfig.from_dict(text_params),
+            vision_config=VisionConfig.from_dict(vision_params),
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+                and k not in ["text_config", "vision_config"]
+            },
+        )