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

mlx_vlm/models/qwen3_omni_moe/thinker.py

@@ -0,0 +1,366 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from mlx_vlm.models.qwen3_omni_moe.audio import AudioModel
+
+from ..base import InputEmbeddingsFeatures
+from .config import ThinkerConfig
+from .language import LanguageModel
+from .vision import VisionModel
+
+
+def masked_scatter(
+    final_embedding: mx.array,
+    image_mask_expanded: mx.array,
+    scaled_image_features: mx.array,
+):
+    final_embedding_shape = final_embedding.shape
+    scaled_image_features_flattened = mx.flatten(scaled_image_features)
+    final_embedding_flattened = mx.flatten(final_embedding)
+    image_mask_expanded_flattened = mx.flatten(image_mask_expanded)
+
+    image_positions = mx.array(np.where(image_mask_expanded_flattened)[0], mx.uint32)
+    final_embedding_flattened[image_positions] = scaled_image_features_flattened
+
+    final_embedding = mx.reshape(final_embedding_flattened, final_embedding_shape)
+
+    return final_embedding
+
+
+class Thinker(nn.Module):
+    def __init__(self, config: ThinkerConfig):
+        super().__init__()
+        self.config = config
+
+        self.vision_tower = VisionModel(config.vision_config)
+
+        self.audio_tower = AudioModel(config.audio_config)
+
+        self.language_model = LanguageModel(config.text_config, config)
+
+    def get_audio_features(
+        self,
+        input_features: mx.array,
+        feature_attention_mask: Optional[mx.array] = None,
+        audio_feature_lengths: Optional[mx.array] = None,
+    ):
+        if feature_attention_mask is not None:
+            input_features = mx.transpose(input_features, (0, 2, 1))
+            audio_feature_lengths = mx.sum(feature_attention_mask, axis=1)
+            mask_bool = feature_attention_mask.astype(mx.bool_)
+            batch_size, seq_len, hidden_dim = input_features.shape
+            input_features_flat = mx.reshape(input_features, (-1, hidden_dim))
+            mask_flat = mx.reshape(mask_bool, (-1,))
+            indices = mx.array(np.where(mask_flat)[0])
+            input_features = mx.take(input_features_flat, indices, axis=0)
+            input_features = mx.transpose(input_features, (1, 0))
+
+        feature_lens = (
+            audio_feature_lengths
+            if audio_feature_lengths is not None
+            else (
+                mx.sum(feature_attention_mask, axis=-1)
+                if feature_attention_mask is not None
+                else None
+            )
+        )
+        audio_outputs = self.audio_tower(
+            input_features,
+            feature_lens=feature_lens,
+        )
+        return audio_outputs
+
+    def get_placeholder_mask(
+        self,
+        input_ids: Optional[mx.array] = None,
+        inputs_embeds: Optional[mx.array] = None,
+        image_features: Optional[mx.array] = None,
+        video_features: Optional[mx.array] = None,
+        audio_features: Optional[mx.array] = None,
+    ):
+        if input_ids is None:
+            image_token_embed = self.language_model.model.embed_tokens(
+                mx.array([self.config.image_token_id], dtype=mx.int32)
+            )
+            video_token_embed = self.language_model.model.embed_tokens(
+                mx.array([self.config.video_token_id], dtype=mx.int32)
+            )
+            audio_token_embed = self.language_model.model.embed_tokens(
+                mx.array([self.config.audio_token_id], dtype=mx.int32)
+            )
+
+            special_image_mask = mx.all(inputs_embeds == image_token_embed, axis=-1)
+            special_video_mask = mx.all(inputs_embeds == video_token_embed, axis=-1)
+            special_audio_mask = mx.all(inputs_embeds == audio_token_embed, axis=-1)
+        else:
+            special_image_mask = input_ids == self.config.image_token_id
+            special_video_mask = input_ids == self.config.video_token_id
+            special_audio_mask = input_ids == self.config.audio_token_id
+
+        n_image_tokens = mx.sum(special_image_mask)
+        special_image_mask = mx.expand_dims(special_image_mask, axis=-1)
+        special_image_mask = mx.broadcast_to(special_image_mask, inputs_embeds.shape)
+        if (
+            image_features is not None
+            and mx.sum(special_image_mask) != image_features.size
+        ):
+            raise ValueError(
+                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}"
+            )
+
+        n_video_tokens = mx.sum(special_video_mask)
+        special_video_mask = mx.expand_dims(special_video_mask, axis=-1)
+        special_video_mask = mx.broadcast_to(special_video_mask, inputs_embeds.shape)
+        if (
+            video_features is not None
+            and mx.sum(special_video_mask) != video_features.size
+        ):
+            raise ValueError(
+                f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {video_features.shape[0]}"
+            )
+
+        special_audio_mask = mx.expand_dims(special_audio_mask, axis=-1)
+        special_audio_mask = mx.broadcast_to(special_audio_mask, inputs_embeds.shape)
+        return special_image_mask, special_video_mask, special_audio_mask
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        pixel_values_videos: Optional[mx.array] = None,
+        image_grid_thw: Optional[mx.array] = None,
+        video_grid_thw: Optional[mx.array] = None,
+        input_features: Optional[mx.array] = None,
+        feature_attention_mask: Optional[mx.array] = None,
+        audio_feature_lengths: Optional[mx.array] = None,
+        **kwargs,
+    ):
+        inputs_embeds = self.language_model.model.embed_tokens(input_ids)
+        visual_pos_masks = None
+        deepstack_visual_embeds = None
+
+        if input_features is not None:
+            audio_features = self.get_audio_features(
+                input_features,
+                feature_attention_mask=feature_attention_mask,
+                audio_feature_lengths=audio_feature_lengths,
+            )
+            audio_features = audio_features.astype(inputs_embeds.dtype)
+            _, _, audio_mask = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds
+            )
+            if audio_features.ndim == 2:
+                audio_mask_flat = mx.reshape(
+                    audio_mask[..., 0] if audio_mask.ndim == 3 else audio_mask, (-1,)
+                )
+                num_audio_tokens = int(mx.sum(audio_mask_flat))
+                if num_audio_tokens > 0:
+                    if audio_features.shape[0] != num_audio_tokens:
+                        if audio_features.shape[0] > num_audio_tokens:
+                            audio_features = audio_features[:num_audio_tokens]
+                        else:
+                            padding = mx.zeros(
+                                (
+                                    num_audio_tokens - audio_features.shape[0],
+                                    audio_features.shape[1],
+                                ),
+                                dtype=audio_features.dtype,
+                            )
+                            audio_features = mx.concatenate(
+                                [audio_features, padding], axis=0
+                            )
+            inputs_embeds = masked_scatter(inputs_embeds, audio_mask, audio_features)
+
+        if pixel_values is not None:
+            dtype = self.vision_tower.patch_embed.proj.weight.dtype
+            pixel_values = pixel_values.astype(dtype)
+            vision_output = self.vision_tower(pixel_values, image_grid_thw)
+            if isinstance(vision_output, tuple):
+                image_embeds, image_embeds_multiscale = vision_output
+            else:
+                image_embeds = vision_output
+                image_embeds_multiscale = None
+            image_embeds = image_embeds.astype(inputs_embeds.dtype)
+            image_mask, _, _ = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+            )
+            inputs_embeds = masked_scatter(inputs_embeds, image_mask, image_embeds)
+            visual_pos_masks = (
+                image_mask[..., 0] if image_mask.ndim == 3 else image_mask
+            )
+            visual_embeds_multiscale = image_embeds_multiscale
+
+        if pixel_values_videos is not None:
+            dtype = self.vision_tower.patch_embed.proj.weight.dtype
+            pixel_values_videos = pixel_values_videos.astype(dtype)
+            vision_output = self.vision_tower(pixel_values_videos, video_grid_thw)
+            if isinstance(vision_output, tuple):
+                video_embeds, video_embeds_multiscale = vision_output
+            else:
+                video_embeds = vision_output
+                video_embeds_multiscale = None
+            video_embeds = video_embeds.astype(inputs_embeds.dtype)
+            _, video_mask, _ = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+            )
+            inputs_embeds = masked_scatter(inputs_embeds, video_mask, video_embeds)
+
+            if visual_embeds_multiscale is None:
+                visual_embeds_multiscale = video_embeds_multiscale
+                visual_pos_masks = (
+                    video_mask[..., 0] if video_mask.ndim == 3 else video_mask
+                )
+            else:
+                visual_pos_masks = (
+                    (video_mask | image_mask)[..., 0]
+                    if video_mask.ndim == 3
+                    else (video_mask | image_mask)
+                )
+                visual_mask_flat = mx.reshape(visual_pos_masks, (-1,))
+                image_mask_flat = mx.reshape(
+                    image_mask[..., 0] if image_mask.ndim == 3 else image_mask, (-1,)
+                )
+                video_mask_flat = mx.reshape(
+                    video_mask[..., 0] if video_mask.ndim == 3 else video_mask, (-1,)
+                )
+                visual_indices_flat = mx.where(visual_mask_flat)[0]
+                image_mask_on_visual = mx.take(
+                    image_mask_flat, visual_indices_flat, axis=0
+                )
+                video_mask_on_visual = mx.take(
+                    video_mask_flat, visual_indices_flat, axis=0
+                )
+                image_indices = mx.where(image_mask_on_visual)[0]
+                video_indices = mx.where(video_mask_on_visual)[0]
+
+                visual_embeds_multiscale_joint = []
+                for img_embed, vid_embed in zip(
+                    visual_embeds_multiscale, video_embeds_multiscale
+                ):
+                    embed_joint = mx.zeros(
+                        (len(visual_indices_flat), img_embed.shape[-1]),
+                        dtype=img_embed.dtype,
+                    )
+                    if len(image_indices) > 0:
+                        embed_joint = mx.scatter(
+                            embed_joint,
+                            image_indices,
+                            mx.take(img_embed, image_indices, axis=0),
+                            axis=0,
+                        )
+                    if len(video_indices) > 0:
+                        embed_joint = mx.scatter(
+                            embed_joint,
+                            video_indices,
+                            mx.take(vid_embed, video_indices, axis=0),
+                            axis=0,
+                        )
+                    visual_embeds_multiscale_joint.append(embed_joint)
+                visual_embeds_multiscale = tuple(visual_embeds_multiscale_joint)
+
+        return InputEmbeddingsFeatures(
+            inputs_embeds=inputs_embeds,
+            visual_pos_masks=visual_pos_masks,
+            deepstack_visual_embeds=deepstack_visual_embeds,
+        )
+
+    @staticmethod
+    def merge_input_ids_with_image_features(
+        image_features, inputs_embeds, input_ids, image_token_index, video_token_index
+    ):
+        special_image_mask = input_ids == image_token_index
+        special_video_mask = input_ids == video_token_index
+        special_image_mask = special_image_mask | special_video_mask
+        n_image_tokens = special_image_mask.sum()
+        special_image_mask = special_image_mask[..., None]
+        special_image_mask = mx.broadcast_to(special_image_mask, inputs_embeds.shape)
+
+        n_image_features = image_features.shape[0]
+        n_image_mask_elements = special_image_mask.sum()
+        if n_image_mask_elements != image_features.size:
+            raise ValueError(
+                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+            )
+
+        inputs_embeds = masked_scatter(
+            inputs_embeds, special_image_mask, image_features
+        )
+
+        return inputs_embeds, special_image_mask
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: Optional[mx.array] = None,
+        pixel_values_videos: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache=None,
+        **kwargs,
+    ):
+        image_grid_thw = kwargs.pop("image_grid_thw", None)
+        video_grid_thw = kwargs.pop("video_grid_thw", None)
+        input_features = kwargs.pop("input_features", None)
+        feature_attention_mask = kwargs.pop("feature_attention_mask", None)
+        audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
+
+        inputs_embeds, visual_pos_masks, deepstack_visual_embeds = (
+            self.get_input_embeddings(
+                input_ids,
+                pixel_values=pixel_values,
+                pixel_values_videos=pixel_values_videos,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                input_features=input_features,
+                feature_attention_mask=feature_attention_mask,
+                audio_feature_lengths=audio_feature_lengths,
+            )
+        )
+
+        kwargs = {
+            "pixel_values": pixel_values,
+            "pixel_values_videos": pixel_values_videos,
+            "image_grid_thw": image_grid_thw,
+            "video_grid_thw": video_grid_thw,
+            "visual_pos_masks": visual_pos_masks,
+            "deepstack_visual_embeds": deepstack_visual_embeds,
+            **kwargs,
+        }
+
+        logits = self.language_model(
+            input_ids, inputs_embeds, mask=mask, cache=cache, **kwargs
+        )
+
+        return logits
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for key, value in weights.items():
+            if "thinker" in key:
+                if "thinker.model" in key:
+                    key = key.replace("thinker.model", "thinker.language_model.model")
+                elif "thinker.visual" in key:
+                    key = key.replace("thinker.visual", "thinker.vision_tower")
+                    if "merger" in key:
+                        key = (
+                            key.replace("ln_q", "norm")
+                            .replace("mlp.0", "linear_fc1")
+                            .replace("mlp.2", "linear_fc2")
+                        )
+                    if "merger_list" in key:
+                        key = key.replace("merger_list", "deepstack_merger_list")
+
+                if "thinker.lm_head" in key:
+                    key = key.replace(
+                        "thinker.lm_head", "thinker.language_model.lm_head"
+                    )
+
+            sanitized_weights[key] = value
+
+        return sanitized_weights