Pixtral: Refactor vision model, update example

Author: turboderp

exllamav2/generator/dynamic_embeddings.py

@@ -35,6 +35,8 @@ class ExLlamaV2MMEmbedding:
     first_index: int
     length: int
 
+    metadata: dict
+
     def __init__(
         self,
         model: ExLlamaV2,
@@ -57,6 +59,7 @@ def __init__(
         self.embeddings = embeddings
         self.text_alias = text_alias
         self.model = model
+        self.metadata = {}
 
         self.length = embeddings.shape[0]
         dim = embeddings.shape[1]

exllamav2/vlm/preprocessor/pixtral.py renamed to exllamav2/vlm/processor/pixtral.py

@@ -1,6 +1,7 @@
 import torch
 import numpy as np
 from PIL import Image
+from exllamav2 import ExLlamaV2, ExLlamaV2Tokenizer
 from exllamav2.config import ExLlamaV2Config
 from exllamav2.vlm.util import (
     convert_to_rgb,
@@ -41,4 +42,32 @@ def preprocess(
 
     image = image.transpose(2, 0, 1)
     image = torch.from_numpy(image).half()
-    return image
+    return image
+
+def postprocess(
+    model: ExLlamaV2,
+    tokenizer: ExLlamaV2Tokenizer,
+    embeddings: torch.Tensor,
+    features_y: int,
+    features_x: int,
+):
+    """
+    Insert [IMG_BREAK] and [IMG_END] tokens in image feature embeddings
+    """
+
+    assert embeddings.shape[0] == features_y * features_x, \
+        "Invalid shape for embeddings"
+
+    id_break = tokenizer.single_id("[IMG_BREAK]")
+    id_end = tokenizer.single_id("[IMG_END]")
+    img_break = model.modules[0].forward(torch.tensor([id_break], dtype=torch.long)).to("cuda:0")
+    img_end = model.modules[0].forward(torch.tensor([id_end], dtype=torch.long)).to("cuda:0")
+
+    dim = embeddings.shape[-1]
+    embeddings = embeddings.view((features_y, features_x, dim))
+    break_col = img_break.expand(features_y, -1, -1)
+    embeddings = torch.cat((embeddings, break_col), dim = 1)
+    embeddings = embeddings.view((features_y * (features_x + 1)), dim)
+    embeddings = torch.cat((embeddings, img_end), dim = 0)
+
+    return embeddings

exllamav2/vlm/vision_tower.py

@@ -4,15 +4,17 @@
 import threading
 
 import torch
-from exllamav2 import ExLlamaV2
+from exllamav2 import ExLlamaV2, ExLlamaV2Tokenizer
 from exllamav2.conv2d import ExLlamaV2Conv2D
 from exllamav2.rmsnorm import ExLlamaV2RMSNorm
 from exllamav2.attn import ExLlamaV2Attention
 from exllamav2.mlp import ExLlamaV2MLP
 from exllamav2.config import ExLlamaV2Config
 from exllamav2.module import ExLlamaV2Module
-from exllamav2.vlm.preprocessor import pixtral
+from exllamav2.vlm.processor import pixtral
 from exllamav2.compat import safe_move_tensor
+from exllamav2.generator import ExLlamaV2MMEmbedding
+from typing import Callable
 
 from PIL.Image import Image
 from exllamav2.vlm.util import position_ids_in_meshgrid
@@ -35,7 +37,8 @@ def __init__(
         # Preprocessor
 
         if cfg.vision_model_type == "pixtral":
-            self.preprocessor = pixtral.preprocess
+            self.preprocess_func = pixtral.preprocess
+            self.postprocess_func = pixtral.postprocess
         else:
             raise ValueError(f"Unknown vision model type: {cfg.vision_model_type}")
 
@@ -90,16 +93,34 @@ def __init__(
             mlp = ExLlamaV2MLP(self, layer_key, layer_idx, archparams = self.archparams)
             self.modules += [attn, mlp]
 
+        # Multimodal projection
+
+        mmp = ExLlamaV2MLP(
+            self,
+            cfg.arch.mmp_prefix,
+            0,
+            archparams = cfg.arch.mmp,
+            in_features = cfg.vision_hidden_size,
+            out_features = cfg.hidden_size,
+            interm_features = cfg.hidden_size,
+            has_norm = False,
+            has_residual = False
+        )
+        self.modules += [mmp]
+
 
     def forward(self, **kwargs):
         raise NotImplementedError()
-
-
-    def preprocess(self, image: Image) -> torch.Tensor:
-        """
-        Preprocess image and prepare for vision tower
-        """
-        return self.preprocessor(self.config, image)
+    def forward_chunk(self, **kwargs):
+        raise NotImplementedError()
+    def load_tp(self, **kwargs):
+        raise ValueError("load_tp not supported for vision model")
+    def load_tp_gen(self, **kwargs):
+        raise ValueError("load_tp not supported for vision model")
+    def load_autosplit(self, **kwargs):
+        raise ValueError("load_autosplit not supported for vision model")
+    def load_autosplit_gen(self, **kwargs):
+        raise ValueError("load_autosplit not supported for vision model")
 
 
     def process(
@@ -134,7 +155,7 @@ def process(
             # Onward
 
             n_device = module.device_idx
-            if n_device is not None and n_device != device and n_device >= 0:
+            if idx == 0 or (n_device is not None and n_device != device and n_device >= 0):
                 hidden_states = safe_move_tensor(hidden_states, n_device, non_blocking = True)
 
             if cos.device != hidden_states.device:
@@ -149,4 +170,62 @@ def process(
                 }
             )
 
-        return hidden_states
+        return hidden_states
+
+
+    def get_image_embeddings(
+        self,
+        model: ExLlamaV2,
+        tokenizer: ExLlamaV2Tokenizer,
+        image: Image,
+        text_alias: str,
+    ) -> ExLlamaV2MMEmbedding:
+        """
+        :param model:
+            Text model for which to produce embeddings
+
+        :param tokenizer:
+            Tokenizer
+
+        :param image:
+            Input PIL image
+
+        :param text_alias:
+            Text string to represent this embedding for tokenizing
+
+        :return:
+            ExLlamaV2MMEmbedding
+        """
+
+        width, height = image.size
+        original_size = (height, width)
+
+        image_tensor = self.preprocess_func(self.config, image)
+        image_size = tuple(image_tensor.shape[1:])
+
+        embedding_tensor = self.process(image_tensor)
+
+        features_y = image_size[0] // 16
+        features_x = image_size[1] // 16
+
+        embedding_tensor = self.postprocess_func(
+            model,
+            tokenizer,
+            embedding_tensor[0],
+            features_y,
+            features_x,
+        )
+
+        mme = ExLlamaV2MMEmbedding(
+            model = model,
+            embeddings = embedding_tensor,
+            text_alias = text_alias
+        )
+
+        mme.metadata.update({
+            "original_size": original_size,
+            "preprocessed_size": image_size,
+            "patches_size": (features_y, features_x),
+        })
+
+        return mme

experimental/multimodal_pixtral_hf.py

@@ -8,52 +8,29 @@
     ExLlamaV2Config,
     ExLlamaV2Cache,
     ExLlamaV2Tokenizer,
-    ExLlamaV2MultimodalProjector,
-    ExLlamaV2VisionTower
+    ExLlamaV2VisionTower,
 )
 
 from exllamav2.generator import (
     ExLlamaV2DynamicGenerator,
     ExLlamaV2Sampler,
-    ExLlamaV2MMEmbedding
 )
 
 from PIL import Image
 import requests
 
-# Get an input image
-
-url = "https://pbs.twimg.com/media/BAeuBsnCIAAUITV.jpg:large"
-image = Image.open(requests.get(url, stream = True).raw)
-
 # Unquantized model used for experiment:
 #
 # https://huggingface.co/mistral-community/pixtral-12b/
 
 model_directory = "/mnt/str/models/pixtral-12b"
 config = ExLlamaV2Config(model_directory)
-config.max_seq_len = 32768  # default is 1M
-
-# Load multimodal projector
-
-multimodal_projector = ExLlamaV2MultimodalProjector(config)
-multimodal_projector.load()
-
-# Load vision tower and preprocessor
-
-vision_tower = ExLlamaV2VisionTower(config)
-vision_tower.load(progress = True)
-
-# Preprocess
-
-image_tensor = vision_tower.preprocess(image)
-image_tensor = image_tensor.cuda()
-image_size = tuple(image_tensor.shape[1:])
+config.max_seq_len = 16384  # default is 1M
 
-# Produce embeddings
+# Load vision model and multimodal projector and initialize preprocessor
 
-embeddings = vision_tower.process(image_tensor)
-embeddings = multimodal_projector.forward(embeddings)[0]
+vision_model = ExLlamaV2VisionTower(config)
+vision_model.load(progress = True)
 
 # Load EXL2 model
 
@@ -62,24 +39,6 @@
 model.load_autosplit(cache, progress = True)
 tokenizer = ExLlamaV2Tokenizer(config)
 
-# Insert [IMG_BREAK] and [IMG_END] tokens.
-
-features_x = image_size[1] // 16
-features_y = image_size[0] // 16
-assert image_size == (features_y * 16, features_x * 16)  # Image should be padded in preprocessing
-
-id_break = tokenizer.single_id("[IMG_BREAK]")
-id_end = tokenizer.single_id("[IMG_END]")
-img_break = model.modules[0].forward(torch.tensor([id_break], dtype = torch.long)).to("cuda:0")
-img_end = model.modules[0].forward(torch.tensor([id_end], dtype = torch.long)).to("cuda:0")
-
-dim = embeddings.shape[-1]
-embeddings = embeddings.view((features_y, features_x, dim))
-break_col = img_break.expand(features_y, -1, -1)
-embeddings = torch.cat((embeddings, break_col), dim = 1)
-embeddings = embeddings.view((features_y * (features_x + 1)), dim)
-embeddings = torch.cat((embeddings, img_end), dim = 0)
-
 # Create generator
 
 generator = ExLlamaV2DynamicGenerator(
@@ -90,15 +49,24 @@
 
 # Create an MMEmbedding for the image features and a prompt containing the placeholder string
 
-image_tokens_a = ExLlamaV2MMEmbedding(
-    model = model,
-    embeddings = embeddings,
-    text_alias = "{{EMBED_A}}"
-)
-
-prompt = "[INST]{{EMBED_A}}\nDescribe the image.[/INST]"
-
-# Pass embeddings to generator
+image_embeddings = [
+    vision_model.get_image_embeddings(
+        model = model,
+        tokenizer = tokenizer,
+        image = img,
+        text_alias = alias,
+    )
+    for (alias, img) in [
+        ("{{IMAGE_1}}", Image.open("test_image_1.jpg")),
+        ("{{IMAGE_2}}", Image.open("test_image_2.jpg")),
+    ]
+]
+
+prompt = "[INST]{{IMAGE_1}}{{IMAGE_2}}\n" + \
+         "What are the similarities and differences between these two experiments?[/INST]"
+
+# Run prompt through generator, with embeddings. The tokenizer will insert preepared image tokens in place
+# of the aliases
 
 output = generator.generate(
     prompt = prompt,
@@ -108,7 +76,7 @@
     decode_special_tokens = True,
     stop_conditions = [tokenizer.eos_token_id],
     gen_settings = ExLlamaV2Sampler.Settings.greedy(),
-    embeddings = [image_tokens_a],
+    embeddings = image_embeddings,
 )
 
 print(output)