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

mlx_vlm/models/glm4v/processing.py

@@ -0,0 +1,220 @@
+from typing import List, Union
+
+import numpy as np
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.processing_utils import ProcessorMixin
+
+
+class Glm46VProcessor(ProcessorMixin):
+    """
+    Processor for GLM-4V 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
+
+        # Get image/video tokens from tokenizer or use defaults
+        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|>")
+
+            # Get token IDs
+            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
+
+        # Pop tokenizer-specific kwargs that shouldn't go to image processor
+        padding = kwargs.pop("padding", False)
+        return_token_type_ids = kwargs.pop("return_token_type_ids", False)
+        return_tensors = kwargs.pop("return_tensors", None)
+
+        # Process images
+        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")
+
+        # Process videos
+        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")
+
+        # Handle text input
+        if text is None:
+            text = [""]
+        elif not isinstance(text, list):
+            text = [text]
+
+        # Make a copy to avoid modifying original
+        text = [t for t in text]
+
+        # Get merge_size from image_processor
+        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
+
+        # Expand image tokens based on grid dimensions
+        if image_grid_thw is not None:
+            index = 0
+            for i in range(len(text)):
+                while self.image_token in text[i]:
+                    # Calculate number of image tokens: prod(grid_thw) / merge_size^2
+                    grid = image_grid_thw[index]
+                    if hasattr(grid, "tolist"):
+                        grid = grid.tolist()
+                    num_image_tokens = int(np.prod(grid) // merge_length)
+
+                    # Replace single image token with correct number of placeholder tokens
+                    text[i] = text[i].replace(
+                        self.image_token,
+                        "<|placeholder|>" * num_image_tokens,
+                        1,
+                    )
+                    index += 1
+                # Replace placeholders back to image tokens
+                text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+        # Expand video tokens based on grid dimensions
+        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]
+                    # Calculate tokens per frame
+                    num_tokens_per_frame = int(
+                        np.prod(grid) // merge_length // num_frames
+                    )
+
+                    # Build video structure with frame tokens
+                    video_structure = ""
+                    for frame_idx in range(num_frames):
+                        # Add image tokens for this frame
+                        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
+
+        # Tokenize text
+        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 AutoTokenizer, Glm4vImageProcessor
+
+        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 = Glm4vImageProcessor.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+
+        return cls(image_processor=image_processor, tokenizer=tokenizer, **kwargs)
+
+
+__all__ = ["Glm46VProcessor"]

mlx_vlm/models/glm4v/vision.py

@@ -0,0 +1,406 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..kernels import grid_sample
+from .config import VisionConfig
+
+
+def check_array_shape(arr):
+    shape = arr.shape
+
+    # Check if the shape has 4 or 5 dimensions
+    if len(shape) == 4:
+        out_channels, kH, KW, _ = shape
+        # Check if out_channels is the largest, and kH and KW are the same
+        return (out_channels >= kH) and (out_channels >= KW) and (kH == KW)
+    elif len(shape) == 5:
+        B, out_channels, kH, KW, t = shape
+        # Special case for temporal dimension
+        if t == 3:
+            return True
+        # Check if out_channels is the largest, and kH and KW are the same
+        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(tensor, freqs) -> mx.array:
+    orig_dtype = tensor.dtype
+
+    cos = mx.cos(freqs)
+    sin = mx.sin(freqs)
+
+    cos = mx.expand_dims(cos, axis=1)  # Equivalent to unsqueeze(1)
+    cos = mx.tile(cos, (1, 1, 2))  # Equivalent to repeat(1, 1, 2)
+    cos = mx.expand_dims(cos, axis=0)  # Equivalent to [None, ...]
+
+    sin = mx.expand_dims(sin, axis=1)  # Equivalent to unsqueeze(1)
+    sin = mx.tile(sin, (1, 1, 2))  # Equivalent to repeat(1, 1, 2)
+    sin = mx.expand_dims(sin, axis=0)  # Equivalent to [None, ...]
+
+    output = (tensor * cos) + (rotate_half(tensor) * sin)
+    return output.astype(orig_dtype)
+
+
+class Glm4vVisionRotaryEmbedding(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.item(), dtype=inv_freq.dtype)
+        freqs = mx.outer(seq, inv_freq)
+        return freqs
+
+
+class Glm4vVisionEmbeddings(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+
+    def __call__(self, embeddings, lengths, image_shapes, h_coords, w_coords):
+
+        # Get position embedding parameters
+        pos_embed_weight = self.position_embedding.weight
+        hidden_size = pos_embed_weight.shape[1]
+        total_seq = h_coords.shape[0]
+
+        # Handle empty sequence case
+        if total_seq == 0:
+            adapted_pos_embed = mx.empty(0, hidden_size, dtype=pos_embed_weight.dtype)
+        else:
+            # Convert inputs to tensors if needed
+            if isinstance(lengths, list):
+                lengths = mx.array(lengths, dtype=mx.int32)
+            if not isinstance(image_shapes, mx.array):
+                image_shapes = mx.array(image_shapes, dtype=mx.int32)
+
+            # Prepare 2D position embedding
+            orig_size_sq = pos_embed_weight.shape[0]
+            orig_size = int(orig_size_sq**0.5)
+            pos_embed_2d = (
+                pos_embed_weight.reshape(orig_size, orig_size, hidden_size)
+                .transpose(2, 0, 1)[None, ...]
+                .astype(mx.float32)
+            )
+
+            # Calculate target dimensions for each patch
+            target_h = mx.concatenate(
+                [mx.repeat(image_shapes[i, 1], lengths[i]) for i in range(len(lengths))]
+            ).astype(mx.float32)
+            target_w = mx.concatenate(
+                [mx.repeat(image_shapes[i, 2], lengths[i]) for i in range(len(lengths))]
+            ).astype(mx.float32)
+
+            # Normalize coordinates to [-1, 1] range for grid_sample
+            h_coords = h_coords.astype(mx.float32)
+            w_coords = w_coords.astype(mx.float32)
+            norm_w = ((w_coords + 0.5) / target_w) * 2 - 1
+            norm_h = ((h_coords + 0.5) / target_h) * 2 - 1
+
+            # Create sampling grid
+            grid = mx.stack((norm_w, norm_h), axis=-1)[None, :, None, ...]
+
+            # Perform bicubic interpolation
+            interpolated_embed_fp32 = grid_sample(
+                pos_embed_2d.transpose(0, 2, 3, 1),
+                grid,
+            )
+
+            # Reshape and convert back to original dtype
+            adapted_pos_embed_fp32 = interpolated_embed_fp32.squeeze(0).squeeze(1)
+            adapted_pos_embed = adapted_pos_embed_fp32.astype(pos_embed_weight.dtype)
+
+        # Add adapted position encoding to embeddings
+        embeddings = embeddings + adapted_pos_embed
+        return embeddings
+
+
+class Glm4vVisionPatchEmbed(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,
+        )
+
+    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 Glm4vVisionPatchMerger(nn.Module):
+    def __init__(self, dim: int, context_dim: int, 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 Glm4vVisionAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 16) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.proj = nn.Linear(dim, dim, bias=False)
+
+    def __call__(
+        self, x: mx.array, cu_seqlens: mx.array, rotary_pos_emb: mx.array = None
+    ) -> mx.array:
+        seq_length = x.shape[0]
+        qkv = (
+            self.qkv(x).reshape(seq_length, 3, self.num_heads, -1).transpose(1, 0, 2, 3)
+        )
+        q, k, v = mx.split(qkv, 3)
+
+        q = apply_rotary_pos_emb_vision(mx.expand_dims(q, 0), rotary_pos_emb)[0]
+        k = apply_rotary_pos_emb_vision(mx.expand_dims(k, 0), rotary_pos_emb)[0]
+
+        q = q.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)
+        v = v.transpose(0, 2, 1, 3)
+
+        lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+        splits = [
+            mx.split(tensor, [lengths[0], sum(lengths[:2])], axis=2)
+            for tensor in (q, k, v)
+        ]
+
+        attn_outputs = []
+        for q, k, v in zip(*splits):
+            output = mx.fast.scaled_dot_product_attention(
+                q, k, v, scale=self.scale, mask=None
+            )
+            attn_outputs.append(output)
+
+        output = mx.concatenate(attn_outputs, axis=2)
+        output = output.transpose(0, 2, 1, 3).reshape(seq_length, -1)
+        return self.proj(output)
+
+
+class Glm4vVisionMLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Glm4vVisionBlock(nn.Module):
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.norm1 = nn.RMSNorm(config.hidden_size, eps=1e-6)
+        self.norm2 = nn.RMSNorm(config.hidden_size, eps=1e-6)
+
+        self.attn = Glm4vVisionAttention(
+            dim=config.hidden_size, num_heads=config.num_heads
+        )
+        self.mlp = Glm4vVisionMLP(
+            dim=config.hidden_size, hidden_dim=config.out_hidden_size
+        )
+
+    def __call__(self, hidden_states, cu_seqlens, rotary_pos_emb) -> mx.array:
+        hidden_states = hidden_states + self.attn(
+            self.norm1(hidden_states),
+            cu_seqlens=cu_seqlens,
+            rotary_pos_emb=rotary_pos_emb,
+        )
+        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
+        if self.model_type not in ["glm4v", "glm4v_vision"]:
+            raise ValueError(f"Unsupported model type: {self.model_type}")
+        self.spatial_merge_size = config.spatial_merge_size
+
+        self.embeddings = Glm4vVisionEmbeddings(config)
+        self.patch_embed = Glm4vVisionPatchEmbed(
+            config=config,
+        )
+
+        self.window_size = config.window_size
+        self.patch_size = config.patch_size
+        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
+
+        head_dim = config.hidden_size // config.num_heads
+        self.rotary_pos_emb = Glm4vVisionRotaryEmbedding(head_dim // 2)
+
+        self.blocks = [Glm4vVisionBlock(config) for _ in range(config.depth)]
+        self.merger = Glm4vVisionPatchMerger(
+            dim=config.out_hidden_size, context_dim=config.intermediate_size
+        )
+
+        self.post_conv_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        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,
+        )
+        self.post_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def rot_pos_emb(self, grid_thw):
+        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)
+        rotary_pos_emb = rotary_pos_emb_full[pos_ids]
+
+        return rotary_pos_emb.reshape(pos_ids.shape[0], -1), 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)
+        hidden_states = self.post_conv_layernorm(hidden_states)
+        rotary_pos_emb, image_type_ids = self.rot_pos_emb(grid_thw)
+
+        seq_lens = grid_thw[:, 1] * grid_thw[:, 2]
+        repeats = grid_thw[:, 0]
+        repeated_values = []
+        for i, (seq_len, repeat_count) in enumerate(
+            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)
+        seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+        hidden_states = self.embeddings(
+            hidden_states, seqlens, grid_thw, image_type_ids[:, 0], image_type_ids[:, 1]
+        )
+
+        for blk in self.blocks:
+            hidden_states = blk(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                rotary_pos_emb=rotary_pos_emb,
+            )
+
+        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
+        )
+
+        hidden_states = self.merger(hidden_states)
+        return hidden_states
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for k, v in weights.items():
+            if "position_ids" in k:
+                # Remove unused position_ids
+                continue
+            elif "patch_embed.proj.weight" in k or "downsample.weight" in k:
+                # PyTorch conv2d weight tensors have shape:
+                #   [out_channels, in_channels, kH, KW]
+                # MLX conv2d expects the weight be of shape:
+                #   [out_channels, kH, KW, in_channels]
+                if check_array_shape(v):
+                    sanitized_weights[k] = v
+                else:
+                    if v.ndim == 5:
+                        sanitized_weights[k] = v.transpose(0, 2, 3, 4, 1)
+                    if v.ndim == 4:
+                        sanitized_weights[k] = v.transpose(0, 2, 3, 1)
+            else:
+                sanitized_weights[k] = v
+
+        return sanitized_weights

mlx_vlm/models/glm4v_moe/__init__.py

@@ -0,0 +1,3 @@
+from .config import ModelConfig, TextConfig, VisionConfig
+from .glm4v_moe import LanguageModel, Model, VisionModel
+from .processing import Glm46VMoEProcessor