nexaai 1.0.19rc7__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc9__cp310-cp310-macosx_14_0_universal2.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of nexaai might be problematic. Click here for more details.
- nexaai/_stub.cpython-310-darwin.so +0 -0
- nexaai/_version.py +1 -1
- nexaai/binds/libnexa_bridge.dylib +0 -0
- nexaai/mlx_backend/vlm/generate_qwen3_vl.py +14 -31
- nexaai/mlx_backend/vlm/generate_qwen3_vl_moe.py +15 -32
- nexaai/mlx_backend/vlm/modeling/models/qwen3_vl/qwen3vl.py +7 -23
- nexaai/mlx_backend/vlm/modeling/models/qwen3vl_moe/qwen3vl_moe.py +8 -24
- {nexaai-1.0.19rc7.dist-info → nexaai-1.0.19rc9.dist-info}/METADATA +1 -1
- {nexaai-1.0.19rc7.dist-info → nexaai-1.0.19rc9.dist-info}/RECORD +11 -200
- nexaai/binds/nexa_mlx/py-lib/asr/__init__.py +0 -12
- nexaai/binds/nexa_mlx/py-lib/asr/interface.py +0 -122
- nexaai/binds/nexa_mlx/py-lib/common/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/common/utils.py +0 -25
- nexaai/binds/nexa_mlx/py-lib/cv/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/cv/generate.py +0 -195
- nexaai/binds/nexa_mlx/py-lib/cv/interface.py +0 -151
- nexaai/binds/nexa_mlx/py-lib/cv/main.py +0 -81
- nexaai/binds/nexa_mlx/py-lib/cv/modeling/pp_ocr_v4.py +0 -1736
- nexaai/binds/nexa_mlx/py-lib/embedding/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/embedding/generate.py +0 -333
- nexaai/binds/nexa_mlx/py-lib/embedding/interface.py +0 -617
- nexaai/binds/nexa_mlx/py-lib/embedding/main.py +0 -173
- nexaai/binds/nexa_mlx/py-lib/embedding/modeling/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/embedding/modeling/nexa_jina_v2.py +0 -399
- nexaai/binds/nexa_mlx/py-lib/image_gen/__init__.py +0 -1
- nexaai/binds/nexa_mlx/py-lib/image_gen/generate_sd.py +0 -244
- nexaai/binds/nexa_mlx/py-lib/image_gen/interface.py +0 -82
- nexaai/binds/nexa_mlx/py-lib/image_gen/main.py +0 -281
- nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/__init__.py +0 -306
- nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/clip.py +0 -116
- nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/config.py +0 -65
- nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/model_io.py +0 -386
- nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/sampler.py +0 -105
- nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/tokenizer.py +0 -100
- nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/unet.py +0 -460
- nexaai/binds/nexa_mlx/py-lib/image_gen/stable_diffusion/vae.py +0 -274
- nexaai/binds/nexa_mlx/py-lib/llm/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/llm/generate.py +0 -149
- nexaai/binds/nexa_mlx/py-lib/llm/interface.py +0 -764
- nexaai/binds/nexa_mlx/py-lib/llm/main.py +0 -68
- nexaai/binds/nexa_mlx/py-lib/rerank/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/rerank/generate.py +0 -174
- nexaai/binds/nexa_mlx/py-lib/rerank/interface.py +0 -287
- nexaai/binds/nexa_mlx/py-lib/rerank/main.py +0 -127
- nexaai/binds/nexa_mlx/py-lib/rerank/modeling/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/rerank/modeling/nexa_jina_rerank.py +0 -330
- nexaai/binds/nexa_mlx/py-lib/sd/__init__.py +0 -1
- nexaai/binds/nexa_mlx/py-lib/sd/interface.py +0 -362
- nexaai/binds/nexa_mlx/py-lib/sd/main.py +0 -286
- nexaai/binds/nexa_mlx/py-lib/sd/modeling/__init__.py +0 -306
- nexaai/binds/nexa_mlx/py-lib/sd/modeling/clip.py +0 -116
- nexaai/binds/nexa_mlx/py-lib/sd/modeling/config.py +0 -65
- nexaai/binds/nexa_mlx/py-lib/sd/modeling/model_io.py +0 -385
- nexaai/binds/nexa_mlx/py-lib/sd/modeling/sampler.py +0 -105
- nexaai/binds/nexa_mlx/py-lib/sd/modeling/tokenizer.py +0 -100
- nexaai/binds/nexa_mlx/py-lib/sd/modeling/unet.py +0 -460
- nexaai/binds/nexa_mlx/py-lib/sd/modeling/vae.py +0 -274
- nexaai/binds/nexa_mlx/py-lib/tts/__init__.py +0 -12
- nexaai/binds/nexa_mlx/py-lib/tts/interface.py +0 -276
- nexaai/binds/nexa_mlx/py-lib/vlm/__init__.py +0 -3
- nexaai/binds/nexa_mlx/py-lib/vlm/generate.py +0 -572
- nexaai/binds/nexa_mlx/py-lib/vlm/generate_qwen3_vl.py +0 -294
- nexaai/binds/nexa_mlx/py-lib/vlm/generate_qwen3_vl_moe.py +0 -276
- nexaai/binds/nexa_mlx/py-lib/vlm/interface.py +0 -504
- nexaai/binds/nexa_mlx/py-lib/vlm/main.py +0 -320
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/convert.py +0 -68
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/aya_vision.py +0 -193
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/interpolate.py +0 -186
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/language.py +0 -233
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/aya_vision/vision.py +0 -503
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/base.py +0 -202
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/cache.py +0 -230
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/__init__.py +0 -10
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/conversation.py +0 -264
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/deepseek_vl_v2.py +0 -472
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/language.py +0 -591
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/processing_deepsek_vl_v2.py +0 -526
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/deepseek_vl_v2/vision.py +0 -356
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/florence2.py +0 -366
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/language.py +0 -488
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/vision.py +0 -591
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/gemma3.py +0 -213
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/language.py +0 -315
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3/vision.py +0 -238
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/__init__.py +0 -2
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/audio.py +0 -1038
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/config.py +0 -139
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/gemma3n.py +0 -322
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/language.py +0 -629
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/vision.py +0 -1022
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/__init__.py +0 -9
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/idefics2.py +0 -294
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/language.py +0 -191
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics2/vision.py +0 -267
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/idefics3.py +0 -175
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/language.py +0 -192
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/idefics3/vision.py +0 -233
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/__init__.py +0 -9
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/internvl_chat.py +0 -140
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/language.py +0 -220
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/processor.py +0 -393
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/internvl_chat/vision.py +0 -293
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kernels.py +0 -307
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/kimi_vl.py +0 -143
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/language.py +0 -509
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/kimi_vl/vision.py +0 -522
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/language.py +0 -386
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/llama4.py +0 -138
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llama4/vision.py +0 -560
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/language.py +0 -240
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/llava.py +0 -153
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/vision.py +0 -259
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/__init__.py +0 -9
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/language.py +0 -236
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/llava_bunny.py +0 -256
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_bunny/vision.py +0 -303
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/language.py +0 -230
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/llava_next.py +0 -160
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava_next/vision.py +0 -243
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mistral3/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mistral3/mistral3.py +0 -283
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/language.py +0 -416
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/mllama.py +0 -172
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/mllama/vision.py +0 -499
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/language.py +0 -243
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/molmo.py +0 -133
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/vision.py +0 -465
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/__init__.py +0 -10
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/language.py +0 -230
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/multi_modality.py +0 -385
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/sam.py +0 -557
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/multi_modality/vision.py +0 -526
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/language.py +0 -282
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/paligemma.py +0 -160
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/paligemma/vision.py +0 -242
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/language.py +0 -21
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/phi3_v.py +0 -243
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/su_rope.py +0 -71
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/phi3_v/vision.py +0 -324
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/language.py +0 -229
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/pixtral.py +0 -161
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/pixtral/vision.py +0 -320
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/__init__.py +0 -2
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/config.py +0 -108
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/language.py +0 -490
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/qwen2_5_vl.py +0 -168
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_5_vl/vision.py +0 -414
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/__init__.py +0 -2
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/config.py +0 -104
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/language.py +0 -490
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/qwen2_vl.py +0 -167
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen2_vl/vision.py +0 -312
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/base.py +0 -117
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/cache.py +0 -531
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/generate.py +0 -701
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/rope_utils.py +0 -255
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/sample_utils.py +0 -303
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/tokenizer_utils.py +0 -407
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/processor.py +0 -476
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/qwen3vl.py +0 -1223
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/__init__.py +0 -0
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/base.py +0 -117
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/cache.py +0 -531
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/generate.py +0 -701
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/rope_utils.py +0 -255
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/sample_utils.py +0 -303
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/tokenizer_utils.py +0 -407
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/processor.py +0 -476
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/qwen3vl_moe.py +0 -1309
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/switch_layers.py +0 -210
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/smolvlm/__init__.py +0 -8
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/smolvlm/smolvlm.py +0 -62
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/processing_qwen2_5_vl.py +0 -209
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/processing_qwen2_vl.py +0 -215
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/prompt_utils.py +0 -474
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/sample_utils.py +0 -39
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/tokenizer_utils.py +0 -344
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/__init__.py +0 -9
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/lora.py +0 -70
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/trainer.py +0 -296
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/utils.py +0 -160
- nexaai/binds/nexa_mlx/py-lib/vlm/modeling/utils.py +0 -928
- {nexaai-1.0.19rc7.dist-info → nexaai-1.0.19rc9.dist-info}/WHEEL +0 -0
- {nexaai-1.0.19rc7.dist-info → nexaai-1.0.19rc9.dist-info}/top_level.txt +0 -0
|
@@ -1,168 +0,0 @@
|
|
|
1
|
-
import glob
|
|
2
|
-
import json
|
|
3
|
-
from pathlib import Path
|
|
4
|
-
from typing import Optional
|
|
5
|
-
|
|
6
|
-
import mlx.core as mx
|
|
7
|
-
import mlx.nn as nn
|
|
8
|
-
from huggingface_hub import snapshot_download
|
|
9
|
-
|
|
10
|
-
from .config import ModelConfig, TextConfig, VisionConfig
|
|
11
|
-
from .language import LanguageModel
|
|
12
|
-
from .vision import VisionModel
|
|
13
|
-
|
|
14
|
-
|
|
15
|
-
class Model(nn.Module):
|
|
16
|
-
def __init__(self, config: ModelConfig):
|
|
17
|
-
super().__init__()
|
|
18
|
-
self.config = config
|
|
19
|
-
self.vision_tower = VisionModel(config.vision_config)
|
|
20
|
-
self.language_model = LanguageModel(config.text_config, config)
|
|
21
|
-
|
|
22
|
-
def get_input_embeddings(
|
|
23
|
-
self,
|
|
24
|
-
input_ids: Optional[mx.array] = None,
|
|
25
|
-
pixel_values: Optional[mx.array] = None,
|
|
26
|
-
image_grid_thw: Optional[mx.array] = None,
|
|
27
|
-
):
|
|
28
|
-
if pixel_values is None:
|
|
29
|
-
return self.language_model.model.embed_tokens(input_ids)
|
|
30
|
-
|
|
31
|
-
dtype = self.vision_tower.patch_embed.proj.weight.dtype
|
|
32
|
-
pixel_values = pixel_values.astype(dtype)
|
|
33
|
-
|
|
34
|
-
# Get the input embeddings from the language model
|
|
35
|
-
inputs_embeds = self.language_model.model.embed_tokens(input_ids)
|
|
36
|
-
|
|
37
|
-
# Get the ouptut hidden states from the vision model
|
|
38
|
-
hidden_states = self.vision_tower(
|
|
39
|
-
pixel_values, image_grid_thw, output_hidden_states=False
|
|
40
|
-
)
|
|
41
|
-
|
|
42
|
-
# Insert special image tokens in the input_ids
|
|
43
|
-
final_inputs_embeds = self.merge_input_ids_with_image_features(
|
|
44
|
-
self.config.image_token_id,
|
|
45
|
-
self.config.video_token_id,
|
|
46
|
-
hidden_states,
|
|
47
|
-
inputs_embeds,
|
|
48
|
-
input_ids,
|
|
49
|
-
)
|
|
50
|
-
return final_inputs_embeds
|
|
51
|
-
|
|
52
|
-
@staticmethod
|
|
53
|
-
def merge_input_ids_with_image_features(
|
|
54
|
-
image_token_id,
|
|
55
|
-
video_token_id,
|
|
56
|
-
image_features,
|
|
57
|
-
inputs_embeds,
|
|
58
|
-
input_ids,
|
|
59
|
-
):
|
|
60
|
-
"""Merge image features into input embeddings at image token positions.
|
|
61
|
-
|
|
62
|
-
Args:
|
|
63
|
-
image_token_id: The token ID for image placeholders
|
|
64
|
-
video_token_id: The token ID for video placeholders (fallback)
|
|
65
|
-
image_features: Vision features from the vision tower [num_features, hidden_dim]
|
|
66
|
-
inputs_embeds: Input embeddings [batch_size, seq_len, hidden_dim]
|
|
67
|
-
input_ids: Input token IDs [batch_size, seq_len]
|
|
68
|
-
grid_thw: Grid dimensions for each image (optional, not used in simple case)
|
|
69
|
-
|
|
70
|
-
Returns:
|
|
71
|
-
Updated input embeddings with image features inserted
|
|
72
|
-
"""
|
|
73
|
-
# Find positions of image tokens
|
|
74
|
-
image_positions = input_ids == image_token_id
|
|
75
|
-
if mx.sum(image_positions) == 0:
|
|
76
|
-
image_positions = input_ids == video_token_id
|
|
77
|
-
|
|
78
|
-
# Get dimensions
|
|
79
|
-
batch_size, seq_len = input_ids.shape
|
|
80
|
-
|
|
81
|
-
# Process each batch item
|
|
82
|
-
batch_outputs = []
|
|
83
|
-
feature_start_idx = 0
|
|
84
|
-
|
|
85
|
-
for batch_idx in range(batch_size):
|
|
86
|
-
# Get mask for this batch
|
|
87
|
-
image_mask = image_positions[batch_idx]
|
|
88
|
-
num_positions = mx.sum(image_mask).item()
|
|
89
|
-
|
|
90
|
-
if num_positions > 0:
|
|
91
|
-
# Extract features for this batch
|
|
92
|
-
batch_features = image_features[
|
|
93
|
-
feature_start_idx : feature_start_idx + num_positions
|
|
94
|
-
]
|
|
95
|
-
|
|
96
|
-
# Validate we have the right number of features
|
|
97
|
-
if batch_features.shape[0] != num_positions:
|
|
98
|
-
raise ValueError(
|
|
99
|
-
f"Number of image token positions ({num_positions}) does not match "
|
|
100
|
-
f"number of image features ({batch_features.shape[0]}) for batch {batch_idx}"
|
|
101
|
-
)
|
|
102
|
-
|
|
103
|
-
# Create indices for gathering
|
|
104
|
-
cumsum = mx.cumsum(image_mask.astype(mx.int32))
|
|
105
|
-
feature_indices = mx.where(image_mask, cumsum - 1, 0)
|
|
106
|
-
|
|
107
|
-
# Gather features
|
|
108
|
-
gathered_features = batch_features[feature_indices]
|
|
109
|
-
|
|
110
|
-
# Combine with original embeddings
|
|
111
|
-
image_mask_expanded = mx.expand_dims(image_mask, axis=-1)
|
|
112
|
-
batch_output = mx.where(
|
|
113
|
-
image_mask_expanded, gathered_features, inputs_embeds[batch_idx]
|
|
114
|
-
)
|
|
115
|
-
|
|
116
|
-
feature_start_idx += num_positions
|
|
117
|
-
else:
|
|
118
|
-
# No image tokens in this batch item
|
|
119
|
-
batch_output = inputs_embeds[batch_idx]
|
|
120
|
-
|
|
121
|
-
batch_outputs.append(batch_output)
|
|
122
|
-
|
|
123
|
-
# Stack all batch outputs
|
|
124
|
-
return mx.stack(batch_outputs, axis=0)
|
|
125
|
-
|
|
126
|
-
@property
|
|
127
|
-
def layers(self):
|
|
128
|
-
return self.language_model.model.layers
|
|
129
|
-
|
|
130
|
-
def __call__(
|
|
131
|
-
self,
|
|
132
|
-
input_ids: mx.array,
|
|
133
|
-
pixel_values: Optional[mx.array] = None,
|
|
134
|
-
mask: Optional[mx.array] = None,
|
|
135
|
-
cache=None,
|
|
136
|
-
**kwargs,
|
|
137
|
-
):
|
|
138
|
-
image_grid_thw = kwargs.pop("image_grid_thw", None)
|
|
139
|
-
video_grid_thw = kwargs.pop("video_grid_thw", None)
|
|
140
|
-
grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw
|
|
141
|
-
|
|
142
|
-
inputs_embeds = self.get_input_embeddings(input_ids, pixel_values, grid_thw)
|
|
143
|
-
|
|
144
|
-
kwargs = {
|
|
145
|
-
"pixel_values": pixel_values,
|
|
146
|
-
"image_grid_thw": image_grid_thw,
|
|
147
|
-
"video_grid_thw": video_grid_thw,
|
|
148
|
-
**kwargs,
|
|
149
|
-
}
|
|
150
|
-
|
|
151
|
-
logits = self.language_model(
|
|
152
|
-
input_ids, inputs_embeds, mask=mask, cache=cache, **kwargs
|
|
153
|
-
)
|
|
154
|
-
|
|
155
|
-
return logits
|
|
156
|
-
|
|
157
|
-
def sanitize(self, weights):
|
|
158
|
-
def transform_key(key):
|
|
159
|
-
if "vision_tower" not in key:
|
|
160
|
-
key = key.replace("visual", "vision_tower")
|
|
161
|
-
if "language_model" not in key:
|
|
162
|
-
if "model" in key:
|
|
163
|
-
key = key.replace("model", "language_model.model")
|
|
164
|
-
elif "lm_head" in key:
|
|
165
|
-
key = key.replace("lm_head", "language_model.lm_head")
|
|
166
|
-
return key
|
|
167
|
-
|
|
168
|
-
return {transform_key(k): v for k, v in weights.items()}
|
|
@@ -1,414 +0,0 @@
|
|
|
1
|
-
from typing import Optional
|
|
2
|
-
|
|
3
|
-
import mlx.core as mx
|
|
4
|
-
import mlx.nn as nn
|
|
5
|
-
import numpy as np
|
|
6
|
-
|
|
7
|
-
from .config import VisionConfig
|
|
8
|
-
|
|
9
|
-
|
|
10
|
-
def check_array_shape(arr):
|
|
11
|
-
shape = arr.shape
|
|
12
|
-
|
|
13
|
-
# Check if the shape has 4 dimensions
|
|
14
|
-
if len(shape) not in [4, 5]:
|
|
15
|
-
return False
|
|
16
|
-
|
|
17
|
-
B, out_channels, kH, KW, t = shape
|
|
18
|
-
|
|
19
|
-
if t == 3:
|
|
20
|
-
return True
|
|
21
|
-
|
|
22
|
-
# Check if out_channels is the largest, and kH and KW are the same
|
|
23
|
-
if (out_channels >= kH) and (out_channels >= KW) and (kH == KW):
|
|
24
|
-
return True
|
|
25
|
-
else:
|
|
26
|
-
return False
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
def rotate_half(x):
|
|
30
|
-
"""Rotates half the hidden dims of the input."""
|
|
31
|
-
x1 = x[..., : x.shape[-1] // 2]
|
|
32
|
-
x2 = x[..., x.shape[-1] // 2 :]
|
|
33
|
-
return mx.concatenate([-x2, x1], axis=-1)
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
def apply_rotary_pos_emb_vision(tensor, freqs) -> mx.array:
|
|
37
|
-
orig_dtype = tensor.dtype
|
|
38
|
-
|
|
39
|
-
cos = mx.cos(freqs)
|
|
40
|
-
sin = mx.sin(freqs)
|
|
41
|
-
|
|
42
|
-
cos = mx.expand_dims(cos, axis=1) # Equivalent to unsqueeze(1)
|
|
43
|
-
cos = mx.tile(cos, (1, 1, 2)) # Equivalent to repeat(1, 1, 2)
|
|
44
|
-
cos = mx.expand_dims(cos, axis=0) # Equivalent to [None, ...]
|
|
45
|
-
|
|
46
|
-
sin = mx.expand_dims(sin, axis=1) # Equivalent to unsqueeze(1)
|
|
47
|
-
sin = mx.tile(sin, (1, 1, 2)) # Equivalent to repeat(1, 1, 2)
|
|
48
|
-
sin = mx.expand_dims(sin, axis=0) # Equivalent to [None, ...]
|
|
49
|
-
|
|
50
|
-
output = (tensor * cos) + (rotate_half(tensor) * sin)
|
|
51
|
-
return output.astype(orig_dtype)
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
class VisionRotaryEmbedding(nn.Module):
|
|
55
|
-
def __init__(self, dim: int, theta: float = 10000.0) -> None:
|
|
56
|
-
super().__init__()
|
|
57
|
-
self.dim = dim
|
|
58
|
-
self.theta = theta
|
|
59
|
-
|
|
60
|
-
def __call__(self, seqlen: int) -> mx.array:
|
|
61
|
-
inv_freq = 1.0 / (
|
|
62
|
-
self.theta ** (mx.arange(0, self.dim, 2, dtype=mx.float32) / self.dim)
|
|
63
|
-
)
|
|
64
|
-
seq = mx.arange(seqlen.item(), dtype=inv_freq.dtype)
|
|
65
|
-
freqs = mx.outer(seq, inv_freq)
|
|
66
|
-
return freqs
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
class PatchEmbed(nn.Module):
|
|
70
|
-
def __init__(
|
|
71
|
-
self,
|
|
72
|
-
patch_size: int = 14,
|
|
73
|
-
temporal_patch_size: int = 2,
|
|
74
|
-
in_channels: int = 3,
|
|
75
|
-
hidden_size: int = 1152,
|
|
76
|
-
) -> None:
|
|
77
|
-
super().__init__()
|
|
78
|
-
self.patch_size = patch_size
|
|
79
|
-
self.temporal_patch_size = temporal_patch_size
|
|
80
|
-
self.in_channels = in_channels
|
|
81
|
-
self.hidden_size = hidden_size
|
|
82
|
-
|
|
83
|
-
kernel_size = [temporal_patch_size, patch_size, patch_size]
|
|
84
|
-
self.proj = nn.Conv3d(
|
|
85
|
-
in_channels,
|
|
86
|
-
hidden_size,
|
|
87
|
-
kernel_size=kernel_size,
|
|
88
|
-
stride=kernel_size,
|
|
89
|
-
bias=False,
|
|
90
|
-
)
|
|
91
|
-
|
|
92
|
-
def __call__(self, hidden_states: mx.array) -> mx.array:
|
|
93
|
-
hidden_states = hidden_states.reshape(
|
|
94
|
-
-1,
|
|
95
|
-
self.in_channels,
|
|
96
|
-
self.temporal_patch_size,
|
|
97
|
-
self.patch_size,
|
|
98
|
-
self.patch_size,
|
|
99
|
-
).moveaxis(1, 4)
|
|
100
|
-
|
|
101
|
-
hidden_states = self.proj(hidden_states)
|
|
102
|
-
hidden_states = hidden_states.reshape(-1, self.hidden_size)
|
|
103
|
-
return hidden_states
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
class PatchMerger(nn.Module):
|
|
107
|
-
def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None:
|
|
108
|
-
super().__init__()
|
|
109
|
-
self.hidden_size = context_dim * (spatial_merge_size**2)
|
|
110
|
-
self.ln_q = nn.RMSNorm(context_dim, eps=1e-6)
|
|
111
|
-
self.mlp = [
|
|
112
|
-
nn.Linear(self.hidden_size, self.hidden_size),
|
|
113
|
-
nn.GELU(),
|
|
114
|
-
nn.Linear(self.hidden_size, dim),
|
|
115
|
-
]
|
|
116
|
-
|
|
117
|
-
def __call__(self, x: mx.array) -> mx.array:
|
|
118
|
-
x = self.ln_q(x).reshape(-1, self.hidden_size)
|
|
119
|
-
for layer in self.mlp:
|
|
120
|
-
x = layer(x)
|
|
121
|
-
return x
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
class Attention(nn.Module):
|
|
125
|
-
def __init__(self, dim: int, num_heads: int = 16) -> None:
|
|
126
|
-
super().__init__()
|
|
127
|
-
self.num_heads = num_heads
|
|
128
|
-
self.head_dim = head_dim = dim // num_heads
|
|
129
|
-
self.scale = head_dim**-0.5
|
|
130
|
-
self.qkv = nn.Linear(dim, dim * 3, bias=True)
|
|
131
|
-
self.proj = nn.Linear(dim, dim)
|
|
132
|
-
|
|
133
|
-
def __call__(
|
|
134
|
-
self, x: mx.array, cu_seqlens: mx.array, rotary_pos_emb: mx.array = None
|
|
135
|
-
) -> mx.array:
|
|
136
|
-
seq_length = x.shape[0]
|
|
137
|
-
qkv = (
|
|
138
|
-
self.qkv(x).reshape(seq_length, 3, self.num_heads, -1).transpose(1, 0, 2, 3)
|
|
139
|
-
)
|
|
140
|
-
q, k, v = mx.split(qkv, 3)
|
|
141
|
-
|
|
142
|
-
q = apply_rotary_pos_emb_vision(mx.expand_dims(q, 0), rotary_pos_emb)[0]
|
|
143
|
-
k = apply_rotary_pos_emb_vision(mx.expand_dims(k, 0), rotary_pos_emb)[0]
|
|
144
|
-
attention_mask = mx.full(
|
|
145
|
-
(1, seq_length, seq_length), mx.finfo(q.dtype).min, dtype=q.dtype
|
|
146
|
-
)
|
|
147
|
-
|
|
148
|
-
for i in range(1, len(cu_seqlens)):
|
|
149
|
-
start = int(cu_seqlens[i - 1])
|
|
150
|
-
end = int(cu_seqlens[i])
|
|
151
|
-
attention_mask[..., start:end, start:end] = 0
|
|
152
|
-
|
|
153
|
-
q = q.transpose(0, 2, 1, 3)
|
|
154
|
-
k = k.transpose(0, 2, 1, 3)
|
|
155
|
-
v = v.transpose(0, 2, 1, 3)
|
|
156
|
-
|
|
157
|
-
output = mx.fast.scaled_dot_product_attention(
|
|
158
|
-
q, k, v, scale=self.scale, mask=attention_mask
|
|
159
|
-
)
|
|
160
|
-
output = output.transpose(0, 2, 1, 3)
|
|
161
|
-
output = output.reshape(seq_length, -1)
|
|
162
|
-
return self.proj(output)
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
class MLP(nn.Module):
|
|
166
|
-
def __init__(self, dim, hidden_dim):
|
|
167
|
-
super().__init__()
|
|
168
|
-
self.gate_proj = nn.Linear(dim, hidden_dim)
|
|
169
|
-
self.up_proj = nn.Linear(dim, hidden_dim)
|
|
170
|
-
self.down_proj = nn.Linear(hidden_dim, dim)
|
|
171
|
-
|
|
172
|
-
def __call__(self, x: mx.array) -> mx.array:
|
|
173
|
-
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
class Qwen2VLVisionBlock(nn.Module):
|
|
177
|
-
def __init__(self, config: VisionConfig) -> None:
|
|
178
|
-
super().__init__()
|
|
179
|
-
self.norm1 = nn.RMSNorm(config.hidden_size, eps=1e-6)
|
|
180
|
-
self.norm2 = nn.RMSNorm(config.hidden_size, eps=1e-6)
|
|
181
|
-
|
|
182
|
-
self.attn = Attention(dim=config.hidden_size, num_heads=config.num_heads)
|
|
183
|
-
self.mlp = MLP(dim=config.hidden_size, hidden_dim=config.intermediate_size)
|
|
184
|
-
|
|
185
|
-
def __call__(self, hidden_states, cu_seqlens, rotary_pos_emb) -> mx.array:
|
|
186
|
-
hidden_states = hidden_states + self.attn(
|
|
187
|
-
self.norm1(hidden_states),
|
|
188
|
-
cu_seqlens=cu_seqlens,
|
|
189
|
-
rotary_pos_emb=rotary_pos_emb,
|
|
190
|
-
)
|
|
191
|
-
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
|
|
192
|
-
return hidden_states
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
class VisionModel(nn.Module):
|
|
196
|
-
|
|
197
|
-
def __init__(self, config: VisionConfig) -> None:
|
|
198
|
-
super().__init__()
|
|
199
|
-
self.config = config
|
|
200
|
-
self.model_type = config.model_type
|
|
201
|
-
if self.model_type != "qwen2_5_vl":
|
|
202
|
-
raise ValueError(f"Unsupported model type: {self.model_type}")
|
|
203
|
-
self.spatial_merge_size = config.spatial_merge_size
|
|
204
|
-
|
|
205
|
-
self.patch_embed = PatchEmbed(
|
|
206
|
-
patch_size=config.patch_size,
|
|
207
|
-
temporal_patch_size=config.temporal_patch_size,
|
|
208
|
-
in_channels=config.in_channels,
|
|
209
|
-
hidden_size=config.hidden_size,
|
|
210
|
-
)
|
|
211
|
-
|
|
212
|
-
self.window_size = config.window_size
|
|
213
|
-
self.patch_size = config.patch_size
|
|
214
|
-
self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size
|
|
215
|
-
self.fullatt_block_indexes = config.fullatt_block_indexes
|
|
216
|
-
head_dim = config.hidden_size // config.num_heads
|
|
217
|
-
self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
|
|
218
|
-
|
|
219
|
-
self.blocks = [Qwen2VLVisionBlock(config) for _ in range(config.depth)]
|
|
220
|
-
self.merger = PatchMerger(
|
|
221
|
-
dim=config.out_hidden_size, context_dim=config.hidden_size
|
|
222
|
-
)
|
|
223
|
-
|
|
224
|
-
def rot_pos_emb(self, grid_thw):
|
|
225
|
-
pos_ids = []
|
|
226
|
-
|
|
227
|
-
for t, h, w in grid_thw.tolist():
|
|
228
|
-
hpos_ids = mx.expand_dims(mx.arange(h), 1)
|
|
229
|
-
hpos_ids = mx.repeat(hpos_ids, w, axis=1)
|
|
230
|
-
hpos_ids = hpos_ids.reshape(
|
|
231
|
-
h // self.spatial_merge_size,
|
|
232
|
-
self.spatial_merge_size,
|
|
233
|
-
w // self.spatial_merge_size,
|
|
234
|
-
self.spatial_merge_size,
|
|
235
|
-
)
|
|
236
|
-
hpos_ids = mx.transpose(hpos_ids, (0, 2, 1, 3))
|
|
237
|
-
hpos_ids = hpos_ids.flatten()
|
|
238
|
-
|
|
239
|
-
wpos_ids = mx.expand_dims(mx.arange(w), 0)
|
|
240
|
-
wpos_ids = mx.repeat(wpos_ids, h, axis=0)
|
|
241
|
-
wpos_ids = wpos_ids.reshape(
|
|
242
|
-
h // self.spatial_merge_size,
|
|
243
|
-
self.spatial_merge_size,
|
|
244
|
-
w // self.spatial_merge_size,
|
|
245
|
-
self.spatial_merge_size,
|
|
246
|
-
)
|
|
247
|
-
wpos_ids = mx.transpose(wpos_ids, (0, 2, 1, 3))
|
|
248
|
-
wpos_ids = wpos_ids.flatten()
|
|
249
|
-
|
|
250
|
-
stacked_pos_ids = mx.stack([hpos_ids, wpos_ids], axis=-1)
|
|
251
|
-
pos_ids.append(mx.tile(stacked_pos_ids, (t, 1)))
|
|
252
|
-
|
|
253
|
-
pos_ids = mx.concatenate(pos_ids, axis=0)
|
|
254
|
-
max_grid_size = mx.max(grid_thw[:, 1:])
|
|
255
|
-
rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
|
|
256
|
-
rotary_pos_emb = rotary_pos_emb_full[pos_ids]
|
|
257
|
-
|
|
258
|
-
return rotary_pos_emb.reshape(pos_ids.shape[0], -1)
|
|
259
|
-
|
|
260
|
-
def get_window_index(self, grid_thw):
|
|
261
|
-
window_index = []
|
|
262
|
-
cu_window_seqlens = [0]
|
|
263
|
-
window_index_id = 0
|
|
264
|
-
vit_merger_window_size = (
|
|
265
|
-
self.window_size // self.spatial_merge_size // self.patch_size
|
|
266
|
-
)
|
|
267
|
-
|
|
268
|
-
for grid_t, grid_h, grid_w in grid_thw.tolist():
|
|
269
|
-
llm_grid_h = grid_h // self.spatial_merge_size
|
|
270
|
-
llm_grid_w = grid_w // self.spatial_merge_size
|
|
271
|
-
|
|
272
|
-
index = mx.arange(grid_t * llm_grid_h * llm_grid_w).reshape(
|
|
273
|
-
grid_t, llm_grid_h, llm_grid_w
|
|
274
|
-
)
|
|
275
|
-
|
|
276
|
-
pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size
|
|
277
|
-
pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size
|
|
278
|
-
num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
|
|
279
|
-
num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
|
|
280
|
-
|
|
281
|
-
# Replace F.pad with np.pad
|
|
282
|
-
index_padded = mx.pad(
|
|
283
|
-
index,
|
|
284
|
-
((0, 0), (0, pad_h), (0, pad_w)),
|
|
285
|
-
mode="constant",
|
|
286
|
-
constant_values=-100,
|
|
287
|
-
)
|
|
288
|
-
|
|
289
|
-
index_padded = index_padded.reshape(
|
|
290
|
-
grid_t,
|
|
291
|
-
num_windows_h,
|
|
292
|
-
vit_merger_window_size,
|
|
293
|
-
num_windows_w,
|
|
294
|
-
vit_merger_window_size,
|
|
295
|
-
)
|
|
296
|
-
|
|
297
|
-
# Replace permute with np.transpose
|
|
298
|
-
index_padded = mx.transpose(index_padded, (0, 1, 3, 2, 4)).reshape(
|
|
299
|
-
grid_t,
|
|
300
|
-
num_windows_h * num_windows_w,
|
|
301
|
-
vit_merger_window_size,
|
|
302
|
-
vit_merger_window_size,
|
|
303
|
-
)
|
|
304
|
-
|
|
305
|
-
# Replace torch operations with numpy
|
|
306
|
-
seqlens = mx.sum(index_padded != -100, axis=(2, 3)).reshape(-1)
|
|
307
|
-
index_padded = index_padded.reshape(-1)
|
|
308
|
-
index = np.where(index_padded != -100)[
|
|
309
|
-
0
|
|
310
|
-
].tolist() # [i for i, x in enumerate(index_padded) if x != -100]
|
|
311
|
-
index_new = index_padded[index]
|
|
312
|
-
|
|
313
|
-
window_index.append(index_new + window_index_id)
|
|
314
|
-
cu_seqlens_tmp = (
|
|
315
|
-
mx.cumsum(seqlens, axis=0) * self.spatial_merge_unit
|
|
316
|
-
+ cu_window_seqlens[-1]
|
|
317
|
-
)
|
|
318
|
-
cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
|
|
319
|
-
window_index_id += int(grid_t * llm_grid_h * llm_grid_w)
|
|
320
|
-
|
|
321
|
-
# Replace torch.cat with np.concatenate
|
|
322
|
-
window_index = mx.concatenate(window_index, axis=0)
|
|
323
|
-
cu_window_seqlens = mx.array(cu_window_seqlens)
|
|
324
|
-
|
|
325
|
-
return window_index, cu_window_seqlens
|
|
326
|
-
|
|
327
|
-
def __call__(
|
|
328
|
-
self,
|
|
329
|
-
hidden_states: mx.array,
|
|
330
|
-
grid_thw: mx.array,
|
|
331
|
-
output_hidden_states: Optional[bool] = None,
|
|
332
|
-
) -> mx.array:
|
|
333
|
-
|
|
334
|
-
hidden_states = self.patch_embed(hidden_states)
|
|
335
|
-
rotary_pos_emb = self.rot_pos_emb(grid_thw)
|
|
336
|
-
window_index, cu_window_seqlens = self.get_window_index(grid_thw)
|
|
337
|
-
|
|
338
|
-
# Get indices of first occurrence of each unique value
|
|
339
|
-
seen = set()
|
|
340
|
-
idx = []
|
|
341
|
-
for i, x in enumerate(cu_window_seqlens):
|
|
342
|
-
if x not in seen:
|
|
343
|
-
seen.add(x)
|
|
344
|
-
idx.append(i)
|
|
345
|
-
|
|
346
|
-
idx = mx.array(idx, dtype=mx.int32)
|
|
347
|
-
cu_window_seqlens = cu_window_seqlens[idx]
|
|
348
|
-
|
|
349
|
-
seq_len, _ = hidden_states.shape
|
|
350
|
-
hidden_states = hidden_states.reshape(
|
|
351
|
-
seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1
|
|
352
|
-
)
|
|
353
|
-
hidden_states = hidden_states[window_index, :, :]
|
|
354
|
-
hidden_states = hidden_states.reshape(seq_len, -1)
|
|
355
|
-
rotary_pos_emb = rotary_pos_emb.reshape(
|
|
356
|
-
seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1
|
|
357
|
-
)
|
|
358
|
-
rotary_pos_emb = rotary_pos_emb[window_index, :, :]
|
|
359
|
-
rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
|
|
360
|
-
|
|
361
|
-
# Assuming grid_thw has shape (batch_size, 3)
|
|
362
|
-
batch_size = grid_thw.shape[0]
|
|
363
|
-
|
|
364
|
-
# Calculate cu_seqlens for each item in the batch
|
|
365
|
-
cu_seqlens = []
|
|
366
|
-
for i in range(batch_size):
|
|
367
|
-
seq_len = grid_thw[i, 1] * grid_thw[i, 2]
|
|
368
|
-
cu_seqlens.append(mx.repeat(seq_len, grid_thw[i, 0]))
|
|
369
|
-
|
|
370
|
-
# Concatenate the cu_seqlens for all items in the batch
|
|
371
|
-
cu_seqlens = mx.concatenate(cu_seqlens)
|
|
372
|
-
|
|
373
|
-
cu_seqlens = mx.cumsum(cu_seqlens.astype(mx.int32), axis=0)
|
|
374
|
-
cu_seqlens = mx.pad(cu_seqlens, (1, 0), mode="constant", constant_values=0)
|
|
375
|
-
|
|
376
|
-
encoder_states = (hidden_states,) if output_hidden_states else None
|
|
377
|
-
|
|
378
|
-
for layer_num, blk in enumerate(self.blocks):
|
|
379
|
-
if layer_num in self.fullatt_block_indexes:
|
|
380
|
-
cu_seqlens_now = cu_seqlens
|
|
381
|
-
else:
|
|
382
|
-
cu_seqlens_now = cu_window_seqlens
|
|
383
|
-
|
|
384
|
-
hidden_states = blk(
|
|
385
|
-
hidden_states, cu_seqlens=cu_seqlens_now, rotary_pos_emb=rotary_pos_emb
|
|
386
|
-
)
|
|
387
|
-
|
|
388
|
-
if output_hidden_states:
|
|
389
|
-
encoder_states = encoder_states + (hidden_states,)
|
|
390
|
-
|
|
391
|
-
hidden_states = self.merger(hidden_states)
|
|
392
|
-
reverse_indices = mx.argsort(window_index, axis=0)
|
|
393
|
-
hidden_states = hidden_states[reverse_indices, :]
|
|
394
|
-
return hidden_states
|
|
395
|
-
|
|
396
|
-
def sanitize(self, weights):
|
|
397
|
-
sanitized_weights = {}
|
|
398
|
-
for k, v in weights.items():
|
|
399
|
-
if "position_ids" in k:
|
|
400
|
-
# Remove unused position_ids
|
|
401
|
-
continue
|
|
402
|
-
elif "patch_embed.proj.weight" in k:
|
|
403
|
-
# PyTorch conv2d weight tensors have shape:
|
|
404
|
-
# [out_channels, in_channels, kH, KW]
|
|
405
|
-
# MLX conv2d expects the weight be of shape:
|
|
406
|
-
# [out_channels, kH, KW, in_channels]
|
|
407
|
-
if check_array_shape(v):
|
|
408
|
-
sanitized_weights[k] = v
|
|
409
|
-
else:
|
|
410
|
-
sanitized_weights[k] = v.transpose(0, 2, 3, 4, 1)
|
|
411
|
-
else:
|
|
412
|
-
sanitized_weights[k] = v
|
|
413
|
-
|
|
414
|
-
return sanitized_weights
|
|
@@ -1,104 +0,0 @@
|
|
|
1
|
-
import inspect
|
|
2
|
-
from dataclasses import dataclass
|
|
3
|
-
from typing import Dict, List, Optional, Union
|
|
4
|
-
|
|
5
|
-
|
|
6
|
-
@dataclass
|
|
7
|
-
class VisionConfig:
|
|
8
|
-
model_type: str = "qwen2_vl"
|
|
9
|
-
depth: int = 32
|
|
10
|
-
embed_dim: int = 1280
|
|
11
|
-
hidden_size: int = 1536
|
|
12
|
-
num_heads: int = 16
|
|
13
|
-
image_size: int = 384
|
|
14
|
-
patch_size: int = 14
|
|
15
|
-
vocab_size: int = 32000
|
|
16
|
-
mlp_ratio: float = 4.0
|
|
17
|
-
in_channels: int = 3
|
|
18
|
-
layer_norm_eps: float = 1e-6
|
|
19
|
-
spatial_patch_size: int = 14
|
|
20
|
-
spatial_merge_size: int = 2
|
|
21
|
-
temporal_patch_size: int = 2
|
|
22
|
-
|
|
23
|
-
@classmethod
|
|
24
|
-
def from_dict(cls, params):
|
|
25
|
-
return cls(
|
|
26
|
-
**{
|
|
27
|
-
k: v
|
|
28
|
-
for k, v in params.items()
|
|
29
|
-
if k in inspect.signature(cls).parameters
|
|
30
|
-
}
|
|
31
|
-
)
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
@dataclass
|
|
35
|
-
class TextConfig:
|
|
36
|
-
model_type: str
|
|
37
|
-
hidden_size: int
|
|
38
|
-
num_hidden_layers: int
|
|
39
|
-
intermediate_size: int
|
|
40
|
-
num_attention_heads: int
|
|
41
|
-
rms_norm_eps: float
|
|
42
|
-
vocab_size: int
|
|
43
|
-
num_key_value_heads: Optional[int] = 8
|
|
44
|
-
max_position_embeddings: Optional[int] = 40960
|
|
45
|
-
rope_theta: float = 1000000.0
|
|
46
|
-
rope_traditional: bool = False
|
|
47
|
-
rope_scaling: Optional[Dict[str, Union[float, str]]] = None
|
|
48
|
-
tie_word_embeddings: bool = False
|
|
49
|
-
sliding_window: int = 32768
|
|
50
|
-
use_sliding_window: bool = False
|
|
51
|
-
use_cache: bool = True
|
|
52
|
-
|
|
53
|
-
def __post_init__(self):
|
|
54
|
-
if self.num_key_value_heads is None:
|
|
55
|
-
self.num_key_value_heads = self.num_attention_heads
|
|
56
|
-
|
|
57
|
-
if self.rope_scaling:
|
|
58
|
-
required_keys = {"mrope_section", "type"}
|
|
59
|
-
if not all(key in self.rope_scaling for key in required_keys):
|
|
60
|
-
raise ValueError(f"rope_scaling must contain keys {required_keys}")
|
|
61
|
-
|
|
62
|
-
if not self.rope_scaling["type"] in ["mrope", "default"]:
|
|
63
|
-
raise ValueError(f"rope_scaling type must be 'mrope' or 'default'")
|
|
64
|
-
|
|
65
|
-
@classmethod
|
|
66
|
-
def from_dict(cls, params):
|
|
67
|
-
return cls(
|
|
68
|
-
**{
|
|
69
|
-
k: v
|
|
70
|
-
for k, v in params.items()
|
|
71
|
-
if k in inspect.signature(cls).parameters
|
|
72
|
-
}
|
|
73
|
-
)
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
@dataclass
|
|
77
|
-
class ModelConfig:
|
|
78
|
-
text_config: TextConfig
|
|
79
|
-
vision_config: VisionConfig
|
|
80
|
-
model_type: str
|
|
81
|
-
ignore_index: int = -100
|
|
82
|
-
image_token_id: int = 151655
|
|
83
|
-
video_token_id: int = 151656
|
|
84
|
-
vision_start_token_id: int = 151652
|
|
85
|
-
vision_feature_select_strategy: str = "default"
|
|
86
|
-
vision_feature_layer: int = -2
|
|
87
|
-
vocab_size: int = 32000
|
|
88
|
-
eos_token_id: Optional[List[int]] = None
|
|
89
|
-
|
|
90
|
-
@classmethod
|
|
91
|
-
def from_dict(cls, params):
|
|
92
|
-
# Copy text config parameters from root level
|
|
93
|
-
excluded_keys = {"vision_config"}
|
|
94
|
-
params["text_config"] = dict(
|
|
95
|
-
filter(lambda x: x[0] not in excluded_keys, params.items())
|
|
96
|
-
)
|
|
97
|
-
|
|
98
|
-
return cls(
|
|
99
|
-
**{
|
|
100
|
-
k: v
|
|
101
|
-
for k, v in params.items()
|
|
102
|
-
if k in inspect.signature(cls).parameters
|
|
103
|
-
}
|
|
104
|
-
)
|