Merge pull request #24 from Blaizzy/pc/quantise-irregular

Add support for PaliGemma and Quant Siglip

mlx_vlm/chat_ui.py

@@ -54,7 +54,10 @@ def generate(
     else:
         tokenizer = processor.tokenizer
 
-    input_ids, pixel_values = prepare_inputs(image_processor, processor, image, prompt)
+    image_token_index = model.config.image_token_index
+    input_ids, pixel_values = prepare_inputs(
+        image_processor, processor, image, prompt, image_token_index
+    )
     logits, cache = model(input_ids, pixel_values)
     logits = logits[:, -1, :]
     y, _ = sample(logits, temp, top_p)
@@ -89,8 +92,7 @@ def generate(
 def chat(message, history, temperature, max_tokens):
 
     chat = []
-
-    if message["files"]:
+    if len(message["files"]) >= 0:
         chat.append(get_message_json(config["model_type"], message["text"]))
     else:
         raise Exception("Please upload an image. Text only chat is not supported.")
@@ -103,16 +105,20 @@ def chat(message, history, temperature, max_tokens):
         )
 
     elif "tokenizer" in processor.__dict__.keys():
-        messages = processor.tokenizer.apply_chat_template(
-            chat,
-            tokenize=False,
-            add_generation_prompt=True,
-        )
+        if processor.tokenizer.chat_template:
+            messages = processor.tokenizer.apply_chat_template(
+                chat,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+        else:
+            messages = message["text"]
+
     response = ""
     for chunk in generate(
         model,
         processor,
-        message["files"][0],
+        message["files"][-1],
         messages,
         image_processor,
         temperature,
@@ -130,7 +136,7 @@ def chat(message, history, temperature, max_tokens):
     ),
     additional_inputs=[
         gr.Slider(
-            minimum=0, maximum=1, step=0.1, value=0.9, label="Temperature", render=False
+            minimum=0, maximum=1, step=0.1, value=0.1, label="Temperature", render=False
         ),
         gr.Slider(
             minimum=128,

mlx_vlm/generate.py

@@ -78,11 +78,12 @@ def main():
         )
 
     elif "tokenizer" in processor.__dict__.keys():
-        prompt = processor.tokenizer.apply_chat_template(
-            [get_message_json(config["model_type"], prompt)],
-            tokenize=False,
-            add_generation_prompt=True,
-        )
+        if model.config.model_type != "paligemma":
+            prompt = processor.tokenizer.apply_chat_template(
+                [get_message_json(config["model_type"], prompt)],
+                tokenize=False,
+                add_generation_prompt=True,
+            )
 
     else:
         ValueError(

mlx_vlm/models/idefics2/idefics2.py

@@ -244,7 +244,9 @@ def _prepare_inputs_for_multimodal(self, image_features, inputs_embeds, input_id
 
         return inputs_embeds
 
-    def __call__(self, input_ids: mx.array, pixel_values: mx.array, cache=None):
+    def __call__(
+        self, input_ids: mx.array, pixel_values: mx.array, mask: mx.array, cache=None
+    ):
         input_embeddings = self.get_input_embeddings(input_ids, pixel_values)
         logits, cache = self.language_model(
             inputs=input_ids, cache=cache, inputs_embeds=input_embeddings

mlx_vlm/models/idefics2/language.py

@@ -147,6 +147,7 @@ def __call__(
         inputs: mx.array,
         cache=None,
         inputs_embeds=None,
+        mask: Optional[mx.array] = None,
     ):
         # for passing merged input embeddings
         if inputs_embeds is None:

mlx_vlm/models/llava/language.py

@@ -198,6 +198,7 @@ def __call__(
         inputs: mx.array,
         cache=None,
         inputs_embeds=None,
+        mask: Optional[mx.array] = None,
     ):
         out, cache = self.model(inputs, cache, inputs_embeds)
         return self.lm_head(out), cache

mlx_vlm/models/llava/llava.py

@@ -18,6 +18,7 @@
 class ModelConfig:
     text_config: TextConfig
     vision_config: VisionConfig
+    model_type: str
     ignore_index: int = -100
     image_token_index: int = 32000
     vision_feature_select_strategy: str = "default"
@@ -130,7 +131,9 @@ def _merge_input_ids_with_image_features(
         # (1, num_image_patches*num_images + sequence_len, embed_dim)
         return mx.concatenate(final_embeddings, axis=1)
 
-    def __call__(self, input_ids: mx.array, pixel_values: mx.array, cache=None):
+    def __call__(
+        self, input_ids: mx.array, pixel_values: mx.array, mask: mx.array, cache=None
+    ):
         input_embddings = self.get_input_embeddings(input_ids, pixel_values)
         logits, cache = self.language_model(
             input_ids, cache=cache, inputs_embeds=input_embddings

mlx_vlm/models/llava/vision.py

@@ -5,6 +5,7 @@
 
 import mlx.core as mx
 import mlx.nn as nn
+import numpy as np
 
 
 @dataclass
@@ -172,8 +173,10 @@ def __call__(self, x: mx.array) -> mx.array:
         cls_embeddings = mx.broadcast_to(
             self.class_embedding, (batch_size, 1, embed_dim)
         )
+        position_ids = mx.array(np.arange(self.num_positions)[None, :])
+
         embeddings = mx.concatenate((cls_embeddings, patch_embeddings), axis=1)
-        embeddings += self.position_embedding.weight
+        embeddings += self.position_embedding(position_ids)
         return embeddings
 
 
@@ -219,8 +222,7 @@ def __call__(
     ) -> mx.array:
         return self.vision_model(x, output_hidden_states)
 
-    @staticmethod
-    def sanitize(weights):
+    def sanitize(self, weights):
         sanitized_weights = {}
         for k, v in weights.items():
             if "position_ids" in k:

mlx_vlm/models/nanoLlava/language.py

@@ -193,6 +193,7 @@ def __call__(
         inputs: mx.array,
         cache=None,
         inputs_embeds=None,
+        mask: Optional[mx.array] = None,
     ):
         out, cache = self.model(inputs, cache, inputs_embeds=inputs_embeds)
         return out, cache

mlx_vlm/models/nanoLlava/nanoLlava.py

@@ -177,7 +177,13 @@ def _prepare_inputs_for_multimodal(self, image_features, inputs_embeds, input_id
         # (1, num_image_patches*num_images + sequence_len, embed_dim)
         return mx.concatenate(final_embeddings, axis=1)
 
-    def __call__(self, input_ids: mx.array, pixel_values: mx.array, cache=None):
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: mx.array,
+        mask: Optional[mx.array] = None,
+        cache=None,
+    ):
         input_embeddings = self.get_input_embeddings(input_ids, pixel_values)
         logits, cache = self.language_model(
             inputs=input_ids, cache=cache, inputs_embeds=input_embeddings

mlx_vlm/models/nanoLlava/vision.py

@@ -4,6 +4,7 @@
 
 import mlx.core as mx
 import mlx.nn as nn
+import numpy as np
 
 
 @dataclass
@@ -206,9 +207,9 @@ def __call__(self, x: mx.array) -> mx.array:
         batch_size = x.shape[0]
         patch_embeddings = self.patch_embedding(x)
         patch_embeddings = mx.flatten(patch_embeddings, start_axis=1, end_axis=2)
-
+        self.position_ids = mx.array(np.arange(self.num_positions)[None, :])
         embeddings = patch_embeddings
-        embeddings += self.position_embedding.weight
+        embeddings += self.position_embedding(self.position_ids)
         return embeddings
 
 

mlx_vlm/models/paligemma/__init__.py

@@ -0,0 +1,8 @@
+from .paligemma import (
+    LanguageModel,
+    Model,
+    ModelConfig,
+    TextConfig,
+    VisionConfig,
+    VisionModel,
+)